feat: type-level BM25 search corpus and search_types()

.github/workflows/schema-update.yml

@@ -65,7 +65,7 @@ jobs:
         run: |
           git config user.name "github-actions[bot]"
           git config user.email "github-actions[bot]@users.noreply.github.com"
-          git add src/rsc/schema.py src/rsc/mcp_index.json src/rsc/mcp_types.json src/rsc/mcp_bm25_corpus.json src/rsc/mcp_fields_corpus.json pyproject.toml
+          git add src/rsc/schema.py src/rsc/mcp_index.json src/rsc/mcp_types.json src/rsc/mcp_bm25_corpus.json src/rsc/mcp_fields_corpus.json src/rsc/mcp_types_bm25_corpus.json pyproject.toml
           git restore --staged .github/
           git commit -m "chore: generate schema for ${{ steps.schema.outputs.date }}"
           git pull --rebase origin main

src/rsc/__init__.py

@@ -2,6 +2,7 @@
 from .fields import field_index_schema_version, search_fields
 from .index import (
     search_operations,
+    search_types,
     describe_operation,
     describe_type,
     list_queries,
@@ -19,4 +20,5 @@
     "list_types",
     "search_fields",
     "search_operations",
+    "search_types",
 ]

src/rsc/index.py

@@ -28,6 +28,8 @@
 _types_index: dict | None = None
 _bm25_index: object | None = None  # rank_bm25.BM25Okapi once loaded
 _bm25_meta: list[dict] | None = None
+_types_bm25_index: object | None = None  # rank_bm25.BM25Okapi once loaded
+_types_bm25_meta: list[dict] | None = None
 
 _CAMEL_RE = re.compile(r"[A-Z][a-z]+|[a-z]+|[A-Z]+(?=[A-Z]|$)|\d+")
 
@@ -78,6 +80,19 @@ def _get_bm25():
     return _bm25_index, _bm25_meta
 
 
+def _get_types_bm25():
+    global _types_bm25_index, _types_bm25_meta
+    if _types_bm25_index is None:
+        from rank_bm25 import BM25Okapi  # type: ignore
+
+        data = json.loads(
+            (importlib.resources.files("rsc") / "mcp_types_bm25_corpus.json").read_text()
+        )
+        _types_bm25_meta = data["meta"]
+        _types_bm25_index = BM25Okapi(data["corpus"])
+    return _types_bm25_index, _types_bm25_meta
+
+
 def search_operations(search: str, operation_type: str = "all") -> list[dict]:
     """Search queries and/or mutations by BM25 relevance with camelCase tokenization.
 
@@ -119,6 +134,40 @@ def search_operations(search: str, operation_type: str = "all") -> list[dict]:
     return results
 
 
+def search_types(search: str) -> list[dict]:
+    """Search the schema's type graph by BM25 relevance.
+
+    Each result is a GraphQL type whose fields semantically match the query.
+    The 'ops' field lists operations that return this type — use as candidates
+    when finding an operation for a domain concept (e.g. 'cluster', 'SLA domain').
+
+    Args:
+        search: Natural-language query or keywords.
+
+    Returns:
+        List of dicts with keys: name, ops, score.
+    """
+    index, meta = _get_types_bm25()
+    raw_tokens = _split_camel(search) + search.lower().split()
+    query_tokens = [_stem(t) for t in raw_tokens]
+    scores = index.get_scores(query_tokens)
+
+    candidates = [(i, float(scores[i])) for i in range(len(meta))]
+    candidates.sort(key=lambda x: x[1], reverse=True)
+
+    results = []
+    for i, score in candidates[:10]:
+        if score <= 0:
+            break
+        m = meta[i]
+        results.append({
+            "name": m["name"],
+            "ops": m["ops"],
+            "score": round(score, 4),
+        })
+    return results
+
+
 def describe_operation(name: str, operation_type: str) -> dict:
     """Return the full signature for a query or mutation.
 

src/rsc/mcp_indexer.py

@@ -263,6 +263,137 @@ def build_fields_corpus(types: dict, schema_version: str, out_dir: Path) -> None
     print(f"  mcp_fields_corpus.json: {len(corpus)} fields ({corpus_path.stat().st_size // 1024}KB)", flush=True)
 
 
+_TYPE_SKIP_FIELDS = frozenset({
+    "edges", "nodes", "pageInfo", "endCursor", "startCursor",
+    "hasNextPage", "hasPreviousPage", "cursor",
+})
+
+_SYNONYMS: dict[str, list[str]] = {
+    "org": ["organization", "organizations"],
+    "vm": ["virtual", "machine"],
+}
+
+
+def _expand_synonyms(tokens: list[str]) -> list[str]:
+    """Return tokens with synonym expansions appended for known abbreviations."""
+    result = list(tokens)
+    for t in tokens:
+        if t in _SYNONYMS:
+            result.extend(_SYNONYMS[t])
+    return result
+
+
+def build_types_bm25_corpus(ops: dict, types: dict, out_dir: Path) -> None:
+    """Build per-type BM25 search corpus and save mcp_types_bm25_corpus.json.
+
+    Indexes every object and interface type that is reachable from at least one
+    operation. Each document aggregates field-name tokens (with synonym expansion)
+    and the first 8 words of each field description. Connection fields are followed
+    one level deep so that, e.g., a ``clustersConnection`` field on a parent type
+    also contributes Cluster's own field vocabulary to the parent document.
+
+    Ops resolution per type:
+      - Direct: operations whose return type is TypeName or TypeNameConnection.
+      - Node: if XConnection.nodes → X, X inherits ops returning XConnection.
+      - Interface: if TypeA implements InterfaceB, TypeA inherits ops returning
+        InterfaceB or InterfaceBConnection.
+    """
+    # ------------------------------------------------------------------
+    # Step 1: build bare-type-name → op-name mapping
+    # ------------------------------------------------------------------
+    type_to_ops: dict[str, list[str]] = {}
+
+    for _op_type, pool in [("query", ops["queries"]), ("mutation", ops["mutations"])]:
+        for op_name, op_info in pool.items():
+            bare = op_info["return_type"].strip("[]!").strip()
+            if not bare:
+                continue
+            if bare.endswith("Connection"):
+                # The Connection type itself gets the op.
+                type_to_ops.setdefault(bare, []).append(op_name)
+                # Resolve Connection → node type via Connection.nodes field.
+                conn_td = types.get(bare, {})
+                nodes_field = conn_td.get("fields", {}).get("nodes", {})
+                node_bare = nodes_field.get("type", "").strip("[]!").strip()
+                if node_bare and node_bare not in _SCALARS:
+                    type_to_ops.setdefault(node_bare, []).append(op_name)
+            else:
+                type_to_ops.setdefault(bare, []).append(op_name)
+
+    # ------------------------------------------------------------------
+    # Step 2: interface inheritance — propagate interface ops to implementors
+    # ------------------------------------------------------------------
+    for type_name, td in types.items():
+        if td.get("kind") != "type":
+            continue
+        for iface_name in td.get("implements", []):
+            # Direct interface ops
+            for op_name in type_to_ops.get(iface_name, []):
+                type_to_ops.setdefault(type_name, []).append(op_name)
+            # Interface connection ops (InterfaceNameConnection)
+            iface_conn = iface_name + "Connection"
+            for op_name in type_to_ops.get(iface_conn, []):
+                type_to_ops.setdefault(type_name, []).append(op_name)
+
+    # ------------------------------------------------------------------
+    # Step 3: build corpus — one document per reachable type
+    # ------------------------------------------------------------------
+    corpus: list[list[str]] = []
+    meta: list[dict] = []
+
+    for type_name, td in types.items():
+        if td.get("kind") not in ("type", "interface"):
+            continue
+        op_list = type_to_ops.get(type_name, [])
+        if not op_list:
+            continue
+
+        # Deduplicate while preserving insertion order.
+        seen_ops: set[str] = set()
+        deduped_ops: list[str] = []
+        for op in op_list:
+            if op not in seen_ops:
+                seen_ops.add(op)
+                deduped_ops.append(op)
+
+        # Token construction:
+        # 1. Type name tokens with synonym expansion.
+        tokens = _expand_synonyms(_split_camel(type_name))
+
+        for field_name, finfo in td.get("fields", {}).items():
+            if field_name in _TYPE_SKIP_FIELDS:
+                continue
+            # 2a. Field name tokens with synonym expansion.
+            tokens.extend(_expand_synonyms(_split_camel(field_name)))
+            # 2b. First 8 words of field description.
+            desc = finfo.get("description")
+            if desc:
+                tokens.extend(desc.lower().split()[:8])
+            # 3. Follow Connection fields one level deep (depth=0, no recursion).
+            field_bare = finfo.get("type", "").strip("[]!").strip()
+            if field_bare.endswith("Connection"):
+                conn_td = types.get(field_bare, {})
+                nodes_field = conn_td.get("fields", {}).get("nodes", {})
+                node_bare = nodes_field.get("type", "").strip("[]!").strip()
+                if node_bare and node_bare not in _SCALARS:
+                    node_td = types.get(node_bare, {})
+                    for sub_fname in node_td.get("fields", {}).keys():
+                        if sub_fname in _TYPE_SKIP_FIELDS:
+                            continue
+                        tokens.extend(_expand_synonyms(_split_camel(sub_fname)))
+
+        corpus.append(_stem_all(tokens))
+        meta.append({"name": type_name, "ops": deduped_ops})
+
+    out = {"meta": meta, "corpus": corpus}
+    corpus_path = out_dir / "mcp_types_bm25_corpus.json"
+    corpus_path.write_text(json.dumps(out, separators=(",", ":")))
+    print(
+        f"  mcp_types_bm25_corpus.json: {len(corpus)} types ({corpus_path.stat().st_size // 1024}KB)",
+        flush=True,
+    )
+
+
 def main() -> None:
     repo_root = Path(__file__).parent.parent.parent
     schemas_dir = repo_root / "schemas"
@@ -289,6 +420,7 @@ def main() -> None:
 
     build_bm25_corpus(ops, types, out_dir)
     build_fields_corpus(types, schema_version, out_dir)
+    build_types_bm25_corpus(ops, types, out_dir)
 
 
 if __name__ == "__main__":

tests/test_search_types.py

@@ -0,0 +1,68 @@
+"""Tests for search_types() — type-level BM25 search."""
+
+from rsc import search_types
+
+
+def test_search_types_returns_list():
+    results = search_types("cluster")
+    assert isinstance(results, list)
+
+
+def test_results_have_name_and_ops_keys():
+    results = search_types("cluster")
+    assert len(results) > 0
+    for r in results:
+        assert "name" in r
+        assert "ops" in r
+        assert "score" in r
+
+
+def test_cluster_search_returns_cluster_type():
+    results = search_types("cluster")
+    names = [r["name"] for r in results]
+    assert any("Cluster" in n for n in names), f"Expected a Cluster type in results, got: {names}"
+
+
+def test_cluster_ops_contain_cluster_operations():
+    results = search_types("cluster")
+    # ClusterConnection is the canonical connection type returned by cluster queries.
+    cluster_result = next((r for r in results if r["name"] == "ClusterConnection"), None)
+    assert cluster_result is not None, (
+        f"ClusterConnection not found in search_types('cluster') top-10; got: "
+        f"{[r['name'] for r in results]}"
+    )
+    assert isinstance(cluster_result["ops"], list)
+    assert len(cluster_result["ops"]) > 0
+    # At least one op should reference cluster semantics
+    ops_lower = [op.lower() for op in cluster_result["ops"]]
+    assert any("cluster" in op for op in ops_lower), (
+        f"Expected at least one cluster-related op, got: {cluster_result['ops'][:10]}"
+    )
+
+
+def test_scores_positive_for_relevant_results():
+    results = search_types("cluster")
+    assert len(results) > 0
+    for r in results:
+        assert r["score"] > 0, f"Expected positive score, got {r['score']} for {r['name']}"
+
+
+def test_results_capped_at_ten():
+    results = search_types("cluster")
+    assert len(results) <= 10
+
+
+def test_ops_is_list_of_strings():
+    results = search_types("sla domain policy")
+    assert len(results) > 0
+    for r in results:
+        assert isinstance(r["ops"], list)
+        for op in r["ops"]:
+            assert isinstance(op, str)
+
+
+def test_empty_query_returns_list():
+    # Even a poor query must return a list (possibly empty), not raise.
+    results = search_types("zzzzzunlikelytermzzzzz")
+    assert isinstance(results, list)
+    assert len(results) == 0