codegraph-engine 2.1.1__py3-none-any.whl

codegraph/planner.py

@@ -0,0 +1,417 @@
+"""Deterministic Retrieval Planner for CodeGraph MCP.
+
+Translates a grounded TaskSpec into an intent-driven, deterministic RetrievalPlan.
+Query groups distinguish between primary targets, callers, callees, dependencies,
+tests, git changes, framework routes, and architecture slices without guessing.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import sqlite3
+from dataclasses import asdict, dataclass
+
+from codegraph.task import AmbiguityStatus, TaskAmbiguity, TaskIntent, TaskSpec
+
+PLAN_SCHEMA_VERSION = "1.0"
+
+
+class QueryPurpose:
+    PRIMARY_TARGET = "PRIMARY_TARGET"
+    ENTRY_POINT = "ENTRY_POINT"
+    RELATED_SYMBOLS = "RELATED_SYMBOLS"
+    CALLERS = "CALLERS"
+    CALLEES = "CALLEES"
+    DEPENDENCIES = "DEPENDENCIES"
+    IMPLEMENTATIONS = "IMPLEMENTATIONS"
+    TESTS = "TESTS"
+    FRAMEWORK = "FRAMEWORK"
+    GIT = "GIT"
+    ARCHITECTURE = "ARCHITECTURE"
+
+
+@dataclass(frozen=True)
+class QueryGroup:
+    purpose: str
+    queries: tuple[str, ...]
+    weight: float = 1.0
+    required: bool = True
+
+    def as_dict(self) -> dict[str, object]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class RetrievalPlan:
+    schema_version: str = PLAN_SCHEMA_VERSION
+    task_spec_hash: str = ""
+    task_fingerprint: str = ""
+    intent: str = ""
+
+    query_groups: tuple[QueryGroup, ...] = ()
+
+    required_relationships: tuple[str, ...] = ()
+    optional_relationships: tuple[str, ...] = ()
+
+    include_tests: bool = True
+    include_git: bool = False
+    include_framework: bool = True
+    include_architecture: bool = False
+
+    entry_point_strategy: str = "DIRECT"  # DIRECT | TRACE_FROM_ENDPOINT | TRACE_HIERARCHY | GIT_DIFF_ENTRY
+
+    max_depth: int = 3
+    max_nodes: int = 300
+    max_edges: int = 1200
+    token_budget: int = 20_000
+
+    freshness_policy: str = "REQUIRE_FRESH"  # REQUIRE_FRESH | PERMIT_PARTIAL | BEST_EFFORT
+    evidence_policy: str = "VERIFY_SOURCE_HASH"
+    ambiguity_policy: str = "REPORT_OR_ASSUME"  # REPORT_OR_ASSUME | STRICT_FAIL | BEST_EFFORT
+
+    def as_dict(self) -> dict[str, object]:
+        return {
+            "schema_version": self.schema_version,
+            "task_spec_hash": self.task_spec_hash or self.task_fingerprint,
+            "task_fingerprint": self.task_fingerprint or self.task_spec_hash,
+            "intent": self.intent,
+            "query_groups": [q.as_dict() for q in self.query_groups],
+            "required_relationships": list(self.required_relationships),
+            "optional_relationships": list(self.optional_relationships),
+            "include_tests": self.include_tests,
+            "include_git": self.include_git,
+            "include_framework": self.include_framework,
+            "include_architecture": self.include_architecture,
+            "entry_point_strategy": self.entry_point_strategy,
+            "max_depth": self.max_depth,
+            "max_nodes": self.max_nodes,
+            "max_edges": self.max_edges,
+            "token_budget": self.token_budget,
+            "freshness_policy": self.freshness_policy,
+            "evidence_policy": self.evidence_policy,
+            "ambiguity_policy": self.ambiguity_policy,
+        }
+
+
+def compute_task_spec_hash(task_spec: TaskSpec) -> str:
+    """Compute a deterministic SHA-256 hash of a TaskSpec."""
+    canonical_payload = {
+        "intent": task_spec.intent,
+        "goal": task_spec.goal,
+        "targets": sorted(task_spec.targets),
+        "entities": sorted(task_spec.entities),
+        "operations": sorted(task_spec.operations),
+        "constraints": sorted(task_spec.constraints),
+        "exclusions": sorted(task_spec.exclusions),
+        "scope_paths": sorted(task_spec.scope_paths),
+        "scope_modules": sorted(task_spec.scope_modules),
+        "frameworks": sorted(task_spec.frameworks),
+        "time_scope": task_spec.time_scope or "",
+    }
+    raw_bytes = json.dumps(canonical_payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    return hashlib.sha256(raw_bytes).hexdigest()[:16]
+
+
+def build_retrieval_plan(
+    task_spec: TaskSpec,
+    ambiguities: tuple[TaskAmbiguity, ...] = (),
+    con: sqlite3.Connection | None = None,
+    token_budget: int = 20_000,
+    resource_mode: str = "BALANCED",
+) -> RetrievalPlan:
+    """Construct an intent-specific, repository-aware RetrievalPlan for a TaskSpec."""
+    spec_hash = compute_task_spec_hash(task_spec)
+    intent = task_spec.intent.upper()
+
+    # Base query items from targets
+    primary_queries = tuple(task_spec.priority_targets or task_spec.targets or (task_spec.goal,))
+    related_queries: list[str] = []
+
+    if con and primary_queries:
+        from codegraph.query_expansion import expand_query_terms, get_search_queries
+
+        expansions = expand_query_terms(primary_queries[:5], con, max_expansions=8)
+        expanded_terms = get_search_queries(expansions)
+        for term in expanded_terms:
+            if term not in primary_queries and term not in related_queries:
+                related_queries.append(term)
+
+    query_groups: list[QueryGroup] = []
+
+    # Primary target group
+    query_groups.append(
+        QueryGroup(
+            purpose=QueryPurpose.PRIMARY_TARGET,
+            queries=primary_queries,
+            weight=1.0,
+            required=True,
+        )
+    )
+
+    if related_queries:
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.RELATED_SYMBOLS,
+                queries=tuple(related_queries[:8]),
+                weight=0.75,
+                required=False,
+            )
+        )
+
+    # Strategy configuration based on intent
+    include_tests = True
+    include_git = False
+    include_framework = True
+    include_architecture = False
+    entry_point_strategy = "DIRECT"
+    max_depth = 3
+    required_relationships: list[str] = []
+    optional_relationships: list[str] = ["CALLS", "IMPORTS"]
+
+    # Check explicit time_scope or operations in TaskSpec
+    if task_spec.time_scope or "INSPECT_RECENT_CHANGES" in task_spec.operations:
+        include_git = True
+
+    # Adaptive depth calculation based on query complexity and resource mode
+    goal_lower = (task_spec.goal or "").lower()
+    is_shallow = any(pattern in goal_lower for pattern in ("where is", "where's", "find definition", "show definition", "locate"))
+    mode_upper = (resource_mode or "BALANCED").upper()
+
+    if is_shallow:
+        max_depth = 1
+    elif mode_upper == "FAST":
+        max_depth = min(max_depth, 2)
+    elif mode_upper == "DEEP":
+        max_depth = max(max_depth, 5)
+
+    if mode_upper == "FAST":
+        max_nodes = 100
+        max_edges = 400
+    elif mode_upper == "DEEP":
+        max_nodes = 600
+        max_edges = 2400
+    else:
+        max_nodes = 300
+        max_edges = 1200
+
+    if intent in (TaskIntent.UNDERSTAND.value, TaskIntent.EXPLAIN.value):
+        include_architecture = True
+        required_relationships = ["HANDLED_BY", "CALLS"]
+        optional_relationships = ["IMPORTS", "EXTENDS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.ARCHITECTURE,
+                queries=primary_queries[:3],
+                weight=0.6,
+                required=False,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.TESTS,
+                queries=primary_queries[:3],
+                weight=0.5,
+                required=False,
+            )
+        )
+
+    elif intent == TaskIntent.DEBUG.value:
+        include_git = True
+        entry_point_strategy = "TRACE_FROM_ENDPOINT"
+        max_depth = 3
+        required_relationships = ["HANDLED_BY", "CALLS"]
+        optional_relationships = ["IMPORTS", "EXTENDS", "POSSIBLE_CALLS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.CALLERS,
+                queries=primary_queries[:4],
+                weight=0.85,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.TESTS,
+                queries=primary_queries[:3],
+                weight=0.7,
+                required=True,
+            )
+        )
+        if include_git:
+            query_groups.append(
+                QueryGroup(
+                    purpose=QueryPurpose.GIT,
+                    queries=primary_queries[:3],
+                    weight=0.8,
+                    required=False,
+                )
+            )
+
+    elif intent in (TaskIntent.CHANGE.value, TaskIntent.REFACTOR.value):
+        entry_point_strategy = "DIRECT"
+        max_depth = 2
+        required_relationships = ["CALLS"]
+        optional_relationships = ["EXTENDS", "IMPLEMENTS", "IMPORTS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.CALLERS,
+                queries=primary_queries[:4],
+                weight=0.9,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.CALLEES,
+                queries=primary_queries[:4],
+                weight=0.85,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.TESTS,
+                queries=primary_queries[:3],
+                weight=0.75,
+                required=True,
+            )
+        )
+
+    elif intent == TaskIntent.TRACE.value:
+        entry_point_strategy = "TRACE_HIERARCHY"
+        max_depth = 4
+        required_relationships = ["HANDLED_BY", "CALLS"]
+        optional_relationships = ["IMPORTS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.FRAMEWORK,
+                queries=primary_queries[:4],
+                weight=0.95,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.CALLEES,
+                queries=primary_queries[:4],
+                weight=0.9,
+                required=True,
+            )
+        )
+
+    elif intent == TaskIntent.IMPACT.value:
+        entry_point_strategy = "DIRECT"
+        max_depth = 4
+        include_architecture = True
+        required_relationships = ["CALLS", "IMPORTS"]
+        optional_relationships = ["HANDLED_BY", "EXTENDS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.CALLERS,
+                queries=primary_queries[:5],
+                weight=0.95,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.DEPENDENCIES,
+                queries=primary_queries[:5],
+                weight=0.8,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.TESTS,
+                queries=primary_queries[:3],
+                weight=0.75,
+                required=True,
+            )
+        )
+
+    elif intent == TaskIntent.REVIEW.value:
+        include_git = True
+        entry_point_strategy = "GIT_DIFF_ENTRY"
+        max_depth = 2
+        required_relationships = ["CALLS"]
+        optional_relationships = ["IMPORTS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.GIT,
+                queries=primary_queries or ("HEAD~1..HEAD",),
+                weight=1.0,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.TESTS,
+                queries=primary_queries[:3],
+                weight=0.8,
+                required=True,
+            )
+        )
+
+    elif intent == TaskIntent.ARCHITECTURE.value:
+        include_architecture = True
+        include_tests = False
+        max_depth = 2
+        required_relationships = ["HANDLED_BY", "IMPORTS"]
+        optional_relationships = ["CALLS", "EXTENDS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.ARCHITECTURE,
+                queries=primary_queries,
+                weight=1.0,
+                required=True,
+            )
+        )
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.FRAMEWORK,
+                queries=primary_queries,
+                weight=0.85,
+                required=False,
+            )
+        )
+
+    elif intent == TaskIntent.TEST.value:
+        include_tests = True
+        max_depth = 2
+        required_relationships = ["TESTED_BY", "CALLS"]
+        optional_relationships = ["IMPORTS"]
+        query_groups.append(
+            QueryGroup(
+                purpose=QueryPurpose.TESTS,
+                queries=primary_queries,
+                weight=1.0,
+                required=True,
+            )
+        )
+
+    # Ambiguity policy handling
+    has_ambiguity = any(a.status == AmbiguityStatus.AMBIGUOUS for a in ambiguities)
+    ambiguity_policy = "REPORT_OR_ASSUME" if has_ambiguity else "CLEAR"
+
+    return RetrievalPlan(
+        schema_version=PLAN_SCHEMA_VERSION,
+        task_spec_hash=spec_hash,
+        task_fingerprint=spec_hash,
+        intent=intent,
+        query_groups=tuple(query_groups),
+        required_relationships=tuple(dict.fromkeys(required_relationships)),
+        optional_relationships=tuple(dict.fromkeys(optional_relationships)),
+        include_tests=include_tests,
+        include_git=include_git,
+        include_framework=include_framework,
+        include_architecture=include_architecture,
+        entry_point_strategy=entry_point_strategy,
+        max_depth=max_depth,
+        max_nodes=max_nodes,
+        max_edges=max_edges,
+        token_budget=token_budget,
+        freshness_policy="REQUIRE_FRESH",
+        evidence_policy="VERIFY_SOURCE_HASH",
+        ambiguity_policy=ambiguity_policy,
+    )

codegraph/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561.

codegraph/query_expansion.py

@@ -0,0 +1,199 @@
+"""Qualified Query Expansion for CodeGraph MCP v2.1.
+
+Converts raw target terms into precise, explainable search queries.
+
+Key rules:
+  1. If a term resolves to a canonical_id, use canonical_id as the search term
+  2. If a term resolves to a qualified_name, use qualified_name
+  3. If a term matches a route, also expand to the handler symbol
+  4. Generic method names MUST be qualified when qualified context exists:
+     e.g. "get" → "AdminDashboardView.get" (never bare "get")
+  5. Every expansion is recorded as a QueryExpansion with full provenance
+
+Every expansion is source-backed. No fabricated terms.
+"""
+from __future__ import annotations
+
+import sqlite3
+from dataclasses import dataclass
+
+from codegraph.target_resolver import TargetResolution, TargetType, resolve_target
+
+# ---------------------------------------------------------------------------
+# Generic name set — these must never be emitted as bare search terms
+# ---------------------------------------------------------------------------
+
+_GENERIC_NAMES: frozenset[str] = frozenset({
+    "get", "post", "put", "delete", "patch", "head", "options",
+    "run", "call", "start", "stop", "test", "init", "main",
+    "handler", "login", "logout", "register", "index", "home",
+    "create", "update", "destroy", "show", "list", "detail",
+    "load", "save", "close", "open", "reset", "check",
+})
+
+
+# ---------------------------------------------------------------------------
+# Result type
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class QueryExpansion:
+    original_term: str
+    expanded_term: str
+    reason: str
+    source_symbol: str | None  # canonical_id of the source, if any
+    qualification_level: str   # CANONICAL | QUALIFIED | SHORT | LEXICAL | ROUTE
+    confidence: str            # HIGH | MEDIUM | LOW
+
+
+# ---------------------------------------------------------------------------
+# Main expansion function
+# ---------------------------------------------------------------------------
+
+def expand_query_terms(
+    raw_terms: list[str] | tuple[str, ...],
+    con: sqlite3.Connection,
+    target_resolutions: dict[str, TargetResolution] | None = None,
+    max_expansions: int = 8,
+) -> list[QueryExpansion]:
+    """Expand raw target terms to precise, qualified search queries.
+
+    Args:
+        raw_terms: Terms from task_spec.priority_targets or task_spec.targets.
+        con: Active SQLite connection to the indexed repository.
+        target_resolutions: Pre-computed TargetResolution map (optional).
+            If provided, avoids duplicate resolution work.
+        max_expansions: Maximum total expansions to emit.
+
+    Returns:
+        List of QueryExpansion records, deduplicated by expanded_term.
+    """
+    expansions: list[QueryExpansion] = []
+    seen_expanded: set[str] = set()
+    resolutions = target_resolutions or {}
+
+    def _emit(
+        original: str,
+        expanded: str,
+        reason: str,
+        source_symbol: str | None,
+        level: str,
+        confidence: str,
+    ) -> None:
+        if expanded and expanded not in seen_expanded:
+            seen_expanded.add(expanded)
+            expansions.append(QueryExpansion(
+                original_term=original,
+                expanded_term=expanded,
+                reason=reason,
+                source_symbol=source_symbol,
+                qualification_level=level,
+                confidence=confidence,
+            ))
+
+    for term in raw_terms:
+        if len(expansions) >= max_expansions:
+            break
+
+        clean = term.strip()
+        if not clean or len(clean) < 2:
+            continue
+
+        is_generic = clean.lower() in _GENERIC_NAMES
+
+        # Use pre-computed resolution if available, otherwise resolve fresh
+        resolution = resolutions.get(clean) or resolve_target(clean, con)
+
+        # ── 1. Canonical ID ───────────────────────────────────────────────
+        if resolution.canonical_id and resolution.confidence in ("HIGH", "MEDIUM"):
+            _emit(
+                clean, resolution.canonical_id,
+                f"Resolved canonical ID for '{clean}'",
+                resolution.canonical_id,
+                "CANONICAL", resolution.confidence,
+            )
+
+        # ── 2. Qualified name (if different from canonical_id) ─────────────
+        if (
+            resolution.qualified_name
+            and resolution.qualified_name != resolution.canonical_id
+            and resolution.confidence in ("HIGH", "MEDIUM")
+        ):
+            _emit(
+                clean, resolution.qualified_name,
+                f"Qualified name resolution for '{clean}'",
+                resolution.canonical_id,
+                "QUALIFIED", resolution.confidence,
+            )
+
+        # ── 3. Route → handler expansion ──────────────────────────────────
+        if resolution.target_type == TargetType.API_ENDPOINT:
+            route_rows = con.execute(
+                "SELECT handler_name, file_path FROM framework_routes "
+                "WHERE route_path=? OR endpoint_id=? LIMIT 3",
+                (clean, resolution.canonical_id or clean),
+            ).fetchall()
+            for rr in route_rows:
+                handler = rr["handler_name"]
+                if handler:
+                    _emit(
+                        clean, handler,
+                        f"Route '{clean}' handler expansion",
+                        handler, "ROUTE", "HIGH",
+                    )
+
+        # ── 4. Generic names: qualify or suppress ──────────────────────────
+        if is_generic:
+            # We already emitted the qualified form via resolution above.
+            # If resolution didn't find a qualified form, search for one.
+            if not resolution.is_resolved():
+                # Try to find any symbol whose short name == clean
+                qrows = con.execute(
+                    "SELECT canonical_id, qualified_name, name FROM symbols "
+                    "WHERE name=? ORDER BY kind, path LIMIT 5",
+                    (clean,),
+                ).fetchall()
+                for qr in qrows:
+                    qname = qr["qualified_name"]
+                    cid = qr["canonical_id"]
+                    if qname and "." in qname:
+                        # Qualified form exists — emit it, not the bare name
+                        _emit(
+                            clean, qname,
+                            f"Generic name '{clean}' qualified to '{qname}'",
+                            cid, "QUALIFIED", "MEDIUM",
+                        )
+                    # Never emit the bare generic name if qualified form exists
+            # If still nothing: skip the bare generic name entirely
+            continue  # do NOT fall through to SHORT expansion
+
+        # ── 5. Short name fallback (non-generic) ──────────────────────────
+        if not resolution.is_resolved():
+            # Use bare term only for non-generic identifiers with reasonable length
+            if len(clean) >= 3:
+                _emit(
+                    clean, clean,
+                    f"Lexical short name for unresolved term '{clean}'",
+                    None, "LEXICAL", "LOW",
+                )
+        elif clean not in seen_expanded:
+            # Term already resolved but short name wasn't emitted separately
+            # — only emit if not generic
+            _emit(
+                clean, clean,
+                f"Short name '{clean}'",
+                resolution.canonical_id, "SHORT", "MEDIUM",
+            )
+
+    return expansions[:max_expansions]
+
+
+def get_search_queries(expansions: list[QueryExpansion]) -> list[str]:
+    """Extract the deduplicated list of search query strings from expansions."""
+    seen: set[str] = set()
+    queries: list[str] = []
+    for e in expansions:
+        if e.expanded_term not in seen:
+            seen.add(e.expanded_term)
+            queries.append(e.expanded_term)
+    return queries