graph-n4j 0.1.0__py3-none-any.whl

graph_n4j/__init__.py

@@ -0,0 +1,36 @@
+"""
+graph-n4j — Bridge LLM agents with code repositories via Neo4j graph databases.
+
+Based on the CodexGraph NAACL 2025 paper:
+"CodexGraph: Bridging Large Language Models and Code Repositories via Code Graph Databases"
+
+Quick Start:
+    from graph_n4j import GraphN4J
+
+    g = GraphN4J.from_env()              # reads GROQ_API_KEY, NEO4J_* from .env
+    g.index("./my-python-project")       # index a local repo
+    result = g.ask("What does the Foo class do?")
+    print(result.answer)
+"""
+
+__version__ = "0.1.0"
+
+from graph_n4j.core import GraphN4J
+from graph_n4j.graph_db import GraphDB
+from graph_n4j.llm.groq_client import GroqClient
+from graph_n4j.schema import SCHEMA_DESCRIPTION
+from graph_n4j.cache import SkeletonCache, QueryCache, ContextTracker
+from graph_n4j.repo_manager import RepoManager, RepoMetadata
+
+__all__ = [
+    "GraphN4J",
+    "GraphDB",
+    "GroqClient",
+    "SCHEMA_DESCRIPTION",
+    "SkeletonCache",
+    "QueryCache",
+    "ContextTracker",
+    "RepoManager",
+    "RepoMetadata",
+    "__version__",
+]

graph_n4j/agents/__init__.py

@@ -0,0 +1 @@
+# graph_n4j.agents — Primary Agent, Translation Agent, and iterative pipeline

graph_n4j/agents/pipeline.py

@@ -0,0 +1,336 @@
+"""
+pipeline.py -- Iterative dual-agent pipeline with token efficiency optimizations.
+
+Per Section 3.2 of the paper, plus B1-B7 token efficiency improvements:
+  B1. Repo skeleton injection (--include-skeleton)
+  B2. Query/translation caching (--enable-cache)
+  B3. Context pruning — seen code replaced with reference tags (--compact-context)
+  B4. Result truncation to top-K rows (--truncate-results)
+  B5. Cheaper model for Translation Agent (GROQ_MODEL_TRANSLATION)
+  B6. Early-stopping on empty/duplicate results (--early-stop)
+  B7. All caches keyed by (repo_id, commit_sha)
+"""
+
+import os
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+
+from graph_n4j.graph_db import GraphDB
+from graph_n4j.llm.groq_client import GroqClient
+from graph_n4j.agents.primary_agent import PrimaryAgent
+from graph_n4j.agents.translation_agent import TranslationAgent
+from graph_n4j.cache import (
+    SkeletonCache, QueryCache, ContextTracker,
+    truncate_results, should_early_stop,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PipelineResult:
+    """Result from a complete pipeline run."""
+    answer: str
+    rounds_used: int
+    total_queries: int
+    neo4j_roundtrips: int = 0              # B2: actual DB hits (excl. cache hits)
+    context_history: list[dict] = field(default_factory=list)
+    token_usage: dict = field(default_factory=dict)
+    cache_stats: dict = field(default_factory=dict)   # B2: hit/miss rates
+    errors: list[str] = field(default_factory=list)
+    early_stopped: bool = False             # B6: whether early-stop fired
+    early_stop_reason: str = ""
+
+
+class CodexGraphPipeline:
+    """
+    Orchestrates the iterative dual-agent loop with token efficiency optimizations.
+
+    All B-flags default ON except B5 (cheaper translation model), which is opt-in.
+    """
+
+    def __init__(
+        self,
+        groq_client: GroqClient,
+        graph_db: GraphDB,
+        max_rounds: int = 5,
+        query_mode: str = "single_query",
+        task_instruction: str = "",
+        no_translation_agent: bool = False,
+        no_edges: bool = False,
+        # ── B-flag controls ──
+        repo_id: str = "",
+        commit_sha: str = "",
+        include_skeleton: bool = True,      # B1
+        enable_cache: bool = True,          # B2
+        compact_context: bool = True,       # B3
+        truncate_max_rows: int = 15,        # B4 (0 = disabled)
+        translation_model: Optional[str] = None,  # B5
+        early_stop: bool = True,            # B6
+        # Shared cache objects (reuse across pipeline calls in a session)
+        skeleton_cache: Optional[SkeletonCache] = None,
+        query_cache: Optional[QueryCache] = None,
+    ):
+        self.client = groq_client
+        self.db = graph_db
+        self.max_rounds = max_rounds
+        self.query_mode = query_mode
+        self.no_translation_agent = no_translation_agent
+        self.no_edges = no_edges
+
+        # B-flag settings
+        self.repo_id = repo_id
+        self.commit_sha = commit_sha
+        self.include_skeleton = include_skeleton
+        self.enable_cache = enable_cache
+        self.compact_context = compact_context
+        self.truncate_max_rows = truncate_max_rows
+        self.translation_model = translation_model or os.getenv("GROQ_MODEL_TRANSLATION")
+        self.early_stop = early_stop
+
+        # Cache objects — shared across calls for session-level reuse (B7)
+        self.skeleton_cache = skeleton_cache or SkeletonCache()
+        self.query_cache = query_cache or QueryCache(enabled=enable_cache)
+
+        # Build task instruction with ablation modifications
+        full_instruction = task_instruction
+        if no_translation_agent:
+            full_instruction += (
+                "\n\nIMPORTANT: You must write Cypher queries directly instead of "
+                "natural language queries. Output valid Cypher in the 'queries' array."
+            )
+        if no_edges:
+            full_instruction += (
+                "\n\nIMPORTANT: You may ONLY query node attributes (properties). "
+                "Do NOT traverse any edges/relationships. Only use MATCH on single "
+                "node labels with property filters."
+            )
+
+        # B1: Build/load skeleton and inject into system prompt
+        skeleton_text = ""
+        if include_skeleton:
+            skeleton = self.skeleton_cache.get(repo_id, commit_sha)
+            if not skeleton:
+                try:
+                    skeleton = self.skeleton_cache.build_from_db(graph_db, repo_id, commit_sha)
+                except Exception as e:
+                    logger.warning("Failed to build skeleton: %s", e)
+            if skeleton:
+                skeleton_text = (
+                    "\n\n## Repository Structure (skeleton)\n"
+                    "Use this overview to plan targeted queries instead of exploring blindly.\n\n"
+                    + skeleton.to_prompt_text()
+                )
+                logger.info("B1: Skeleton injected (%d chars)", len(skeleton_text))
+
+        self.primary_agent = PrimaryAgent(
+            groq_client=groq_client,
+            query_mode=query_mode,
+            task_instruction=full_instruction + skeleton_text,
+        )
+
+        self.translation_agent = TranslationAgent(
+            groq_client=groq_client,
+            model_override=self.translation_model,  # B5
+        )
+
+    def run(self, question: str) -> PipelineResult:
+        """Execute the full iterative pipeline for a user question."""
+        initial_tokens = self.client.usage.total_tokens
+        context_history: list[dict] = []
+        total_queries = 0
+        neo4j_roundtrips = 0
+        errors: list[str] = []
+        previous_results: list[str] = []
+
+        # B3: Context tracker for this session
+        ctx_tracker = ContextTracker(enabled=self.compact_context)
+
+        logger.info("Pipeline started -- question: '%s'", question[:100])
+
+        for round_num in range(1, self.max_rounds + 1):
+            logger.info("--- Round %d/%d ---", round_num, self.max_rounds)
+
+            # -- Step 1: Primary Agent generates analysis + queries --
+            try:
+                primary_response = self.primary_agent.run(question, context_history)
+            except Exception as e:
+                error_msg = f"Round {round_num}: Primary Agent error -- {e}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+                break
+
+            # Check if sufficient context
+            if primary_response.get("sufficient_context", False):
+                answer = primary_response.get("final_answer", primary_response.get("analysis", ""))
+                logger.info("Primary Agent signals sufficient context at round %d.", round_num)
+                return PipelineResult(
+                    answer=answer,
+                    rounds_used=round_num,
+                    total_queries=total_queries,
+                    neo4j_roundtrips=neo4j_roundtrips,
+                    context_history=context_history,
+                    token_usage=self._get_token_delta(initial_tokens),
+                    cache_stats=self.query_cache.stats() if self.enable_cache else {},
+                    errors=errors,
+                )
+
+            queries = primary_response.get("queries", [])
+            if not queries:
+                logger.warning("No queries generated in round %d. Ending pipeline.", round_num)
+                break
+
+            # -- Step 2: Translate + Execute each query --
+            round_results: list[str] = []
+            for nl_query in queries:
+                total_queries += 1
+
+                # Get Cypher (from cache or Translation Agent)
+                try:
+                    if self.no_translation_agent:
+                        cypher = nl_query
+                    else:
+                        # B2: Check translation cache first
+                        cached_cypher = self.query_cache.get_translation(nl_query)
+                        if cached_cypher:
+                            cypher = cached_cypher
+                            logger.info("B2: Translation cache hit for '%s'", nl_query[:50])
+                        else:
+                            cypher = self.translation_agent.translate(nl_query)
+                            self.query_cache.put_translation_with_nl(nl_query, cypher)
+                except Exception as e:
+                    error_msg = f"Translation error for '{nl_query[:50]}': {e}"
+                    logger.error(error_msg)
+                    errors.append(error_msg)
+                    round_results.append(f"(Translation error: {e})")
+                    continue
+
+                # Execute against Neo4j (with B2 result cache)
+                try:
+                    # B2: Check result cache
+                    cached_result = self.query_cache.get_cypher_result(
+                        cypher, self.repo_id, self.commit_sha
+                    )
+                    if cached_result is not None:
+                        records = cached_result
+                        logger.info("B2: Result cache hit (%d records)", len(records))
+                    else:
+                        records = self.db.run_query(cypher, parameters={"repo_id": self.repo_id})
+                        neo4j_roundtrips += 1
+                        # Cache the result
+                        try:
+                            self.query_cache.put_cypher_result(
+                                cypher, self.repo_id, self.commit_sha, records
+                            )
+                        except Exception:
+                            pass  # Don't fail on cache write errors
+
+                    # B4: Truncate large result sets
+                    if self.truncate_max_rows > 0:
+                        records, truncation_note = truncate_results(
+                            records, max_rows=self.truncate_max_rows
+                        )
+                    else:
+                        truncation_note = None
+
+                    formatted = self.db.format_results(records)
+
+                    # B3: Compact already-seen code
+                    if self.compact_context:
+                        formatted = ctx_tracker.compact_results(formatted, round_num)
+                        for rec in records:
+                            ctx_tracker.mark_seen(rec, round_num)
+
+                    if truncation_note:
+                        formatted += f"\n\n{truncation_note}"
+
+                    round_results.append(formatted)
+                    logger.info("Query returned %d records.", len(records))
+
+                except Exception as e:
+                    error_msg = f"Neo4j query error -- Cypher: '{cypher[:80]}' -- {e}"
+                    logger.error(error_msg)
+                    errors.append(error_msg)
+                    round_results.append(f"(Query execution error: {e})")
+
+            # B6: Early-stopping heuristic
+            if self.early_stop:
+                stop, reason = should_early_stop(round_results, previous_results)
+                if stop:
+                    logger.info("B6: Early stop triggered -- %s", reason)
+                    # Store context and break to final synthesis
+                    context_history.append({
+                        "round": round_num,
+                        "queries": queries,
+                        "results": round_results,
+                    })
+                    return PipelineResult(
+                        answer=self._synthesize_final(question, context_history),
+                        rounds_used=round_num,
+                        total_queries=total_queries,
+                        neo4j_roundtrips=neo4j_roundtrips,
+                        context_history=context_history,
+                        token_usage=self._get_token_delta(initial_tokens),
+                        cache_stats=self.query_cache.stats() if self.enable_cache else {},
+                        errors=errors,
+                        early_stopped=True,
+                        early_stop_reason=reason,
+                    )
+
+            previous_results = round_results
+
+            # -- Step 3: Store round context --
+            context_history.append({
+                "round": round_num,
+                "queries": queries,
+                "results": round_results,
+            })
+
+        # -- Fallback: max rounds reached --
+        logger.info("Max rounds reached. Requesting final synthesis.")
+        answer = self._synthesize_final(question, context_history)
+
+        return PipelineResult(
+            answer=answer,
+            rounds_used=self.max_rounds,
+            total_queries=total_queries,
+            neo4j_roundtrips=neo4j_roundtrips,
+            context_history=context_history,
+            token_usage=self._get_token_delta(initial_tokens),
+            cache_stats=self.query_cache.stats() if self.enable_cache else {},
+            errors=errors,
+        )
+
+    def _synthesize_final(self, question: str, context_history: list[dict]) -> str:
+        """Ask the LLM for a final answer given all gathered context."""
+        try:
+            final_messages = [{
+                "role": "system",
+                "content": (
+                    "You have reached the maximum number of retrieval rounds. "
+                    "Based on all the context gathered so far, provide your best "
+                    "answer to the user's question. Respond with valid JSON: "
+                    '{"analysis": "...", "final_answer": "..."}'
+                ),
+            }, {
+                "role": "user",
+                "content": f"Question: {question}\n\nContext gathered:\n" +
+                    "\n".join(
+                        f"Round {e['round']}: {e['results']}"
+                        for e in context_history
+                    ),
+            }]
+            final = self.client.chat_json(final_messages, temperature=0.3)
+            return final.get("final_answer", final.get("analysis", "Unable to determine answer."))
+        except Exception as e:
+            logger.error("Final synthesis failed: %s", e)
+            return "Unable to synthesize a final answer. See context history for partial results."
+
+    def _get_token_delta(self, initial_total: int) -> dict:
+        """Calculate tokens used in this pipeline run."""
+        delta = self.client.usage.total_tokens - initial_total
+        return {
+            "pipeline_total_tokens": delta,
+            "cumulative": self.client.usage.to_dict(),
+        }

graph_n4j/agents/primary_agent.py

@@ -0,0 +1,185 @@
+"""
+primary_agent.py — The Primary Agent in CodexGraph's dual-agent architecture.
+
+Per Section 3.2 of the paper, the Primary Agent:
+  - Receives the user's code question + accumulated context from prior rounds
+  - Produces: (a) natural-language analysis, (b) one or more NL queries describing
+    what to retrieve from the graph, (c) a decision on whether enough context exists
+  - Supports both 'single_query' and 'multiple_queries' modes
+"""
+
+import logging
+from typing import Optional
+
+from graph_n4j.llm.groq_client import GroqClient
+
+logger = logging.getLogger(__name__)
+
+
+# ═══════════════════════════════════════════
+# System prompts
+# ═══════════════════════════════════════════
+
+PRIMARY_SYSTEM_PROMPT = """You are the Primary Agent in the CodexGraph system. Your job is to help \
+answer questions about a Python code repository by deciding what information to \
+retrieve from a code graph database.
+
+The graph database contains the following node and edge types:
+
+**Nodes:** MODULE (name, file_path), CLASS (name, file_path, signature, code), \
+FUNCTION (name, file_path, signature, code), METHOD (name, file_path, class_name, \
+signature, code), FIELD (name, file_path, class_name), GLOBAL_VARIABLE (name, \
+file_path, code)
+
+**Edges:** CONTAINS (MODULE->CLASS/FUNCTION/GLOBAL_VARIABLE), \
+HAS_METHOD (CLASS->METHOD), HAS_FIELD (CLASS->FIELD), \
+INHERITS (CLASS->CLASS), USES (FUNCTION/METHOD->GLOBAL_VARIABLE/FIELD, \
+with source_association_type and target_association_type)
+
+## Your Process
+
+Each round, you will:
+1. Analyze what information you still need
+2. Write natural language queries describing what to retrieve from the graph
+3. Decide if you have gathered enough context to answer the question
+
+## Output Format
+
+You MUST respond with valid JSON in this exact format:
+{{
+  "analysis": "Your analysis of what information is needed and why",
+  "queries": [
+    "Natural language description of what to retrieve from the graph database"
+  ],
+  "sufficient_context": false
+}}
+
+When sufficient_context is true, include a "final_answer" field:
+{{
+  "analysis": "Final analysis synthesizing all gathered context",
+  "queries": [],
+  "sufficient_context": true,
+  "final_answer": "Your comprehensive answer to the user's question"
+}}
+
+## Rules
+- Each query should describe ONE specific retrieval operation
+- Queries should reference specific node types, names, and relationships from the schema
+- Be precise: "Find the class named X" is better than "find the relevant class"
+- {query_mode_instruction}
+- Use information from previous rounds to refine your queries
+- Do NOT write Cypher — write natural language descriptions of what to retrieve
+"""
+
+SINGLE_QUERY_INSTRUCTION = "Generate exactly ONE query per round (best for complex reasoning tasks)"
+MULTIPLE_QUERIES_INSTRUCTION = "Generate up to 3 queries per round (best for gathering broad context quickly)"
+
+
+def get_system_prompt(
+    query_mode: str = "single_query",
+    task_specific_instruction: str = "",
+) -> str:
+    """
+    Build the Primary Agent's system prompt.
+
+    Args:
+        query_mode: "single_query" or "multiple_queries"
+        task_specific_instruction: Additional instructions for specific tasks
+            (e.g., Code Debugger, Code Generator, etc.)
+    """
+    instruction = (
+        SINGLE_QUERY_INSTRUCTION if query_mode == "single_query"
+        else MULTIPLE_QUERIES_INSTRUCTION
+    )
+
+    prompt = PRIMARY_SYSTEM_PROMPT.format(query_mode_instruction=instruction)
+
+    if task_specific_instruction:
+        prompt += f"\n\n## Task-Specific Instructions\n{task_specific_instruction}"
+
+    return prompt
+
+
+# ═══════════════════════════════════════════
+# Primary Agent class
+# ═══════════════════════════════════════════
+
+class PrimaryAgent:
+    """
+    The Primary Agent: analyzes the user's question, generates NL graph queries,
+    and synthesizes a final answer when sufficient context is gathered.
+    """
+
+    def __init__(
+        self,
+        groq_client: GroqClient,
+        query_mode: str = "single_query",
+        task_instruction: str = "",
+        temperature: float = 0.3,
+    ):
+        self.client = groq_client
+        self.query_mode = query_mode
+        self.temperature = temperature
+        self.system_prompt = get_system_prompt(query_mode, task_instruction)
+
+    def run(
+        self,
+        user_question: str,
+        context_history: list[dict],
+    ) -> dict:
+        """
+        Execute one round of the Primary Agent.
+
+        Args:
+            user_question: The original user question.
+            context_history: List of prior round results:
+                [{"round": 1, "queries": [...], "results": [...]}, ...]
+
+        Returns:
+            Parsed JSON dict with keys: analysis, queries, sufficient_context,
+            and optionally final_answer.
+        """
+        messages = [{"role": "system", "content": self.system_prompt}]
+
+        # Build the user message with accumulated context
+        user_content = f"## User Question\n{user_question}\n"
+
+        if context_history:
+            user_content += "\n## Context from Previous Rounds\n"
+            for entry in context_history:
+                user_content += f"\n### Round {entry['round']}\n"
+                for i, (query, result) in enumerate(
+                    zip(entry["queries"], entry["results"]), 1
+                ):
+                    user_content += f"**Query {i}:** {query}\n"
+                    user_content += f"**Result:** {result}\n\n"
+        else:
+            user_content += "\n(This is the first round — no prior context.)\n"
+
+        messages.append({"role": "user", "content": user_content})
+
+        # Call Groq with JSON mode
+        response = self.client.chat_json(
+            messages=messages,
+            temperature=self.temperature,
+        )
+
+        # Validate response structure
+        if "analysis" not in response:
+            response["analysis"] = ""
+        if "queries" not in response:
+            response["queries"] = []
+        if "sufficient_context" not in response:
+            response["sufficient_context"] = False
+
+        # Enforce query mode limits
+        if self.query_mode == "single_query" and len(response["queries"]) > 1:
+            response["queries"] = response["queries"][:1]
+
+        logger.info(
+            "Primary Agent — round analysis: %s queries, sufficient=%s",
+            len(response["queries"]),
+            response["sufficient_context"],
+        )
+
+        return response

graph_n4j/agents/translation_agent.py

@@ -0,0 +1,98 @@
+"""
+translation_agent.py — The Translation Agent in CodexGraph's dual-agent architecture.
+
+Per Section 3.2 ("write then translate" strategy), the Translation Agent:
+  - Receives the full graph schema + a natural language query from the Primary Agent
+  - Outputs a syntactically valid Cypher query
+  - Uses temperature=0.0 for deterministic translations
+  - The system prompt embeds the full schema so the model translates reliably
+"""
+
+import logging
+
+from graph_n4j.llm.groq_client import GroqClient
+from graph_n4j.schema import SCHEMA_DESCRIPTION
+
+logger = logging.getLogger(__name__)
+
+
+# ═══════════════════════════════════════════
+# System prompt
+# ═══════════════════════════════════════════
+
+TRANSLATION_SYSTEM_PROMPT = f"""You are the Translation Agent in the CodexGraph system. Your ONLY job \
+is to translate natural language queries about a code repository into valid Cypher \
+queries for Neo4j.
+
+{SCHEMA_DESCRIPTION}
+
+## Rules
+
+1. Output ONLY the Cypher query — no explanations, no markdown fences, no comments.
+2. Use MATCH patterns that align with the schema above.
+3. Always use parameterized property matching where possible.
+4. Return useful properties in your RETURN clause (name, file_path, signature, code).
+5. Use OPTIONAL MATCH when a relationship might not exist.
+6. Limit results to 20 rows unless the query implies exhaustive retrieval.
+7. For inheritance chains, use variable-length paths: `-[:INHERITS*]->`.
+8. When searching by name, use case-sensitive exact match by default.
+9. If the query mentions "all methods" or "all fields" of a class, include inherited ones via INHERITS edges.
+10. Always include RETURN — never write a query that only MATCHes without returning.
+
+## Examples
+
+Query: "Find the class named Shape and its methods"
+MATCH (c:CLASS {{name: 'Shape'}})-[:HAS_METHOD]->(m:METHOD) RETURN c.name, c.file_path, c.signature, m.name, m.signature, m.code
+
+Query: "Find all classes that inherit from Shape"
+MATCH (c:CLASS)-[:INHERITS]->(base:CLASS {{name: 'Shape'}}) RETURN c.name, c.file_path, c.signature
+
+Query: "Find the module containing the function calculate_area"
+MATCH (m:MODULE)-[:CONTAINS]->(f:FUNCTION {{name: 'calculate_area'}}) RETURN m.name, m.file_path, f.signature
+
+Query: "Find all global variables used by the function process_data"
+MATCH (f:FUNCTION {{name: 'process_data'}})-[:USES]->(gv:GLOBAL_VARIABLE) RETURN gv.name, gv.file_path, gv.code
+
+Query: "Find the complete inheritance chain for class Circle"
+MATCH (c:CLASS {{name: 'Circle'}})-[:INHERITS*]->(parent:CLASS) RETURN parent.name, parent.file_path, parent.signature
+"""
+
+
+class TranslationAgent:
+    """
+    The Translation Agent: converts natural language graph queries into Cypher.
+    Uses temperature=0.0 for deterministic output.
+
+    B5: Supports a separate model_override so you can run a cheaper model
+    for translation while keeping the Primary Agent on the strong model.
+    Configure via GROQ_MODEL_TRANSLATION env var or pass model_override=.
+    """
+
+    def __init__(self, groq_client: GroqClient, model_override: str | None = None):
+        self.client = groq_client
+        self.model_override = model_override
+
+    def translate(self, nl_query: str) -> str:
+        """
+        Translate a natural language query into a Cypher query.
+
+        Args:
+            nl_query: Natural language description of what to retrieve.
+
+        Returns:
+            A Cypher query string (fences stripped).
+        """
+        messages = [
+            {"role": "system", "content": TRANSLATION_SYSTEM_PROMPT},
+            {"role": "user", "content": nl_query},
+        ]
+
+        cypher = self.client.chat_cypher(
+            messages=messages,
+            temperature=0.0,
+            max_tokens=2048,
+            model_override=self.model_override,
+        )
+
+        logger.info("Translation Agent -- NL: '%s' -> Cypher: '%s'", nl_query[:80], cypher[:80])
+        return cypher

graph_n4j/apps/__init__.py

@@ -0,0 +1 @@
+# graph_n4j.apps — 5 real-world applications from Section 6 of the paper