""" GraphRAG: AST-based code graph stored in PostgreSQL. Supports Python (ast module). Multi-hop traversal for architecture queries. """ import ast, os, json from pathlib import Path from uuid import UUID from db import get_pool from ollama_client import embed def _vec_str(v: list[float]) -> str: return "[" + ",".join(f"{x:.6f}" for x in v) + "]" # ── AST Parser ──────────────────────────────────────────────────────────────── class PythonGraphBuilder(ast.NodeVisitor): def __init__(self, file_path: str, source: str): self.file_path = file_path self.source_lines = source.splitlines() self.nodes: list[dict] = [] self.edges: list[dict] = [] self._class_stack: list[str] = [] self._func_stack: list[str] = [] def _full_name(self, name: str) -> str: parts = self._class_stack + self._func_stack + [name] return ".".join(parts) def _src(self, node) -> str: try: return ast.get_source_segment("\n".join(self.source_lines), node) or "" except Exception: return "" def visit_ClassDef(self, node): full = self._full_name(node.name) self.nodes.append({"type": "class", "name": node.name, "full_name": full, "file": self.file_path, "line_start": node.lineno, "line_end": node.end_lineno, "content": self._src(node)}) self._class_stack.append(node.name) self.generic_visit(node) self._class_stack.pop() def visit_FunctionDef(self, node): full = self._full_name(node.name) parent = ".".join(self._class_stack + self._func_stack) or None self.nodes.append({"type": "function", "name": node.name, "full_name": full, "file": self.file_path, "line_start": node.lineno, "line_end": node.end_lineno, "content": self._src(node)}) if parent: self.edges.append({"source": parent, "target": full, "rel": "has_method"}) self._func_stack.append(node.name) self.generic_visit(node) self._func_stack.pop() visit_AsyncFunctionDef = visit_FunctionDef def visit_Call(self, node): caller = ".".join(self._class_stack + self._func_stack) if not caller: self.generic_visit(node) return if isinstance(node.func, ast.Attribute) and hasattr(node.func.value, "id"): callee = f"{node.func.value.id}.{node.func.attr}" elif isinstance(node.func, ast.Name): callee = node.func.id else: self.generic_visit(node) return self.edges.append({"source": caller, "target": callee, "rel": "calls"}) self.generic_visit(node) def visit_Import(self, node): file_node = self.file_path for alias in node.names: self.edges.append({"source": file_node, "target": alias.name, "rel": "imports"}) self.generic_visit(node) def visit_ImportFrom(self, node): if node.module: file_node = self.file_path self.edges.append({"source": file_node, "target": node.module, "rel": "imports"}) self.generic_visit(node) def parse_python_file(path: str) -> tuple[list[dict], list[dict]]: try: source = Path(path).read_text(encoding="utf-8", errors="ignore") tree = ast.parse(source) builder = PythonGraphBuilder(path, source) builder.visit(tree) # add file-level node builder.nodes.insert(0, {"type": "file", "name": path, "full_name": path, "file": path, "line_start": 1, "line_end": len(source.splitlines()), "content": source[:500]}) return builder.nodes, builder.edges except SyntaxError: return [], [] # ── Indexer ─────────────────────────────────────────────────────────────────── async def index_project(project_id: str, root_dir: str, extensions: list[str] = None): if extensions is None: extensions = [".py"] pool = await get_pool() # clear old data async with pool.acquire() as conn: await conn.execute("DELETE FROM code_edges WHERE project_id=$1", project_id) await conn.execute("DELETE FROM code_nodes WHERE project_id=$1", project_id) node_name_to_id: dict[str, UUID] = {} for ext in extensions: for path in Path(root_dir).rglob(f"*{ext}"): if ".git" in str(path) or "venv" in str(path) or "__pycache__" in str(path): continue nodes, edges = parse_python_file(str(path)) for n in nodes: content_for_embed = f"{n['type']} {n['full_name']}\n{n['content'][:800]}" vec = await embed(content_for_embed) async with pool.acquire() as conn: uid = await conn.fetchval(""" INSERT INTO code_nodes (project_id, node_type, name, full_name, content, embedding, file_path, line_start, line_end) VALUES ($1,$2,$3,$4,$5,$6::vector,$7,$8,$9) RETURNING id """, project_id, n["type"], n["name"], n["full_name"], n["content"][:2000], _vec_str(vec), n["file"], n.get("line_start"), n.get("line_end")) node_name_to_id[n["full_name"]] = uid # store edges (resolve names to IDs later in a second pass) for e in edges: async with pool.acquire() as conn: src_id = node_name_to_id.get(e["source"]) tgt_id = node_name_to_id.get(e["target"]) if src_id and tgt_id: await conn.execute(""" INSERT INTO code_edges (project_id, source_id, target_id, relation_type) VALUES ($1,$2,$3,$4) ON CONFLICT DO NOTHING """, project_id, src_id, tgt_id, e["rel"]) return {"indexed_nodes": len(node_name_to_id)} # ── Graph search ────────────────────────────────────────────────────────────── async def search_code(query: str, project_id: str, limit: int = 5) -> list[dict]: """Vector search over code nodes.""" pool = await get_pool() vec = await embed(query) async with pool.acquire() as conn: rows = await conn.fetch(""" SELECT id, node_type, full_name, content, file_path, line_start, 1-(embedding<=>$1::vector) AS score FROM code_nodes WHERE project_id=$2 ORDER BY embedding<=>$1::vector LIMIT $3 """, _vec_str(vec), project_id, limit) return [dict(r) for r in rows] async def get_neighbors(node_id: UUID, project_id: str, hops: int = 2) -> list[dict]: """Multi-hop graph traversal starting from node_id.""" pool = await get_pool() async with pool.acquire() as conn: # recursive CTE for graph traversal rows = await conn.fetch(""" WITH RECURSIVE graph AS ( SELECT id, node_type, full_name, content, file_path, 0 AS depth FROM code_nodes WHERE id=$1 UNION SELECT n.id, n.node_type, n.full_name, n.content, n.file_path, g.depth+1 FROM code_nodes n JOIN code_edges e ON (e.target_id=n.id AND e.source_id=g.id) OR (e.source_id=n.id AND e.target_id=g.id) JOIN graph g ON TRUE WHERE g.depth < $2 AND n.project_id=$3 ) SELECT DISTINCT id, node_type, full_name, content, file_path, depth FROM graph ORDER BY depth LIMIT 30 """, node_id, hops, project_id) return [dict(r) for r in rows] async def graphrag_context(query: str, project_id: str) -> str: """Full GraphRAG: vector search + multi-hop expansion.""" seed_nodes = await search_code(query, project_id, limit=3) if not seed_nodes: return "" all_nodes: dict[str, dict] = {str(n["id"]): n for n in seed_nodes} for n in seed_nodes: neighbors = await get_neighbors(n["id"], project_id, hops=2) for nb in neighbors: all_nodes[str(nb["id"])] = nb parts = [] for n in sorted(all_nodes.values(), key=lambda x: x.get("depth", 0)): parts.append( f"[{n['node_type']}] {n['full_name']} ({n.get('file_path','')}" f":{n.get('line_start','')})\n{(n.get('content') or '')[:400]}" ) return "\n\n".join(parts[:10])