pretensor 0.1.0__py3-none-any.whl

pretensor/mcp/tools/_rank.py

@@ -0,0 +1,244 @@
+"""Shared ranking helpers for the ``query`` (hybrid) and ``semantic_search`` tools.
+
+Module-private to the ``pretensor.mcp.tools`` package: centralizes the
+``KuzuStore.iter_table_embeddings`` + ``cosine_similarity`` scan
+so the hybrid BM25+cosine fusion path in ``query`` and the
+cosine-only ``semantic_search`` tool share one implementation.
+
+Also hosts the Reciprocal Rank Fusion helper and its ``RRF_K`` constant.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from pretensor.core.registry import RegistryEntry
+from pretensor.core.store import TableEmbeddingRow
+from pretensor.intelligence.embeddings import cosine_similarity
+from pretensor.visibility.filter import VisibilityFilter
+
+from ..service_registry import open_store_for_entry, release_store
+
+__all__ = [
+    "RRF_K",
+    "CosineHit",
+    "cosine_rank_tables",
+    "fusion_key",
+    "rrf_fuse",
+    "score_embedding_rows",
+]
+
+logger = logging.getLogger(__name__)
+
+# Reciprocal Rank Fusion constant (Cormack, Clarke & Büttcher, 2009).
+# 60 is the value from the canonical paper; kept as a named constant so a
+# future tweak (e.g. after an L2 benchmark regression) is a one-line change.
+RRF_K: int = 60
+
+
+@dataclass(frozen=True, slots=True)
+class CosineHit:
+    """One cosine-ranked ``SchemaTable`` hit with its fusion key pre-computed."""
+
+    key: str
+    node_id: str
+    connection_name: str
+    database: str
+    schema_name: str
+    table_name: str
+    description: str
+    score: float
+    cluster_id: str | None
+
+
+def fusion_key(
+    *,
+    connection_name: str,
+    database: str,
+    schema_name: str,
+    table_name: str,
+) -> str:
+    """Stable key shared by BM25 and cosine rankings for RRF fusion.
+
+    The pipe delimiter cannot collide with legal identifier characters, so
+    ``"a|b|c.d"`` maps 1:1 to ``(a, b, c, d)``.
+    """
+    return f"{connection_name}|{database}|{schema_name}.{table_name}"
+
+
+def score_embedding_rows(
+    rows: Iterable[TableEmbeddingRow],
+    qvec: list[float],
+    *,
+    visibility_filter: VisibilityFilter | None,
+) -> tuple[list[CosineHit], bool]:
+    """Score a stream of ``TableEmbeddingRow`` against ``qvec``.
+
+    Shared inner loop used by :func:`cosine_rank_tables` (which chains rows
+    from freshly-opened stores) and by ``context.similar_tables_for_table``
+    (which feeds rows from an already-open ``KuzuStore``). Rows without an
+    embedding are skipped; rows filtered by ``visibility_filter`` are
+    dropped **after** flipping ``any_candidate_had_vector`` so callers can
+    distinguish "no tables carry vectors" from "hidden by visibility".
+
+    Returns ``(scored, any_candidate_had_vector)`` **unsorted and
+    un-truncated** — callers sort + slice to taste.
+    """
+    # ``KuzuStore.iter_table_embeddings(cluster=None)`` joins through
+    # ``IN_CLUSTER`` via OPTIONAL MATCH, which yields one row per cluster
+    # a table belongs to.  A table in N clusters would therefore show up
+    # N times here, double-counting in downstream RRF fusion and clobbering
+    # the cosine_by_key dict in ``_maybe_fuse``.  Dedup by node_id, keeping
+    # the first occurrence — the embedding (and therefore the cosine
+    # score) is identical across the duplicate rows; only ``cluster_id``
+    # varies and we already preserve the first cluster encountered.
+    scored: list[CosineHit] = []
+    seen_nids: set[str] = set()
+    any_candidate_had_vector = False
+
+    for row in rows:
+        if row.embedding is None:
+            continue
+        if row.node_id in seen_nids:
+            continue
+        # Mark the node as seen regardless of whether we end up emitting
+        # a hit for it. A subsequent (table, cluster) row for the same
+        # ``node_id`` would otherwise bypass the dedup — Kuzu's OPTIONAL
+        # MATCH yields one row per cluster the table is in, and we want
+        # exactly one decision per table here.
+        seen_nids.add(row.node_id)
+        any_candidate_had_vector = True
+        if visibility_filter is not None and not visibility_filter.is_table_visible(
+            row.connection_name, row.schema_name, row.table_name
+        ):
+            continue
+        try:
+            score = cosine_similarity(qvec, row.embedding)
+        except ValueError as exc:
+            logger.warning(
+                "score_embedding_rows: dim mismatch for %s: %s",
+                row.node_id,
+                exc,
+            )
+            continue
+        scored.append(
+            CosineHit(
+                key=fusion_key(
+                    connection_name=row.connection_name,
+                    database=row.database,
+                    schema_name=row.schema_name,
+                    table_name=row.table_name,
+                ),
+                node_id=row.node_id,
+                connection_name=row.connection_name,
+                database=row.database,
+                schema_name=row.schema_name,
+                table_name=row.table_name,
+                description=row.description,
+                score=score,
+                cluster_id=row.cluster_id,
+            )
+        )
+
+    return scored, any_candidate_had_vector
+
+
+def any_entry_has_vectors(entries: Iterable[RegistryEntry]) -> bool:
+    """Cheap LIMIT-1 probe: does any entry's store carry a table vector?
+
+    Callers run this BEFORE embedding the query text so a store with no
+    vectors never pays the embedding-client cost — on a cold process the
+    first ``embed()`` call triggers the ONNX model download + session
+    init, which must not happen when the result would be discarded by the
+    "no tables carry vectors" fallback anyway.
+    """
+    for entry in entries:
+        try:
+            store = open_store_for_entry(entry)
+        except RuntimeError as exc:
+            logger.warning(
+                "any_entry_has_vectors: skipping %s (%s)",
+                entry.connection_name,
+                exc,
+            )
+            continue
+        try:
+            if store.has_any_table_embeddings():
+                return True
+        finally:
+            release_store(store)
+    return False
+
+
+def cosine_rank_tables(
+    entries: Iterable[RegistryEntry],
+    qvec: list[float],
+    *,
+    database: str | None,
+    cluster: str | None,
+    visibility_filter: VisibilityFilter | None,
+    k: int,
+) -> tuple[list[CosineHit], bool]:
+    """Top-``k`` ``SchemaTable`` rows by cosine similarity to ``qvec``.
+
+    Iterates ``KuzuStore.iter_table_embeddings`` for each registry entry,
+    applies the ``visibility_filter``, scores rows whose ``embedding`` is
+    present, sorts descending by score, and truncates to ``k``.
+
+    Returns ``(hits, any_candidate_had_vector)``. The boolean preserves the
+    "no tables carry vectors" fallback signal without a second store scan —
+    callers in ``query`` and ``semantic_search`` use it to route to the
+    documented fallback envelope instead of returning an empty result set
+    that looks indistinguishable from "scored but nothing matched".
+
+    Entry selection (including the "unknown database" error path) is left to
+    callers since ``semantic_search`` and ``query`` want different behavior
+    there (error envelope vs. silent empty fusion).
+    """
+    scored: list[CosineHit] = []
+    any_candidate_had_vector = False
+
+    for entry in entries:
+        try:
+            store = open_store_for_entry(entry)
+        except RuntimeError as exc:
+            logger.warning(
+                "cosine_rank_tables: skipping %s (%s)", entry.connection_name, exc
+            )
+            continue
+        try:
+            entry_scored, entry_had_vector = score_embedding_rows(
+                store.iter_table_embeddings(database=database, cluster=cluster),
+                qvec,
+                visibility_filter=visibility_filter,
+            )
+        finally:
+            release_store(store)
+        scored.extend(entry_scored)
+        any_candidate_had_vector = any_candidate_had_vector or entry_had_vector
+
+    scored.sort(key=lambda h: h.score, reverse=True)
+    return scored[:k], any_candidate_had_vector
+
+
+def rrf_fuse(
+    *,
+    bm25_keys: list[str],
+    cosine_keys: list[str],
+    k_constant: int = RRF_K,
+) -> list[tuple[str, float]]:
+    """Reciprocal Rank Fusion of two ranked key lists.
+
+    Input lists are ordered best-first (rank 1 = index 0). Returns a list of
+    ``(key, fused_score)`` ordered by fused score descending. A key appearing
+    in only one list contributes only that list's term — missing means rank
+    infinity, i.e. zero contribution.
+    """
+    scores: dict[str, float] = {}
+    for rank, key in enumerate(bm25_keys, start=1):
+        scores[key] = scores.get(key, 0.0) + 1.0 / (k_constant + rank)
+    for rank, key in enumerate(cosine_keys, start=1):
+        scores[key] = scores.get(key, 0.0) + 1.0 / (k_constant + rank)
+    return sorted(scores.items(), key=lambda item: item[1], reverse=True)

pretensor/mcp/tools/_timed.py

@@ -0,0 +1,26 @@
+"""Shared timing context manager for MCP tool handlers."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import AbstractContextManager
+from pathlib import Path
+from typing import Any
+
+from pretensor.observability import log_timed_operation
+
+
+def timed_tool(
+    tool_name: str, graph_dir: Path, **fields: Any
+) -> AbstractContextManager[None]:
+    # Logger name is intentionally pinned to "pretensor.mcp.server" (not
+    # __name__) so these handler-timing records keep emitting under the same
+    # logger they used when the handlers lived in server.py — preserving any
+    # downstream log routing/filtering. Do not switch to getLogger(__name__).
+    return log_timed_operation(
+        logging.getLogger("pretensor.mcp.server"),
+        event="mcp.tool_handler",
+        tool=tool_name,
+        graph_dir=str(graph_dir),
+        **fields,
+    )

pretensor/mcp/tools/compile_metric.py

@@ -0,0 +1,144 @@
+"""MCP ``compile_metric`` tool payload — compile YAML metrics to validated SQL."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from pretensor.introspection.models.semantic import (
+    SemanticLayer as SemanticLayerModel,
+)
+from pretensor.mcp.tool_registry import McpTool
+from pretensor.semantic.compiler import MetricCompileError, MetricSqlCompiler
+
+from ..service_registry import (
+    load_registry,
+    open_store_for_entry,
+    release_store,
+    resolve_registry_entry,
+)
+
+__all__ = ["compile_metric_payload", "create_tool"]
+
+
+def compile_metric_payload(
+    graph_dir: Path,
+    *,
+    semantic_yaml: str,
+    metric: str,
+    database: str,
+) -> dict[str, Any]:
+    """Build JSON-serializable payload for ``compile_metric``.
+
+    Args:
+        graph_dir: Root directory of the Pretensor graph store.
+        semantic_yaml: Full YAML text of the semantic layer (inline, not a path).
+        metric: Name of the metric to compile.
+        database: ``connection_name`` or logical ``database`` key.
+    """
+    if not semantic_yaml.strip():
+        return {"error": "Missing or empty `semantic_yaml`"}
+    if not metric.strip():
+        return {"error": "Missing or empty `metric`"}
+    if not database.strip():
+        return {"error": "Missing or empty `database`"}
+
+    reg = load_registry(graph_dir)
+    entry = resolve_registry_entry(reg, database)
+    if entry is None:
+        return {"error": f"No registry entry matches database {database!r}"}
+
+    try:
+        layer = SemanticLayerModel.from_yaml(semantic_yaml)
+    except Exception as exc:
+        return {"error": f"Failed to parse semantic_yaml: {exc}"}
+
+    store = open_store_for_entry(entry)
+    try:
+        compiler = MetricSqlCompiler(
+            store,
+            connection_name=entry.connection_name,
+            database_key=entry.database,
+        )
+        try:
+            compiled = compiler.compile(layer, metric)
+        except MetricCompileError as exc:
+            return {"error": str(exc)}
+    finally:
+        release_store(store)
+
+    return {
+        "metric": compiled.metric,
+        "entity": compiled.entity,
+        "sql": compiled.sql,
+        "dialect": compiled.dialect,
+        "warnings": compiled.warnings,
+        "valid": compiled.validation.valid,
+        "syntax_errors": compiled.validation.syntax_errors,
+        "missing_tables": compiled.validation.missing_tables,
+        "missing_columns": compiled.validation.missing_columns,
+        "invalid_joins": [
+            {
+                "message": j.message,
+                "left_table": j.left_table,
+                "right_table": j.right_table,
+            }
+            for j in compiled.validation.invalid_joins
+        ],
+        "suggestions": compiled.validation.suggestions,
+    }
+
+
+def create_tool(graph_dir: Path) -> McpTool:
+    from ._timed import timed_tool
+
+    async def _handle(args: dict) -> dict:
+        yaml_s = str(args.get("semantic_yaml", ""))
+        metric_s = str(args.get("metric", "")).strip()
+        db_t = str(args.get("database", "")).strip()
+        if not yaml_s.strip():
+            return {"error": "Missing `semantic_yaml`"}
+        if not metric_s:
+            return {"error": "Missing `metric`"}
+        if not db_t:
+            return {"error": "Missing `database`"}
+        with timed_tool("compile_metric", graph_dir, database=db_t, metric=metric_s):
+            return compile_metric_payload(
+                graph_dir,
+                semantic_yaml=yaml_s,
+                metric=metric_s,
+                database=db_t,
+            )
+
+    return McpTool(
+        name="compile_metric",
+        description=(
+            "Compile a user-authored YAML metric to validated SQL against the "
+            "indexed graph. Accepts the semantic-layer YAML inline and a metric "
+            "name; returns SQL plus a validation report (missing tables/columns, "
+            "invalid joins, suggestions)."
+        ),
+        input_schema={
+            "type": "object",
+            "properties": {
+                "semantic_yaml": {
+                    "type": "string",
+                    "description": (
+                        "Full YAML text of the semantic layer "
+                        "(metric definitions, joins, and base tables)."
+                    ),
+                },
+                "metric": {
+                    "type": "string",
+                    "description": "Name of the metric to compile",
+                },
+                "database": {
+                    "type": "string",
+                    "description": "Connection name or logical database",
+                },
+            },
+            "required": ["semantic_yaml", "metric", "database"],
+            "additionalProperties": False,
+        },
+        handler=_handle,
+    )

pretensor/mcp/tools/consumers.py

@@ -0,0 +1,161 @@
+"""MCP ``consumers`` tool: external code consumers of a table."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from pretensor.mcp.tool_registry import McpTool
+from pretensor.visibility.filter import VisibilityFilter
+
+from ..service_context import get_effective_visibility_filter
+from ..service_registry import (
+    graph_path_for_entry,
+    load_registry,
+    open_store_for_entry,
+    release_store,
+    resolve_registry_entry,
+)
+from .context import resolve_table_node_id
+
+__all__ = ["create_tool", "consumers_payload"]
+
+
+def consumers_payload(
+    graph_dir: Path,
+    *,
+    table: str,
+    database: str,
+    op: str | None = None,
+    min_confidence: float = 0.6,
+    visibility_filter: VisibilityFilter | None = None,
+) -> dict[str, Any]:
+    """External code locations that read or write ``table``, with provenance."""
+    reg = load_registry(graph_dir)
+    entry = resolve_registry_entry(reg, database)
+    if entry is None:
+        return {"error": "Unknown database connection or name; pass `database`."}
+    gp = graph_path_for_entry(entry)
+    if not gp.exists():
+        return {"error": f"Graph file missing: {gp}"}
+
+    db_key = str(entry.database)
+    vf = visibility_filter or get_effective_visibility_filter()
+    store = open_store_for_entry(entry)
+    try:
+        start_id, err = resolve_table_node_id(
+            store, table, db_key, visibility_filter=vf
+        )
+        if err is not None:
+            try:
+                return json.loads(err)
+            except json.JSONDecodeError:
+                return {"error": err}
+        if start_id is None:
+            return {
+                "error": "Could not resolve table; check the table name and database."
+            }
+
+        rows = store.query_all_rows(
+            """
+            MATCH (c:ExternalConsumer)-[r:CONSUMES]->(t:SchemaTable {node_id: $nid})
+            WHERE r.confidence >= $minc
+              AND ($op IS NULL OR r.op = $op)
+            RETURN c.service_name, c.file_path, c.line_start, c.line_end, c.kind,
+                   c.language, r.op, r.confidence, c.sql_fingerprint
+            ORDER BY r.confidence DESC, c.service_name, c.file_path
+            """,
+            {"nid": start_id, "minc": min_confidence, "op": op},
+        )
+
+        consumers: list[dict[str, Any]] = []
+        reads = writes = 0
+        for service, file_path, ls, le, kind, lang, edge_op, conf, fp in rows:
+            consumers.append(
+                {
+                    "service": service,
+                    "file": file_path,
+                    "line_start": ls,
+                    "line_end": le,
+                    "kind": kind,
+                    "language": lang,
+                    "op": edge_op,
+                    "confidence": conf,
+                    "sql_fingerprint": fp,
+                }
+            )
+            if edge_op == "write":
+                writes += 1
+            else:
+                reads += 1
+
+        return {
+            "table": table,
+            "database": database,
+            "consumers": consumers,
+            "counts": {"read": reads, "write": writes, "total": len(consumers)},
+        }
+    finally:
+        release_store(store)
+
+
+def create_tool(graph_dir: Path) -> McpTool:
+    from ._timed import timed_tool
+
+    async def _handle(args: dict) -> dict:
+        tbl = str(args.get("table", "")).strip()
+        db_t = str(args.get("database", "")).strip()
+        if not tbl:
+            return {"error": "Missing `table`"}
+        if not db_t:
+            return {"error": "Missing `database`"}
+        op = args.get("op")
+        op_s = str(op).strip() if op is not None else None
+        if op_s == "":
+            op_s = None
+        min_conf = float(args.get("min_confidence", 0.6))
+        with timed_tool("consumers", graph_dir, table=tbl, database=db_t, op=op_s):
+            return consumers_payload(
+                graph_dir,
+                table=tbl,
+                database=db_t,
+                op=op_s,
+                min_confidence=min_conf,
+            )
+
+    return McpTool(
+        name="consumers",
+        description=(
+            "External code locations (services) that read or write a table, from "
+            "`pretensor analyze`. Each entry carries service, file, line range, op, "
+            "kind, and confidence."
+        ),
+        input_schema={
+            "type": "object",
+            "properties": {
+                "table": {
+                    "type": "string",
+                    "description": "Table name or schema.table",
+                },
+                "database": {
+                    "type": "string",
+                    "description": "Connection name or logical database",
+                },
+                "op": {
+                    "type": ["string", "null"],
+                    "enum": ["read", "write", None],
+                    "description": "Filter to 'read' or 'write' consumers only",
+                },
+                "min_confidence": {
+                    "type": "number",
+                    "default": 0.6,
+                    "minimum": 0.0,
+                    "maximum": 1.0,
+                },
+            },
+            "required": ["table", "database"],
+            "additionalProperties": False,
+        },
+        handler=_handle,
+    )