KB Semantic Search

1"""kb-semantic-search — semantic (vector) search over the KB's document chunks.23The thematic entry of the unified-retrieval flow: embeds ``query`` and runs a4chunk kNN over the Milvus-backed knowledge index via the gais.knowledge SDK5(``Gais.knowledge.search``), which carries the per-execution USER_TOKEN and6talks to the permissioned www-emblema v2 knowledge API — this node imports no7HTTP client. Permissions are enforced SERVER-SIDE: the endpoint derives the8accessible scope from the caller's own drive permissions and only ever9narrows it.1011Scoping (1.3.0): ``driveItemIds`` is the ONLY scope input — Drive item ids,12folders and/or files (folders expand server-side to their whole subtree,13files to their document). "KB folder id" was the legacy pre-Drive concept;14a folder's drive-item id IS the kb id. Absent-vs-explicit-empty is15preserved: an unwired input means "no narrowing", a wired input that16produced ZERO ids scopes to nothing (clean empty output) — it never17silently broadens. Any input shape is accepted (id string, DriveItem18object, lists of either, JSON-encoded forms) — a scope input never crashes19on shape.2021Outputs (flat):22- ``chunks``      — the raw ranked chunk dicts for structured downstream use;23- ``chunkTexts``  — self-describing texts (provenance header + body) shaped24  for Knowledge Synthesizer's ``context`` port;25- ``documentIds`` — deduped, order-preserving document ids for the26  Ontology From Documents bridge node.2728An empty result set is a CLEAN empty output (three empty lists), never an29error — downstream nodes degrade gracefully.30"""3132from __future__ import annotations3334import json35import re36import sys37import traceback38from typing import Any3940from gais import Gais4142# 1.4.0: fetch WIDE by default (k=100) and let the server-side cross-encoder43# rerank + threshold decide what survives — the reranker, not k, is the44# precision instrument.45MAX_K = 10046DEFAULT_K = 100474849def _parse_drive_item_ids(raw: Any) -> list[str] | None:50    """Parse the driveItemIds scope input, accepting every shape the canvas51    actually delivers, while PRESERVING absent-vs-explicit-empty.5253    Investigate 2026-07-22: a Folder picker's single DriveItem output wired54    into the DriveItem[] port arrives as a BARE id string — the JSON-only55    parser crashed the node. A scope input must never crash on shape; it56    normalizes:57    - None / ""                -> None: "no narrowing" (input not wired).58    - JSON string              -> parsed, then normalized as below.59    - bare string              -> the id itself (comma/whitespace-separated60                                  strings contribute one id each).61    - dict (a DriveItem object) -> its "id".62    - list/tuple               -> per item: strings kept, dicts contribute63                                  their "id"; blanks/junk dropped.64    - [] / "[]" / only-blanks  -> []: an EXPLICIT empty scope ("these ZERO65      items") — the wired-input-found-nothing case must scope to nothing,66      never broaden (N1, eng review 2026-07-22).67    """68    if raw is None or raw == "":69        return None70    if isinstance(raw, str):71        try:72            raw = json.loads(raw)73        except json.JSONDecodeError:74            # Not JSON: a bare id (or several, separator-joined) from a75            # single-DriveItem edge. Never a crash.76            return [tok for tok in re.split(r"[,\s]+", raw.strip()) if tok]77    if isinstance(raw, str):78        # A JSON-encoded bare string ('"<id>"').79        return [raw.strip()] if raw.strip() else []80    if isinstance(raw, dict):81        raw = [raw]82    if not isinstance(raw, (list, tuple)):83        raise ValueError(84            "'driveItemIds' must be ids (a string, a list, or DriveItem objects)"85        )86    ids: list[str] = []87    for item in raw:88        if isinstance(item, str) and item.strip():89            ids.append(item.strip())90        elif isinstance(item, dict):91            item_id = item.get("id")92            if isinstance(item_id, str) and item_id.strip():93                ids.append(item_id.strip())94    return ids959697def _chunk_text(index: int, chunk: dict) -> str:98    """One self-describing context entry: provenance header + chunk body."""99    document_id = chunk.get("documentId") or ""100    try:101        score = round(float(chunk.get("score") or 0.0), 4)102    except (TypeError, ValueError):103        score = 0.0104    text = chunk.get("text") or ""105    return (106        f"[Chunk {index} | Document ID: {document_id} | Relevance: {score}]\n{text}"107    )108109110def process(inputs: dict[str, Any]) -> dict[str, Any]:111    query = inputs.get("query")112    if not isinstance(query, str) or not query.strip():113        raise ValueError("Required input 'query' not provided")114    query = query.strip()115116    # Clamp k to the server's cap (100). Default = fetch wide, rerank filters.117    try:118        k = int(inputs.get("k") or DEFAULT_K)119    except (TypeError, ValueError):120        k = DEFAULT_K121    k = max(1, min(k, MAX_K))122123    # Rerank threshold (0-1): forwarded to the server's cross-encoder filter;124    # absent/junk -> None (the server applies its own default, 0.01).125    rerank_threshold: float | None126    try:127        raw_threshold = inputs.get("rerankThreshold")128        rerank_threshold = None if raw_threshold in (None, "") else float(raw_threshold)129    except (TypeError, ValueError):130        rerank_threshold = None131    if rerank_threshold is not None:132        rerank_threshold = max(0.0, min(rerank_threshold, 1.0))133134    # 1.3.0: driveItemIds is the ONLY scope input. "KB folder id" was the135    # legacy pre-Drive concept — a folder's drive-item id IS the kb id (Milvus136    # kb_id = document.parentId), and the drive path additionally expands137    # folder SUBTREES, which the legacy kbIds path never did. A stale `kbIds`138    # value from an old canvas is ignored.139    drive_item_ids = _parse_drive_item_ids(inputs.get("driveItemIds"))140141    # Explicitly scoped to zero items -> clean empty output WITHOUT an API142    # round-trip (the server would answer the same uniform empty).143    if drive_item_ids is not None and len(drive_item_ids) == 0:144        print(145            "[kb-semantic-search] explicit empty scope — clean empty output",146            file=sys.stderr,147        )148        return {"chunks": [], "chunkTexts": [], "documentIds": []}149150    chunks = Gais.knowledge.search(151        query=query,152        k=k,153        drive_item_ids=drive_item_ids,154        rerank_threshold=rerank_threshold,155    )156    if not isinstance(chunks, list):157        chunks = []158    chunks = [chunk for chunk in chunks if isinstance(chunk, dict)]159160    chunk_texts = [_chunk_text(i, chunk) for i, chunk in enumerate(chunks, 1)]161162    # Deduped, order-preserving document ids (dict preserves insertion order).163    document_ids = list(164        dict.fromkeys(165            chunk["documentId"]166            for chunk in chunks167            if isinstance(chunk.get("documentId"), str) and chunk["documentId"]168        )169    )170171    print(172        f"[kb-semantic-search] k={k} "173        f"driveItemIds={len(drive_item_ids) if drive_item_ids else 0} "174        f"chunks={len(chunks)} documents={len(document_ids)}",175        file=sys.stderr,176    )177178    return {179        "chunks": chunks,180        "chunkTexts": chunk_texts,181        "documentIds": document_ids,182    }183184185def main() -> None:186    try:187        envelope = json.loads(sys.stdin.read() or "{}")188        if not isinstance(envelope, dict):189            envelope = {}190        inputs = envelope.get("inputs", {}) or {}191        json.dump(process(inputs), sys.stdout)192    except Exception as e:193        print(194            json.dumps({195                "error": str(e),196                "errorType": type(e).__name__,197                "traceback": traceback.format_exc(),198            }),199            file=sys.stderr,200        )201        sys.exit(1)202203204if __name__ == "__main__":205    main()