salim0986/graph-bug-ai

1"""2M8 — Agent tools for the reflection and agentic review nodes.3 4These tools wrap the existing graph/vector infrastructure with safe,5template-based interfaces so the LLM can never execute arbitrary Cypher.6 7Three tools are exposed:8  query_graph        — named Cypher templates, parameterised at call time9  search_similar_code — semantic vector search10  get_file_excerpt   — retrieve entity source code from the graph11 12The LLM can request any combination of these during reflection.13In the current implementation the reflection node calls them14programmatically; full tool-use (LLM-driven dispatch) is the next step.15"""16 17from __future__ import annotations18 19from typing import Any, Dict, List, Optional, TYPE_CHECKING20 21from .logger import setup_logger22 23if TYPE_CHECKING:24    from .graph_builder import GraphBuilder25    from .vector_builder import VectorBuilder26 27logger = setup_logger(__name__)28 29 30# ---------------------------------------------------------------------------31# Safe Cypher templates — the LLM can only reference these by name.32# No arbitrary Cypher is ever passed through.33# ---------------------------------------------------------------------------34 35_SAFE_TEMPLATES: Dict[str, str] = {36    "callers": """37        MATCH (caller:Entity {repo_id: $repo_id})-[:CALLS|MAY_CALL]->38              (target:Entity {repo_id: $repo_id, name: $entity_name})39        RETURN DISTINCT caller.name AS name, caller.file AS file,40               caller.type AS type, caller.start_line AS line41        LIMIT 2042    """,43    "callees": """44        MATCH (source:Entity {repo_id: $repo_id, name: $entity_name})45              -[:CALLS|MAY_CALL]->(callee)46        RETURN DISTINCT callee.name AS name, callee.file AS file,47               callee.type AS type48        LIMIT 2049    """,50    "file_entities": """51        MATCH (e:Entity {repo_id: $repo_id, file: $file_path})52        RETURN e.name AS name, e.type AS type,53               e.start_line AS start_line, e.end_line AS end_line54        ORDER BY e.start_line55        LIMIT 5056    """,57    "cycle_check": """58        MATCH (start:Entity {repo_id: $repo_id, file: $file_path})59        MATCH path = (start)-[:CALLS*2..6]->(start)60        RETURN [n IN nodes(path) | n.name] AS cycle61        LIMIT 562    """,63    "blast_radius": """64        MATCH path = (caller:Entity {repo_id: $repo_id})65                     -[:CALLS|MAY_CALL*1..4]->66                     (target:Entity {repo_id: $repo_id, name: $entity_name})67        RETURN DISTINCT caller.name AS name, caller.file AS file,68               min(length(path)) AS depth69        ORDER BY depth, caller.file70        LIMIT 3071    """,72}73 74 75# ---------------------------------------------------------------------------76# Public tool functions77# ---------------------------------------------------------------------------78 79def query_graph(80    graph_db: "GraphBuilder",81    repo_id: str,82    template_name: str,83    params: Dict[str, Any],84) -> List[Dict[str, Any]]:85    """86    Execute a named Cypher template against the knowledge graph.87 88    `template_name` must be one of the keys in `_SAFE_TEMPLATES`.89    Arbitrary Cypher strings are rejected.  `repo_id` is always injected90    so callers cannot query across tenants.91 92    Returns a list of record dicts; returns [] on any error.93    """94    if template_name not in _SAFE_TEMPLATES:95        logger.warning(96            f"[M8] query_graph: unknown template '{template_name}'. "97            f"Allowed: {list(_SAFE_TEMPLATES.keys())}"98        )99        return []100 101    cypher = _SAFE_TEMPLATES[template_name]102    query_params = {**params, "repo_id": repo_id}103 104    try:105        with graph_db.driver.session() as session:106            result = session.run(cypher, **query_params)107            return [dict(r) for r in result]108    except Exception as e:109        logger.error(f"[M8] query_graph('{template_name}') failed: {e}")110        return []111 112 113def search_similar_code(114    vector_db: "VectorBuilder",115    repo_id: str,116    query: str,117    top_k: int = 5,118) -> List[Dict[str, Any]]:119    """120    Run a semantic similarity search against the vector index.121 122    Returns a list of dicts with keys: file, name, score, snippet.123    Returns [] on any error.124    """125    if not query or not query.strip():126        return []127    try:128        results = vector_db.search(129            repo_id=repo_id,130            query=query,131            top_k=top_k,132        )133        return [134            {135                "file": r.get("file", ""),136                "name": r.get("name", ""),137                "score": r.get("score", 0.0),138                "snippet": r.get("code", "")[:300],139            }140            for r in (results or [])141        ]142    except Exception as e:143        logger.error(f"[M8] search_similar_code failed: {e}")144        return []145 146 147def get_file_excerpt(148    graph_db: "GraphBuilder",149    repo_id: str,150    file_path: str,151    start_line: int,152    end_line: int,153) -> str:154    """155    Retrieve entity source code from the graph for a given line range.156 157    Finds the Entity whose start_line is closest to `start_line` inside158    `file_path` and returns its `raw_code`.  Returns "" on miss or error.159    """160    if not file_path or start_line < 0:161        return ""162    cypher = """163    MATCH (e:Entity {repo_id: $repo_id, file: $file_path})164    WHERE e.start_line >= $start_line AND e.start_line <= $end_line165    RETURN e.raw_code AS code, e.name AS name166    ORDER BY e.start_line167    LIMIT 3168    """169    try:170        with graph_db.driver.session() as session:171            result = session.run(172                cypher,173                repo_id=repo_id,174                file_path=file_path,175                start_line=start_line,176                end_line=end_line,177            )178            rows = [r for r in result]179        if not rows:180            return ""181        return "\n\n".join(182            f"# {r['name']}\n{r['code']}" for r in rows if r["code"]183        )184    except Exception as e:185        logger.error(f"[M8] get_file_excerpt failed: {e}")186        return ""187 188 189def available_templates() -> List[str]:190    """Return the list of valid template names (for prompt injection)."""191    return list(_SAFE_TEMPLATES.keys())192