codecompass-mcp 2.0.0__py3-none-any.whl

main.py

@@ -0,0 +1,306 @@
+"""CodeCompass — code dependency index for LLM coding agents.
+
+Commands:
+    ingest-code <repo_path> --project <name> [--normalize] [--dump-triples <out.json>]
+    load-triples <triples.json> --project <name>
+    watch <repo_path> --project <name>
+    dedupe-edges [--dry-run]
+    setup
+"""
+
+import sys
+from dotenv import load_dotenv
+from rich.console import Console
+
+from graph.code_graph_client import get_client
+from ingestion.code_parser import parse_directory
+from ingestion.hierarchy_builder import build_hierarchy, get_file_id_map
+from ingestion.file_watcher import FileWatcher
+
+load_dotenv(override=True)
+console = Console()
+
+
+def ingest_code(repo_path: str, project_name: str, normalize: bool = False, dump_triples: str | None = None) -> None:
+    """Ingest a codebase into the code knowledge graph.
+
+    Phase 1: Walk the repo and write the Project → Folder → File skeleton.
+    Phase 2: Parse every source file with tree-sitter into CodeTriples.
+    Phase 3: Normalize entity names via Haiku (only if --normalize is passed).
+    Phase 4: Write all triples to the project graph.
+
+    If dump_triples is given, raw triples are written to that JSON file and
+    the ingest stops — use this to normalize externally, then reload with
+    load-triples.
+    """
+    import os
+    import json
+    from tqdm import tqdm
+
+    repo_path = os.path.abspath(repo_path)
+    console.print(f"[bold blue]Ingesting codebase:[/] {repo_path}")
+    console.print(f"[dim]Project name:[/] {project_name}")
+
+    client = get_client(project_name)
+    client.ensure_indexes()
+
+    console.print("[dim]Phase 1/4 — Building hierarchy…[/]")
+    file_id_map = build_hierarchy(repo_path, project_name, client)
+    console.print(f"[dim]  {len(file_id_map)} source files indexed[/]")
+
+    console.print("[dim]Phase 2/4 — Parsing source files…[/]")
+    raw_triples = parse_directory(repo_path, progress=True)
+    console.print(f"[dim]  {len(raw_triples)} raw triples extracted[/]")
+
+    if not raw_triples:
+        console.print("[yellow]No triples extracted — check that the repo contains supported files.[/]")
+        client.close()
+        return
+
+    if dump_triples:
+        data = [
+            {
+                "from_entity": t.from_entity,
+                "from_type": t.from_type,
+                "relation_type": t.relation_type,
+                "to_entity": t.to_entity,
+                "to_type": t.to_type,
+                "source_file": t.source_file,
+                "line_number": t.line_number,
+            }
+            for t in raw_triples
+        ]
+        with open(dump_triples, "w") as f:
+            json.dump(data, f, indent=2)
+        client.close()
+        console.print(f"[bold green]Dumped {len(raw_triples)} raw triples to:[/] {dump_triples}")
+        console.print("[dim]Normalize them, then run: codecompass load-triples <file> --project <name>[/]")
+        return
+
+    if normalize:
+        from ingestion.code_normalizer import normalize_triples
+        console.print("[dim]Phase 3/4 — Normalizing triples via Haiku…[/]")
+        triples = normalize_triples(raw_triples, progress=True)
+        console.print(f"[dim]  {len(triples)} triples after normalization[/]")
+    else:
+        console.print("[dim]Phase 3/4 — Skipping normalization (pass --normalize to enable)[/]")
+        triples = raw_triples
+
+    console.print("[dim]Phase 4/4 — Writing to Neo4j…[/]")
+    written = client.write_code_triples_batch(triples, file_id_map, project_name)
+
+    total_nodes = client.node_count()
+    client.close()
+
+    console.print(
+        f"[bold green]Done.[/] Wrote {written} triples. "
+        f"Graph now has {total_nodes} nodes."
+    )
+    _register_project_agents_md(repo_path, project_name)
+
+
+_CODECOMPASS_START = "<!-- codecompass-code-graph-start -->"
+_CODECOMPASS_END = "<!-- codecompass-code-graph-end -->"
+
+
+def _register_project_agents_md(repo_path: str, project_name: str) -> None:
+    """Write or update the Code graph section in the project's AGENTS.md.
+
+    Uses HTML comment markers so re-ingesting safely replaces only that block.
+    """
+    import os, re
+
+    block = (
+        f"{_CODECOMPASS_START}\n"
+        f"## Code graph\n\n"
+        f"This project is indexed in the CodeCompass code graph as `{project_name}`. "
+        f"Query it before editing to know what to read:\n\n"
+        f"```bash\n"
+        f"# Run from your codecompass install directory:\n"
+        f"codecompass --deps <file> --project {project_name}\n"
+        f"codecompass --impact \"<function>\" --project {project_name}\n"
+        f"codecompass --tree {project_name}\n"
+        f"```\n\n"
+        f"Re-ingest after adding files:\n"
+        f"```bash\n"
+        f"codecompass ingest-code {repo_path} --project {project_name}\n"
+        f"```\n"
+        f"{_CODECOMPASS_END}"
+    )
+
+    agents_md_path = os.path.join(repo_path, "AGENTS.md")
+
+    if os.path.exists(agents_md_path):
+        with open(agents_md_path) as f:
+            content = f.read()
+        if _CODECOMPASS_START in content:
+            pattern = re.escape(_CODECOMPASS_START) + r".*?" + re.escape(_CODECOMPASS_END)
+            new_content = re.sub(pattern, block, content, flags=re.DOTALL)
+        else:
+            new_content = content.rstrip() + f"\n\n---\n\n{block}\n"
+    else:
+        new_content = block + "\n"
+
+    with open(agents_md_path, "w") as f:
+        f.write(new_content)
+
+    console.print(f"[dim]  Registered in {agents_md_path}[/]")
+
+
+def load_triples(triples_file: str, project_name: str) -> None:
+    """Load pre-normalized triples from a JSON file into the code graph."""
+    import json
+    from models.code_types import CodeTriple
+
+    with open(triples_file) as f:
+        data = json.load(f)
+
+    triples = [
+        CodeTriple(
+            from_entity=d["from_entity"],
+            from_type=d["from_type"],
+            relation_type=d["relation_type"],
+            to_entity=d["to_entity"],
+            to_type=d["to_type"],
+            source_file=d["source_file"],
+            line_number=d["line_number"],
+        )
+        for d in data
+    ]
+
+    console.print(f"[bold blue]Loading {len(triples)} triples from:[/] {triples_file}")
+    client = get_client(project_name)
+    file_id_map = get_file_id_map(project_name, client)
+
+    written = client.write_code_triples_batch(triples, file_id_map, project_name)
+
+    total_nodes = client.node_count()
+    client.close()
+    console.print(f"[bold green]Done.[/] Wrote {written} triples. Graph now has {total_nodes} nodes.")
+
+
+def watch_code(repo_path: str, project_name: str) -> None:
+    """Watch a repo for file changes and keep the code graph updated incrementally."""
+    import os
+    repo_path = os.path.abspath(repo_path)
+    client = get_client(project_name)
+    file_id_map = build_hierarchy(repo_path, project_name, client)
+    watcher = FileWatcher(repo_path, project_name, client, file_id_map)
+    watcher.start()
+
+
+def dedupe_edges(dry_run: bool = False) -> None:
+    """Remove duplicate RELATION edges (same from-node, type, to-node)."""
+    from config import neo4j_config
+    from neo4j import GraphDatabase
+
+    cfg = neo4j_config()
+    driver = GraphDatabase.driver(cfg["uri"], auth=(cfg["user"], cfg["password"]))
+
+    with driver.session() as session:
+        result = session.run("""
+            MATCH (a)-[r:RELATION]->(b)
+            WITH a, r.type AS rel_type, b, collect(r) AS rels
+            WHERE size(rels) > 1
+            RETURN count(*) AS dup_groups, sum(size(rels) - 1) AS removable
+        """)
+        row = result.single()
+        dup_groups = row["dup_groups"]
+        removable = row["removable"]
+
+        if removable == 0:
+            console.print("[dim]No duplicate edges found.[/]")
+            driver.close()
+            return
+
+        console.print(f"[dim]Found {dup_groups} duplicate groups, {removable} removable edges.[/]")
+
+        if dry_run:
+            console.print("[dim]Dry run — nothing deleted.[/]")
+        else:
+            session.run("""
+                MATCH (a)-[r:RELATION]->(b)
+                WITH a, r.type AS rel_type, b, collect(r) AS rels
+                WHERE size(rels) > 1
+                FOREACH (r IN tail(rels) | DELETE r)
+            """)
+            console.print(f"[bold green]Done.[/] Removed {removable} duplicate edge(s).")
+
+    driver.close()
+
+
+def main():
+    prog = "codecompass"
+    if len(sys.argv) < 2:
+        console.print("[bold]Usage:[/]")
+        console.print(f"  {prog} ingest-code [italic]<repo_path>[/] --project [italic]<name>[/]")
+        console.print(f"  {prog} ingest-code [italic]<repo_path>[/] --project [italic]<name>[/] --normalize")
+        console.print(f"  {prog} ingest-code [italic]<repo_path>[/] --project [italic]<name>[/] --dump-triples [italic]<out.json>[/]")
+        console.print(f"  {prog} load-triples [italic]<triples.json>[/] --project [italic]<name>[/]")
+        console.print(f"  {prog} watch [italic]<repo_path>[/] --project [italic]<name>[/]")
+        console.print(f"  {prog} dedupe-edges [italic][--dry-run][/]")
+        console.print(f"  {prog} setup")
+        sys.exit(1)
+
+    command = sys.argv[1]
+
+    if command == "ingest-code":
+        args = sys.argv[2:]
+        if not args:
+            console.print(f"[red]Usage: {prog} ingest-code <repo_path> --project <name>[/]")
+            sys.exit(1)
+        repo_path = args[0]
+        project_name = "default"
+        normalize = "--normalize" in args
+        dump_triples = None
+        if "--dump-triples" in args:
+            idx = args.index("--dump-triples")
+            if idx + 1 < len(args):
+                dump_triples = args[idx + 1]
+        if "--project" in args:
+            idx = args.index("--project")
+            if idx + 1 < len(args):
+                project_name = args[idx + 1]
+        ingest_code(repo_path, project_name, normalize=normalize, dump_triples=dump_triples)
+
+    elif command == "load-triples":
+        args = sys.argv[2:]
+        if not args:
+            console.print(f"[red]Usage: {prog} load-triples <triples.json> --project <name>[/]")
+            sys.exit(1)
+        triples_file = args[0]
+        project_name = "default"
+        if "--project" in args:
+            idx = args.index("--project")
+            if idx + 1 < len(args):
+                project_name = args[idx + 1]
+        load_triples(triples_file, project_name)
+
+    elif command == "watch":
+        args = sys.argv[2:]
+        if not args:
+            console.print(f"[red]Usage: {prog} watch <repo_path> --project <name>[/]")
+            sys.exit(1)
+        repo_path = args[0]
+        project_name = "default"
+        if "--project" in args:
+            idx = args.index("--project")
+            if idx + 1 < len(args):
+                project_name = args[idx + 1]
+        watch_code(repo_path, project_name)
+
+    elif command == "dedupe-edges":
+        dry_run = "--dry-run" in sys.argv[2:]
+        dedupe_edges(dry_run=dry_run)
+
+    elif command == "setup":
+        from graph.setup import run_setup
+        run_setup()
+
+    else:
+        console.print(f"[red]Unknown command:[/] {command}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

models/code_types.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class CodeTriple:
+    """A single typed relationship extracted from source code."""
+
+    from_entity: str
+    from_type: str   # function | class | module | css_selector | html_element | scss_mixin | scss_variable
+    relation_type: str  # CALLS | IMPORTS | INHERITS | DEFINED_IN | STYLES | HAS_CLASS | POSTS_TO | INCLUDES | USED_BY
+    to_entity: str
+    to_type: str
+    source_file: str  # relative path from project root
+    line_number: int
+
+
+@dataclass
+class FileNode:
+    """Leaf node in the project hierarchy — a single source file."""
+
+    path: str        # relative path from project root
+    name: str        # filename without directory
+    extension: str   # e.g. ".py", ".ts"
+    depth: int       # nesting depth from project root (root files = 1)
+
+
+@dataclass
+class FolderNode:
+    """Intermediate node in the project hierarchy."""
+
+    path: str    # relative path from project root
+    name: str    # folder name
+    depth: int

models/types.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class Entity:
+    id: str
+    name: str
+    type: str                        # Person, Concept, Place, Event, etc.
+    description: Optional[str] = None
+    source_chunk: Optional[str] = None
+
+
+@dataclass
+class Relation:
+    from_id: str
+    to_id: str
+    type: str                        # HAS_COMPONENT, CAUSES, DEPENDS_ON, etc.
+    weight: float = 1.0              # Confidence/strength 0–1
+    description: Optional[str] = None
+
+
+@dataclass
+class Triple:
+    entity_from: Entity
+    relation: Relation
+    entity_to: Entity
+
+
+@dataclass
+class TraversalStep:
+    node_id: str
+    node_name: str
+    relation_type: str
+    relevance_score: float
+    reasoning: str                   # Why the agent followed this edge
+
+
+@dataclass
+class QueryResult:
+    answer: str
+    reasoning_path: list[TraversalStep]
+    nodes_explored: int
+    nodes_retrieved: int
+    hops_taken: int

utils/formatting.py

@@ -0,0 +1,24 @@
+from models.types import TraversalStep
+
+
+def format_traversal_steps(steps: list[TraversalStep]) -> str:
+    """Human-readable summary of a reasoning path"""
+    if not steps:
+        return "(no traversal steps)"
+    lines = []
+    for i, step in enumerate(steps, 1):
+        score = f"{step.relevance_score:.2f}"
+        lines.append(
+            f"  {i}. [{step.relation_type}] → {step.node_name} "
+            f"(score={score}): {step.reasoning}"
+        )
+    return "\n".join(lines)
+
+
+def format_subgraph_rows(rows: list[dict]) -> str:
+    """Format Neo4j subgraph rows for LLM prompts"""
+    return "\n".join(
+        f"({row['n.name']}) --[{row['r.type']}]--> ({row['m.name']})"
+        + (f": {row['r.description']}" if row.get("r.description") else "")
+        for row in rows
+    )