codecompass-mcp 2.0.0__py3-none-any.whl

ingestion/entity_resolver.py

@@ -0,0 +1,179 @@
+import anthropic
+import json
+import re
+import subprocess
+import sys
+from dotenv import load_dotenv
+from graph.neo4j_client import Neo4jClient
+
+load_dotenv(override=True)
+_client = anthropic.Anthropic()
+
+# Chunk size: how many node names to show Haiku per pass.
+# Larger = fewer API calls but risks missing cross-chunk duplicates.
+CHUNK_SIZE = 80
+
+RESOLVER_SYSTEM = """You are a knowledge graph entity resolver.
+
+Given a list of entity names from a knowledge graph, identify groups of names that clearly refer to the SAME real-world entity or concept.
+
+Examples of duplicates:
+- "knowledge tracing" and "Knowledge Tracing System" (same concept, different verbosity)
+- "BKT" and "Bayesian Knowledge Tracing" (acronym + full name)
+- "student model" and "Student Model" (casing only)
+
+Rules:
+- Be CONSERVATIVE — only group names you are highly confident are the same entity
+- The first name in each group becomes the canonical (kept) name — pick the clearest, most complete form
+- Do not group names that are merely related (e.g. "Python" and "Python library" are NOT the same)
+
+Return ONLY valid JSON, no other text:
+{"groups": [["canonical_name", "duplicate1", "duplicate2"], ...]}
+
+If no clear duplicates exist, return {"groups": []}"""
+
+
+def resolve_entities(graph: Neo4jClient, dry_run: bool = False) -> int:
+    """
+    Identify and merge duplicate entity nodes in the graph.
+
+    Strategy (Facade over Haiku + Neo4j):
+      1. Fetch all node names
+      2. Chunk into groups of CHUNK_SIZE and ask Haiku to find duplicates
+      3. Collect all duplicate groups across chunks
+      4. Merge each group: re-point relationships to canonical node, delete duplicates
+
+    Returns the number of nodes merged (0 on dry_run).
+    """
+    all_nodes = graph.get_all_node_names()
+    if len(all_nodes) < 2:
+        print("[resolver] fewer than 2 nodes — nothing to resolve.")
+        return 0
+
+    print(f"[resolver] scanning {len(all_nodes)} nodes for duplicates...", flush=True)
+
+    all_groups: list[list[str]] = []
+
+    for i in range(0, len(all_nodes), CHUNK_SIZE):
+        chunk = all_nodes[i : i + CHUNK_SIZE]
+        node_list = "\n".join(f"- {n['name']} ({n['type']})" for n in chunk)
+
+        response = _client.messages.create(
+            model="claude-haiku-4-5-20251001",
+            max_tokens=1000,
+            system=RESOLVER_SYSTEM,
+            messages=[{"role": "user", "content": f"Entity names to resolve:\n{node_list}"}],
+        )
+
+        try:
+            raw = response.content[0].text.strip()
+            raw = re.sub(r"^```[a-z]*\n?", "", raw, flags=re.MULTILINE)
+            raw = re.sub(r"```$", "", raw.strip(), flags=re.MULTILINE)
+            groups = json.loads(raw).get("groups", [])
+            # Filter out singleton or empty groups
+            groups = [g for g in groups if isinstance(g, list) and len(g) >= 2]
+            if groups:
+                print(f"  chunk {i//CHUNK_SIZE + 1}: found {len(groups)} duplicate group(s)", flush=True)
+            all_groups.extend(groups)
+        except (json.JSONDecodeError, IndexError) as e:
+            print(f"  [resolver] parse error on chunk {i//CHUNK_SIZE + 1}: {e}")
+            continue
+
+    if not all_groups:
+        print("[resolver] no duplicates found.")
+        return 0
+
+    name_to_node = {n["name"]: n for n in all_nodes}
+    merged_count = 0
+
+    for group in all_groups:
+        canonical_name = group[0]
+        duplicates = group[1:]
+
+        canonical_node = name_to_node.get(canonical_name)
+        if not canonical_node:
+            # Canonical name itself doesn't exist — skip
+            continue
+
+        dup_ids = []
+        for dup_name in duplicates:
+            dup_node = name_to_node.get(dup_name)
+            if not dup_node:
+                continue
+            if dup_node["id"] == canonical_node["id"]:
+                continue
+            tag = "[dry-run] " if dry_run else ""
+            print(f"  {tag}'{dup_name}'  →  '{canonical_name}'")
+            dup_ids.append(dup_node["id"])
+
+        if dup_ids and not dry_run:
+            graph.merge_nodes(canonical_node["id"], canonical_name, dup_ids)
+            merged_count += len(dup_ids)
+
+    return merged_count
+
+
+def resolve_dump(graph: Neo4jClient, out_file: str) -> None:
+    """
+    Phase 1 of native resolve: write all node names to a JSON file so that
+    Claude Code can analyse them and produce a groups file.
+
+    Usage:
+        python main.py resolve --native --dump /tmp/nodes.json
+        # → Claude Code reads /tmp/nodes.json, writes /tmp/groups.json
+        python main.py resolve --native --apply /tmp/groups.json
+    """
+    all_nodes = graph.get_all_node_names()
+    with open(out_file, "w") as f:
+        json.dump(all_nodes, f, indent=2)
+    print(f"[resolver] {len(all_nodes)} nodes written to: {out_file}")
+    print()
+    print("Next step — ask Claude Code:")
+    print(f'  "Read {out_file}, find duplicate entity names, write groups to /tmp/resolve_groups.json"')
+    print()
+    print("Then apply:")
+    print("  python main.py resolve --native --apply /tmp/resolve_groups.json")
+
+
+def resolve_apply(graph: Neo4jClient, groups_file: str, dry_run: bool = False) -> int:
+    """
+    Phase 2 of native resolve: load a groups JSON file produced by Claude Code
+    and merge the duplicates.
+
+    groups.json format:
+        [["canonical_name", "duplicate1", "duplicate2"], ...]
+    """
+    with open(groups_file) as f:
+        all_groups = json.load(f)
+
+    all_groups = [g for g in all_groups if isinstance(g, list) and len(g) >= 2]
+    if not all_groups:
+        print("[resolver] groups file is empty — nothing to merge.")
+        return 0
+
+    all_nodes = graph.get_all_node_names()
+    name_to_node = {n["name"]: n for n in all_nodes}
+    merged_count = 0
+
+    for group in all_groups:
+        canonical_name = group[0]
+        duplicates = group[1:]
+        canonical_node = name_to_node.get(canonical_name)
+        if not canonical_node:
+            print(f"  [skip] canonical not found: {canonical_name!r}")
+            continue
+
+        dup_ids = []
+        for dup_name in duplicates:
+            dup_node = name_to_node.get(dup_name)
+            if not dup_node or dup_node["id"] == canonical_node["id"]:
+                continue
+            tag = "[dry-run] " if dry_run else ""
+            print(f"  {tag}'{dup_name}'  →  '{canonical_name}'")
+            dup_ids.append(dup_node["id"])
+
+        if dup_ids and not dry_run:
+            graph.merge_nodes(canonical_node["id"], canonical_name, dup_ids)
+            merged_count += len(dup_ids)
+
+    return merged_count

ingestion/file_watcher.py

@@ -0,0 +1,165 @@
+"""Incremental file watcher — keeps the code graph fresh without full re-ingestion.
+
+Uses the Observer pattern via watchdog. File system events fire callbacks
+that trigger targeted re-ingestion of only the changed file. The watcher
+is decoupled from the ingestion logic — it only knows how to detect changes
+and hand off file paths.
+
+Usage:
+    watcher = FileWatcher(project_root, project_name, client, file_id_map)
+    watcher.start()
+    # ... runs until KeyboardInterrupt or watcher.stop()
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+from typing import Optional
+
+from watchdog.events import FileSystemEventHandler, FileSystemEvent
+from watchdog.observers import Observer
+
+from graph.code_graph_client import CodeGraphClient
+from ingestion.code_parser import parse_file, SUPPORTED_EXTENSIONS
+
+
+def pid_file_path(project_name: str) -> str:
+    """Return the path of the PID file for a given project's watcher process."""
+    return f"/tmp/codecompass_watcher_{project_name}.pid"
+
+
+# ---------------------------------------------------------------------------
+# Event handler
+# ---------------------------------------------------------------------------
+
+class _CodeFileEventHandler(FileSystemEventHandler):
+    """Handles FS events for source files and triggers incremental re-ingestion."""
+
+    def __init__(
+        self,
+        project_root: str,
+        project_name: str,
+        client: CodeGraphClient,
+        file_id_map: dict[str, str],
+    ) -> None:
+        super().__init__()
+        self._project_root = project_root
+        self._project_name = project_name
+        self._client = client
+        self._file_id_map = file_id_map  # {rel_path: neo4j_file_node_id}
+
+    def on_modified(self, event: FileSystemEvent) -> None:
+        if event.is_directory:
+            return
+        self._handle_change(event.src_path)
+
+    def on_created(self, event: FileSystemEvent) -> None:
+        if event.is_directory:
+            return
+        self._handle_change(event.src_path)
+
+    def on_deleted(self, event: FileSystemEvent) -> None:
+        if event.is_directory:
+            return
+        rel_path = os.path.relpath(event.src_path, self._project_root)
+        self._remove_file_from_graph(rel_path)
+
+    def on_moved(self, event: FileSystemEvent) -> None:
+        if event.is_directory:
+            return
+        src_rel = os.path.relpath(event.src_path, self._project_root)
+        dest_rel = os.path.relpath(event.dest_path, self._project_root)
+        if any(src_rel.endswith(ext) for ext in SUPPORTED_EXTENSIONS):
+            self._remove_file_from_graph(src_rel)
+        if any(dest_rel.endswith(ext) for ext in SUPPORTED_EXTENSIONS):
+            self._handle_change(event.dest_path)
+
+    # ------------------------------------------------------------------
+    # Core delta logic
+    # ------------------------------------------------------------------
+
+    def _handle_change(self, abs_path: str) -> None:
+        """Re-parse a changed file and apply only the delta to Neo4j."""
+        ext = Path(abs_path).suffix.lower()
+        if ext not in SUPPORTED_EXTENSIONS:
+            return
+
+        rel_path = os.path.relpath(abs_path, self._project_root)
+        print(f"[file_watcher] changed: {rel_path}")
+
+        # Delete stale entity nodes (File node stays — file still exists)
+        self._client.delete_file_triples(rel_path, self._project_name)
+
+        new_triples = parse_file(abs_path, self._project_root)
+        if not new_triples:
+            return
+
+        file_node_id = self._file_id_map.get(rel_path, "")
+        for triple in new_triples:
+            self._client.write_code_triple(triple, file_node_id, self._project_name)
+
+        print(f"[file_watcher] wrote {len(new_triples)} triples for {rel_path}")
+
+    def _remove_file_from_graph(self, rel_path: str) -> None:
+        """Purge File node and all entity nodes for a deleted or moved file."""
+        print(f"[file_watcher] removed: {rel_path}")
+        self._client.delete_file(rel_path, self._project_name)
+
+
+# ---------------------------------------------------------------------------
+# Public watcher class
+# ---------------------------------------------------------------------------
+
+class FileWatcher:
+    """Watches a project root and keeps its code graph incrementally updated."""
+
+    def __init__(
+        self,
+        project_root: str,
+        project_name: str,
+        client: CodeGraphClient,
+        file_id_map: dict[str, str],
+    ) -> None:
+        self._project_root = project_root
+        self._pid_file = pid_file_path(project_name)
+        self._handler = _CodeFileEventHandler(
+            project_root=project_root,
+            project_name=project_name,
+            client=client,
+            file_id_map=file_id_map,
+        )
+        self._observer = Observer()
+        self._observer.schedule(self._handler, project_root, recursive=True)
+
+    def start(self) -> None:
+        """Start watching. Blocks until stop() is called or KeyboardInterrupt."""
+        self._observer.start()
+        self._write_pid()
+        print(f"[file_watcher] watching {self._project_root} — Ctrl-C to stop")
+        try:
+            while self._observer.is_alive():
+                time.sleep(1)
+        except KeyboardInterrupt:
+            self.stop()
+
+    def stop(self) -> None:
+        """Stop the observer and wait for it to finish."""
+        self._observer.stop()
+        self._observer.join()
+        self._remove_pid()
+        print("[file_watcher] stopped")
+
+    def _write_pid(self) -> None:
+        try:
+            with open(self._pid_file, "w") as f:
+                f.write(str(os.getpid()))
+        except OSError:
+            pass
+
+    def _remove_pid(self) -> None:
+        try:
+            os.unlink(self._pid_file)
+        except OSError:
+            pass

ingestion/graph_writer.py

@@ -0,0 +1,17 @@
+from tqdm import tqdm
+from graph.neo4j_client import Neo4jClient
+from models.types import Triple
+
+
+def write_triples(client: Neo4jClient, triples: list[Triple]) -> int:
+    """Write triples to Neo4j, deduplicating via (from, rel_type, to) key"""
+    seen: set[tuple] = set()
+    written = 0
+    for triple in tqdm(triples, desc="Writing to Neo4j", unit="triple"):
+        key = (triple.entity_from.id, triple.relation.type, triple.entity_to.id)
+        if key in seen:
+            continue
+        seen.add(key)
+        client.write_triple(triple)
+        written += 1
+    return written

ingestion/hierarchy_builder.py

@@ -0,0 +1,148 @@
+"""Hierarchy builder — walks a repo and writes the Project → Folder → File skeleton to Neo4j.
+
+This runs before code_parser so every file has a node to attach entities to.
+No API calls — purely local filesystem traversal.
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from pathlib import Path
+
+from models.code_types import FileNode, FolderNode
+
+# Directory names skipped during traversal
+_SKIP_DIRS = {
+    ".git", "node_modules", "__pycache__", ".venv", "venv",
+    "dist", "build", ".mypy_cache", ".pytest_cache",
+    "coverage", "tmp", "cache", ".nx", "lcov-report",
+}
+
+# Supported source file extensions (mirrors code_parser.SUPPORTED_EXTENSIONS)
+_SOURCE_EXTENSIONS = {".py", ".js", ".ts", ".tsx", ".html", ".css", ".scss"}
+
+
+def build_hierarchy(project_root: str, project_name: str, client) -> dict[str, str]:
+    """Walk project_root and write Project → Folder → File nodes to Neo4j.
+
+    Returns a mapping of {relative_file_path: neo4j_node_id} so the caller
+    can attach entity nodes to the correct File nodes.
+
+    Args:
+        project_root: Absolute path to the repo.
+        project_name: Human-readable project identifier (e.g. "frontend").
+        client: CodeGraphClient connected to the project's database.
+    """
+    project_id = _stable_id(f"project:{project_name}")
+    client.merge_project_node(project_id, project_name, project_root)
+
+    file_id_map: dict[str, str] = {}
+
+    for dirpath, dirnames, filenames in os.walk(project_root):
+        dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS]
+
+        rel_dir = os.path.relpath(dirpath, project_root)
+        is_root = rel_dir == "."
+
+        if not is_root:
+            folder = _make_folder_node(rel_dir, project_root)
+            folder_id = _stable_id(f"folder:{project_name}:{folder.path}")
+            parent_id = _parent_id(folder, project_name, project_id)
+            client.merge_folder_node(folder_id, folder, project_name)
+            client.merge_contains_edge(parent_id, folder_id)
+
+        for filename in filenames:
+            ext = Path(filename).suffix.lower()
+            if ext not in _SOURCE_EXTENSIONS:
+                continue
+
+            full_path = os.path.join(dirpath, filename)
+            rel_path = os.path.relpath(full_path, project_root)
+            file = _make_file_node(rel_path, project_root)
+            file_id = _stable_id(f"file:{project_name}:{file.path}")
+
+            if is_root:
+                parent_node_id = project_id
+            else:
+                parent_node_id = _stable_id(f"folder:{project_name}:{rel_dir}")
+
+            client.merge_file_node(file_id, file, project_name)
+            client.merge_contains_edge(parent_node_id, file_id)
+            file_id_map[rel_path] = file_id
+
+    return file_id_map
+
+
+def collect_file_nodes(project_root: str) -> list[FileNode]:
+    """Return FileNode objects for every supported source file under project_root.
+
+    Useful for dry-run inspection or passing file lists to other pipeline stages.
+    """
+    nodes: list[FileNode] = []
+    for dirpath, dirnames, filenames in os.walk(project_root):
+        dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS]
+        for filename in filenames:
+            ext = Path(filename).suffix.lower()
+            if ext in _SOURCE_EXTENSIONS:
+                full_path = os.path.join(dirpath, filename)
+                rel_path = os.path.relpath(full_path, project_root)
+                nodes.append(_make_file_node(rel_path, project_root))
+    return nodes
+
+
+def collect_folder_nodes(project_root: str) -> list[FolderNode]:
+    """Return FolderNode objects for every non-skipped directory under project_root."""
+    nodes: list[FolderNode] = []
+    for dirpath, dirnames, _ in os.walk(project_root):
+        dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS]
+        rel_dir = os.path.relpath(dirpath, project_root)
+        if rel_dir != ".":
+            nodes.append(_make_folder_node(rel_dir, project_root))
+    return nodes
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+def _make_file_node(rel_path: str, project_root: str) -> FileNode:
+    depth = len(Path(rel_path).parts)
+    return FileNode(
+        path=rel_path,
+        name=Path(rel_path).name,
+        extension=Path(rel_path).suffix.lower(),
+        depth=depth,
+    )
+
+
+def _make_folder_node(rel_dir: str, project_root: str) -> FolderNode:
+    depth = len(Path(rel_dir).parts)
+    return FolderNode(
+        path=rel_dir,
+        name=Path(rel_dir).name,
+        depth=depth,
+    )
+
+
+def _parent_id(folder: FolderNode, project_name: str, project_id: str) -> str:
+    """Return the node ID of the immediate parent of this folder."""
+    parent_path = str(Path(folder.path).parent)
+    if parent_path == ".":
+        return project_id
+    return _stable_id(f"folder:{project_name}:{parent_path}")
+
+
+def get_file_id_map(project_name: str, client) -> dict[str, str]:
+    """Reconstruct the {relative_path: node_id} map from the existing graph.
+
+    Used by load-triples to attach entity nodes to File nodes that were
+    written during a previous build_hierarchy call.
+    """
+    records = client.get_file_nodes(project_name)
+    return {r["path"]: r["id"] for r in records}
+
+
+def _stable_id(key: str) -> str:
+    """Deterministic UUID from a string key — same key always produces the same ID."""
+    return str(uuid.uuid5(uuid.NAMESPACE_DNS, key))

ingestion/reader_agent.py

@@ -0,0 +1,135 @@
+import anthropic
+import json
+import asyncio
+import uuid
+import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dotenv import load_dotenv
+from tqdm.asyncio import tqdm
+from tqdm import tqdm as tqdm_sync
+from models.types import Entity, Relation, Triple
+
+load_dotenv(override=True)
+# Sync client — each call is run in a thread pool to avoid blocking the event loop
+_client = anthropic.Anthropic()
+
+# Cached system prompt block — same every call, so cache it after the first hit
+_EXTRACTION_SYSTEM = [
+    {
+        "type": "text",
+        "text": """You are a knowledge graph extraction agent.
+Given a text chunk, extract all meaningful entities and the relationships between them.
+
+Rules:
+- Entities: concrete nouns, concepts, people, places, events, systems, components
+- Relations: verbs or relationship types connecting entities (CAUSES, DEPENDS_ON, HAS_COMPONENT, USED_BY, etc.)
+- Weight: confidence in the relation from 0.0 to 1.0
+- Be selective — only extract clear, meaningful relationships
+- Use consistent entity names (no duplicates with different casing)
+
+Return ONLY valid JSON, no other text:
+{
+  "entities": [
+    {"name": "Entity Name", "type": "Concept|Person|Place|Event|System", "description": "brief description"}
+  ],
+  "relations": [
+    {"from": "Entity A", "to": "Entity B", "type": "RELATION_TYPE", "weight": 0.9, "description": "brief explanation"}
+  ]
+}""",
+        "cache_control": {"type": "ephemeral"},
+    }
+]
+
+# Expose plain text for any code that imports EXTRACTION_SYSTEM by name
+EXTRACTION_SYSTEM = _EXTRACTION_SYSTEM[0]["text"]
+
+# How many Haiku calls to run concurrently.
+# Haiku rate limits are generous — 15 keeps throughput high without hitting 429s.
+MAX_CONCURRENT = 15
+
+
+def _extract_triples_sync(chunk: str) -> list[Triple]:
+    """Blocking extraction — called from a thread pool."""
+    response = _client.messages.create(
+        model="claude-haiku-4-5-20251001",
+        max_tokens=4096,
+        system=_EXTRACTION_SYSTEM,
+        messages=[{"role": "user", "content": f"Extract knowledge graph from:\n\n{chunk}"}],
+    )
+
+    try:
+        raw = response.content[0].text
+        raw = re.sub(r"^```[a-z]*\n?", "", raw.strip(), flags=re.MULTILINE)
+        raw = re.sub(r"```$", "", raw.strip(), flags=re.MULTILINE)
+        data = json.loads(raw.strip())
+    except (json.JSONDecodeError, IndexError) as e:
+        print(f"[reader_agent] JSON parse failed: {e} | raw: {raw[:200]!r}")
+        return []
+
+    entity_map: dict[str, Entity] = {}
+    for e in data.get("entities", []):
+        # Normalise: strip whitespace, collapse internal spaces, strip trailing punctuation
+        name = re.sub(r"\s+", " ", e.get("name", "").strip()).strip(".,;:()")
+        if not name:
+            continue
+        eid = str(uuid.uuid5(uuid.NAMESPACE_DNS, name.lower()))
+        entity_map[name] = Entity(
+            id=eid,
+            name=name,
+            type=e.get("type", "Concept"),
+            description=e.get("description"),
+            source_chunk=chunk[:100],
+        )
+
+    triples: list[Triple] = []
+    for r in data.get("relations", []):
+        from_entity = entity_map.get(r.get("from", ""))
+        to_entity = entity_map.get(r.get("to", ""))
+        if not from_entity or not to_entity:
+            continue
+        relation = Relation(
+            from_id=from_entity.id,
+            to_id=to_entity.id,
+            type=r.get("type", "RELATED_TO"),
+            weight=float(r.get("weight", 0.8)),
+            description=r.get("description"),
+        )
+        triples.append(Triple(from_entity, relation, to_entity))
+
+    return triples
+
+
+def extract_triples_parallel_sync(chunks: list[str], max_workers: int = MAX_CONCURRENT) -> list[Triple]:
+    """
+    Run extraction on all chunks in parallel using a thread pool.
+    Used for mid-session ingest_source where we're already in a sync context.
+    """
+    all_triples: list[Triple] = []
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = {executor.submit(_extract_triples_sync, chunk): i for i, chunk in enumerate(chunks)}
+        for future in tqdm_sync(as_completed(futures), total=len(futures), desc="Extracting", unit="chunk"):
+            try:
+                all_triples.extend(future.result())
+            except Exception as e:
+                print(f"[reader_agent] chunk failed: {e}")
+    return all_triples
+
+
+async def extract_triples(chunk: str) -> list[Triple]:
+    """Async wrapper — offloads blocking API call to a thread."""
+    return await asyncio.to_thread(_extract_triples_sync, chunk)
+
+
+async def ingest_chunks_parallel(
+    chunks: list[str], max_concurrent: int = MAX_CONCURRENT
+) -> list[Triple]:
+    """Run extraction on all chunks with bounded concurrency."""
+    semaphore = asyncio.Semaphore(max_concurrent)
+
+    async def bounded_extract(chunk: str) -> list[Triple]:
+        async with semaphore:
+            return await extract_triples(chunk)
+
+    tasks = [bounded_extract(c) for c in chunks]
+    results = await tqdm.gather(*tasks, desc="Extracting triples", unit="chunk")
+    return [triple for batch in results for triple in batch]