code-graph-rag 0.0.79__py3-none-any.whl

codebase_rag/models.py

@@ -0,0 +1,94 @@
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple
+
+from rich.console import Console
+
+from .constants import SupportedLanguage
+from .types_defs import MCPHandlerType, MCPInputSchema, PropertyValue
+
+if TYPE_CHECKING:
+    from tree_sitter import Node
+
+
+@dataclass
+class SessionState:
+    confirm_edits: bool = True
+    log_file: Path | None = None
+    cancelled: bool = False
+
+    def reset_cancelled(self) -> None:
+        self.cancelled = False
+
+
+def _default_console() -> Console:
+    return Console(width=None, force_terminal=True)
+
+
+@dataclass
+class AppContext:
+    session: SessionState = field(default_factory=SessionState)
+    console: Console = field(default_factory=_default_console)
+
+
+@dataclass
+class GraphNode:
+    node_id: int
+    labels: list[str]
+    properties: dict[str, PropertyValue]
+
+
+@dataclass
+class GraphRelationship:
+    from_id: int
+    to_id: int
+    type: str
+    properties: dict[str, PropertyValue]
+
+
+class FQNSpec(NamedTuple):
+    scope_node_types: frozenset[str]
+    function_node_types: frozenset[str]
+    get_name: Callable[["Node"], str | None]
+    file_to_module_parts: Callable[[Path, Path], list[str]]
+
+
+@dataclass(frozen=True)
+class LanguageSpec:
+    language: SupportedLanguage | str
+    file_extensions: tuple[str, ...]
+    function_node_types: tuple[str, ...]
+    class_node_types: tuple[str, ...]
+    module_node_types: tuple[str, ...]
+    call_node_types: tuple[str, ...] = ()
+    import_node_types: tuple[str, ...] = ()
+    import_from_node_types: tuple[str, ...] = ()
+    name_field: str = "name"
+    body_field: str = "body"
+    package_indicators: tuple[str, ...] = ()
+    function_query: str | None = None
+    class_query: str | None = None
+    call_query: str | None = None
+
+
+@dataclass
+class Dependency:
+    name: str
+    spec: str
+    properties: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass
+class MethodModifiersAndAnnotations:
+    modifiers: list[str] = field(default_factory=list)
+    annotations: list[str] = field(default_factory=list)
+
+
+@dataclass
+class ToolMetadata:
+    name: str
+    description: str
+    input_schema: MCPInputSchema
+    handler: MCPHandlerType
+    returns_json: bool

codebase_rag/parser_loader.py

@@ -0,0 +1,292 @@
+import importlib
+import subprocess
+import sys
+from copy import deepcopy
+from pathlib import Path
+
+from loguru import logger
+from tree_sitter import Language, Parser, Query
+
+from . import constants as cs
+from . import exceptions as ex
+from . import logs as ls
+from .language_spec import LANGUAGE_SPECS, LanguageSpec
+from .types_defs import LanguageImport, LanguageLoader, LanguageQueries
+
+
+def _try_load_from_submodule(lang_name: cs.SupportedLanguage) -> LanguageLoader:
+    submodule_path = Path(cs.GRAMMARS_DIR) / f"{cs.TREE_SITTER_PREFIX}{lang_name}"
+    python_bindings_path = (
+        submodule_path / cs.BINDINGS_DIR / cs.SupportedLanguage.PYTHON
+    )
+
+    if not python_bindings_path.exists():
+        return None
+
+    python_bindings_str = str(python_bindings_path)
+    try:
+        if python_bindings_str not in sys.path:
+            sys.path.insert(0, python_bindings_str)
+
+        try:
+            module_name = f"{cs.TREE_SITTER_MODULE_PREFIX}{lang_name.replace('-', '_')}"
+
+            setup_py_path = submodule_path / cs.SETUP_PY
+            if setup_py_path.exists():
+                logger.debug(ls.BUILDING_BINDINGS.format(lang=lang_name))
+                result = subprocess.run(
+                    [sys.executable, cs.SETUP_PY, cs.BUILD_EXT_CMD, cs.INPLACE_FLAG],
+                    check=False,
+                    cwd=str(submodule_path),
+                    capture_output=True,
+                    text=True,
+                )
+
+                if result.returncode != 0:
+                    logger.debug(
+                        ls.BUILD_FAILED.format(
+                            lang=lang_name, stdout=result.stdout, stderr=result.stderr
+                        )
+                    )
+                    return None
+                logger.debug(ls.BUILD_SUCCESS.format(lang=lang_name))
+
+            logger.debug(ls.IMPORTING_MODULE.format(module=module_name))
+            module = importlib.import_module(module_name)
+
+            language_attrs: list[str] = [
+                cs.QUERY_LANGUAGE,
+                f"{cs.LANG_ATTR_PREFIX}{lang_name}",
+                f"{cs.LANG_ATTR_PREFIX}{lang_name.replace('-', '_')}",
+            ]
+
+            for attr_name in language_attrs:
+                if hasattr(module, attr_name):
+                    logger.debug(
+                        ls.LOADED_FROM_SUBMODULE.format(lang=lang_name, attr=attr_name)
+                    )
+                    loader: LanguageLoader = getattr(module, attr_name)
+                    return loader
+
+            logger.debug(
+                ls.NO_LANG_ATTR.format(module=module_name, available=dir(module))
+            )
+
+        finally:
+            if python_bindings_str in sys.path:
+                sys.path.remove(python_bindings_str)
+
+    except Exception as e:
+        logger.debug(ls.SUBMODULE_LOAD_FAILED.format(lang=lang_name, error=e))
+
+    return None
+
+
+def _try_import_language(
+    module_path: str, attr_name: str, lang_name: cs.SupportedLanguage
+) -> LanguageLoader:
+    try:
+        module = importlib.import_module(module_path)
+        loader: LanguageLoader = getattr(module, attr_name)
+        return loader
+    except ImportError:
+        return _try_load_from_submodule(lang_name)
+
+
+def _import_language_loaders() -> dict[cs.SupportedLanguage, LanguageLoader]:
+    language_imports: list[LanguageImport] = [
+        LanguageImport(
+            cs.SupportedLanguage.PYTHON,
+            cs.TreeSitterModule.PYTHON,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.PYTHON,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.JS,
+            cs.TreeSitterModule.JS,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.JS,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.TS,
+            cs.TreeSitterModule.TS,
+            cs.LANG_ATTR_TYPESCRIPT,
+            cs.SupportedLanguage.TS,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.RUST,
+            cs.TreeSitterModule.RUST,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.RUST,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.GO,
+            cs.TreeSitterModule.GO,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.GO,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.SCALA,
+            cs.TreeSitterModule.SCALA,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.SCALA,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.JAVA,
+            cs.TreeSitterModule.JAVA,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.JAVA,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.CPP,
+            cs.TreeSitterModule.CPP,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.CPP,
+        ),
+        LanguageImport(
+            cs.SupportedLanguage.LUA,
+            cs.TreeSitterModule.LUA,
+            cs.QUERY_LANGUAGE,
+            cs.SupportedLanguage.LUA,
+        ),
+    ]
+
+    loaders: dict[cs.SupportedLanguage, LanguageLoader] = {
+        lang_import.lang_key: _try_import_language(
+            lang_import.module_path,
+            lang_import.attr_name,
+            lang_import.submodule_name,
+        )
+        for lang_import in language_imports
+    }
+    for lang_key in LANGUAGE_SPECS:
+        lang_name = cs.SupportedLanguage(lang_key)
+        if lang_name not in loaders or loaders[lang_name] is None:
+            loaders[lang_name] = _try_load_from_submodule(lang_name)
+
+    return loaders
+
+
+_language_loaders = _import_language_loaders()
+
+LANGUAGE_LIBRARIES: dict[cs.SupportedLanguage, LanguageLoader] = _language_loaders
+
+
+def _build_query_pattern(node_types: tuple[str, ...], capture_name: str) -> str:
+    return " ".join([f"({node_type}) @{capture_name}" for node_type in node_types])
+
+
+def _get_locals_pattern(lang_name: cs.SupportedLanguage) -> str | None:
+    match lang_name:
+        case cs.SupportedLanguage.JS:
+            return cs.JS_LOCALS_PATTERN
+        case cs.SupportedLanguage.TS:
+            return cs.TS_LOCALS_PATTERN
+        case _:
+            return None
+
+
+def _build_combined_import_pattern(lang_config: LanguageSpec) -> str:
+    import_patterns = _build_query_pattern(
+        lang_config.import_node_types, cs.CAPTURE_IMPORT
+    )
+    import_from_patterns = _build_query_pattern(
+        lang_config.import_from_node_types, cs.CAPTURE_IMPORT_FROM
+    )
+
+    all_patterns: list[str] = []
+    if import_patterns.strip():
+        all_patterns.append(import_patterns)
+    if import_from_patterns.strip() and import_from_patterns != import_patterns:
+        all_patterns.append(import_from_patterns)
+    return " ".join(all_patterns)
+
+
+def _create_optional_query(language: Language, pattern: str | None) -> Query | None:
+    return Query(language, pattern) if pattern else None
+
+
+def _create_locals_query(
+    language: Language, lang_name: cs.SupportedLanguage
+) -> Query | None:
+    locals_pattern = _get_locals_pattern(lang_name)
+    if not locals_pattern:
+        return None
+    try:
+        return Query(language, locals_pattern)
+    except Exception as e:
+        logger.debug(ls.LOCALS_QUERY_FAILED.format(lang=lang_name, error=e))
+        return None
+
+
+def _create_language_queries(
+    language: Language,
+    parser: Parser,
+    lang_config: LanguageSpec,
+    lang_name: cs.SupportedLanguage,
+) -> LanguageQueries:
+    function_patterns = lang_config.function_query or _build_query_pattern(
+        lang_config.function_node_types, cs.CAPTURE_FUNCTION
+    )
+    class_patterns = lang_config.class_query or _build_query_pattern(
+        lang_config.class_node_types, cs.CAPTURE_CLASS
+    )
+    call_patterns = lang_config.call_query or _build_query_pattern(
+        lang_config.call_node_types, cs.CAPTURE_CALL
+    )
+    combined_import_patterns = _build_combined_import_pattern(lang_config)
+
+    return LanguageQueries(
+        functions=_create_optional_query(language, function_patterns),
+        classes=_create_optional_query(language, class_patterns),
+        calls=_create_optional_query(language, call_patterns),
+        imports=_create_optional_query(language, combined_import_patterns),
+        locals=_create_locals_query(language, lang_name),
+        config=lang_config,
+        language=language,
+        parser=parser,
+    )
+
+
+def _process_language(
+    lang_name: cs.SupportedLanguage,
+    lang_config: LanguageSpec,
+    parsers: dict[cs.SupportedLanguage, Parser],
+    queries: dict[cs.SupportedLanguage, LanguageQueries],
+) -> bool:
+    lang_lib = LANGUAGE_LIBRARIES.get(lang_name)
+    if not lang_lib:
+        logger.debug(ls.LIB_NOT_AVAILABLE.format(lang=lang_name))
+        return False
+
+    try:
+        language = Language(lang_lib())
+        parser = Parser(language)
+        parsers[lang_name] = parser
+        queries[lang_name] = _create_language_queries(
+            language, parser, lang_config, lang_name
+        )
+        logger.success(ls.GRAMMAR_LOADED.format(lang=lang_name))
+        return True
+    except Exception as e:
+        logger.warning(ls.GRAMMAR_LOAD_FAILED.format(lang=lang_name, error=e))
+        return False
+
+
+def load_parsers() -> tuple[
+    dict[cs.SupportedLanguage, Parser], dict[cs.SupportedLanguage, LanguageQueries]
+]:
+    parsers: dict[cs.SupportedLanguage, Parser] = {}
+    queries: dict[cs.SupportedLanguage, LanguageQueries] = {}
+    available_languages: list[cs.SupportedLanguage] = []
+
+    for lang_key, lang_config in deepcopy(LANGUAGE_SPECS).items():
+        lang_name = cs.SupportedLanguage(lang_key)
+        if _process_language(lang_name, lang_config, parsers, queries):
+            available_languages.append(lang_name)
+
+    if not available_languages:
+        raise RuntimeError(ex.NO_LANGUAGES)
+
+    logger.info(ls.INITIALIZED_PARSERS.format(languages=", ".join(available_languages)))
+    return parsers, queries

codebase_rag/prompts.py

@@ -0,0 +1,273 @@
+from typing import TYPE_CHECKING
+
+from .cypher_queries import (
+    CYPHER_EXAMPLE_CLASS_METHODS,
+    CYPHER_EXAMPLE_CONTENT_BY_PATH,
+    CYPHER_EXAMPLE_DECORATED_FUNCTIONS,
+    CYPHER_EXAMPLE_FILES_IN_FOLDER,
+    CYPHER_EXAMPLE_FIND_FILE,
+    CYPHER_EXAMPLE_KEYWORD_SEARCH,
+    CYPHER_EXAMPLE_LIMIT_ONE,
+    CYPHER_EXAMPLE_PYTHON_FILES,
+    CYPHER_EXAMPLE_README,
+    CYPHER_EXAMPLE_TASKS,
+)
+from .schema_builder import GRAPH_SCHEMA_DEFINITION
+from .types_defs import ToolNames
+
+if TYPE_CHECKING:
+    from pydantic_ai import Tool
+
+
+def extract_tool_names(tools: list["Tool"]) -> ToolNames:
+    tool_map = {t.name: t.name for t in tools}
+    return ToolNames(
+        query_graph=tool_map.get(
+            "query_codebase_knowledge_graph", "query_codebase_knowledge_graph"
+        ),
+        read_file=tool_map.get("read_file_content", "read_file_content"),
+        analyze_document=tool_map.get("analyze_document", "analyze_document"),
+        semantic_search=tool_map.get("semantic_code_search", "semantic_code_search"),
+        create_file=tool_map.get("create_new_file", "create_new_file"),
+        edit_file=tool_map.get("replace_code_surgically", "replace_code_surgically"),
+        shell_command=tool_map.get("execute_shell_command", "execute_shell_command"),
+    )
+
+
+CYPHER_QUERY_RULES = """**2. Critical Cypher Query Rules**
+
+- **ALWAYS Return Specific Properties with Aliases**: Do NOT return whole nodes (e.g., `RETURN n`). You MUST return specific properties with clear aliases (e.g., `RETURN n.name AS name`).
+- **Use `STARTS WITH` for Paths**: When matching paths, always use `STARTS WITH` for robustness (e.g., `WHERE n.path STARTS WITH 'workflows/src'`). Do not use `=`.
+- **Use `ENDS WITH` for qualified_name**: The `qualified_name` property contains full paths like `'Project.folder.subfolder.ClassName'`. When users mention a class, function, or method by its short name (e.g., "VatManager"), use `ENDS WITH` to match: `WHERE c.qualified_name ENDS WITH '.VatManager'`. Do NOT use `{name: 'VatManager'}` equality matching.
+- **Use `toLower()` for Searches**: For case-insensitive searching on string properties, use `toLower()`.
+- **Querying Lists**: To check if a list property (like `decorators`) contains an item, use the `ANY` or `IN` clause (e.g., `WHERE 'flow' IN n.decorators`)."""
+
+
+def build_graph_schema_and_rules() -> str:
+    return f"""You are an expert AI assistant for analyzing codebases using a **hybrid retrieval system**: a **Memgraph knowledge graph** for structural queries and a **semantic code search engine** for intent-based discovery.
+
+**1. Graph Schema Definition**
+The database contains information about a codebase, structured with the following nodes and relationships.
+
+{GRAPH_SCHEMA_DEFINITION}
+
+{CYPHER_QUERY_RULES}
+"""
+
+
+GRAPH_SCHEMA_AND_RULES = build_graph_schema_and_rules()
+
+
+def build_rag_orchestrator_prompt(tools: list["Tool"]) -> str:
+    t = extract_tool_names(tools)
+    return f"""You are an expert AI assistant for analyzing codebases. Your answers are based **EXCLUSIVELY** on information retrieved using your tools.
+
+**CRITICAL RULES:**
+1.  **TOOL-ONLY ANSWERS**: You must ONLY use information from the tools provided. Do not use external knowledge.
+2.  **NATURAL LANGUAGE QUERIES**: When using the `{t.query_graph}` tool, ALWAYS use natural language questions. NEVER write Cypher queries directly - the tool will translate your natural language into the appropriate database query.
+3.  **HONESTY**: If a tool fails or returns no results, you MUST state that clearly and report any error messages. Do not invent answers.
+4.  **CHOOSE THE RIGHT TOOL FOR THE FILE TYPE**:
+    - For source code files (.py, .ts, etc.), use `{t.read_file}`.
+    - For documents like PDFs, use the `{t.analyze_document}` tool. This is more effective than trying to read them as plain text.
+
+**Your General Approach:**
+1.  **Analyze Documents**: If the user asks a question about a document (like a PDF), you **MUST** use the `{t.analyze_document}` tool. Provide both the `file_path` and the user's `question` to the tool.
+2.  **Deep Dive into Code**: When you identify a relevant component (e.g., a folder), you must go beyond documentation.
+    a. First, check if documentation files like `README.md` exist and read them for context. For configuration, look for files appropriate to the language (e.g., `pyproject.toml` for Python, `package.json` for Node.js).
+    b. **Then, you MUST dive into the source code.** Explore the `src` directory (or equivalent). Identify and read key files (e.g., `main.py`, `index.ts`, `app.ts`) to understand the implementation details, logic, and functionality.
+    c. Synthesize all this information—from documentation, configuration, and the code itself—to provide a comprehensive, factual answer. Do not just describe the files; explain what the code *does*.
+    d. Only ask for clarification if, after a thorough investigation, the user's intent is still unclear.
+3.  **Choose the Right Search Strategy - SEMANTIC FIRST for Intent**:
+    a. **WHEN TO USE SEMANTIC SEARCH FIRST**: Always start with `{t.semantic_search}` for ANY of these patterns:
+       - "main entry point", "startup", "initialization", "bootstrap", "launcher"
+       - "error handling", "validation", "authentication"
+       - "where is X done", "how does Y work", "find Z logic"
+       - Any question about PURPOSE, INTENT, or FUNCTIONALITY
+
+       **Entry Point Recognition Patterns**:
+       - Python: `if __name__ == "__main__"`, `main()` function, CLI scripts, `app.run()`
+       - JavaScript/TypeScript: `index.js`, `main.ts`, `app.js`, `server.js`, package.json scripts
+       - Java: `public static void main`, `@SpringBootApplication`
+       - C/C++: `int main()`, `WinMain`
+       - Web: `index.html`, routing configurations, startup middleware
+
+    b. **WHEN TO USE GRAPH DIRECTLY**: Only use `{t.query_graph}` directly for pure structural queries:
+       - "What does function X call?" (when you already know X's name)
+       - "List methods of User class" (when you know the exact class name)
+       - "Show files in folder Y" (when you know the exact folder path)
+
+    c. **HYBRID APPROACH (RECOMMENDED)**: For most queries, use this sequence:
+       1. Use `{t.semantic_search}` to find relevant code elements by intent/meaning
+       2. Then use `{t.query_graph}` to explore structural relationships
+       3. **CRITICAL**: Always read the actual files using `{t.read_file}` to examine source code
+       4. For entry points specifically: Look for `if __name__ == "__main__"`, `main()` functions, or CLI entry points
+
+    d. **Tool Chaining Example**: For "main entry point and what it calls":
+       1. `{t.semantic_search}` for focused terms like "main entry startup" (not overly broad)
+       2. `{t.query_graph}` to find specific function relationships
+       3. `{t.read_file}` for main.py with targeted sections (use offset/limit for large files)
+       4. Look for the true application entry point (main function, __main__ block, CLI commands)
+       5. If you find CLI frameworks (typer, click, argparse), read relevant command sections only
+       6. Summarize execution flow concisely rather than showing all details
+4.  **Plan Before Writing or Modifying**:
+    a. Before using `{t.create_file}`, `{t.edit_file}`, or modifying files, you MUST explore the codebase to find the correct location and file structure.
+    b. For shell commands: If `{t.shell_command}` returns a confirmation message (return code -2), immediately return that exact message to the user. When they respond "yes", call the tool again with `user_confirmed=True`.
+5.  **Execute Shell Commands**: The `{t.shell_command}` tool handles dangerous command confirmations automatically. If it returns a confirmation prompt, pass it directly to the user.
+6.  **Complete the Investigation Cycle**: For entry point queries, you MUST:
+    a. Find candidate functions via semantic search
+    b. Explore their relationships via graph queries
+    c. **AUTOMATICALLY read main.py** (or main entry file) - NEVER ask the user for permission
+    d. Look for the ACTUAL startup code: `if __name__ == "__main__"`, CLI commands, `main()` functions
+    e. If CLI framework detected (typer, click, argparse), examine command functions
+    f. Distinguish between helper functions and the real application entry point
+    g. Show the complete execution flow from the true entry point through initialization
+7.  **Token Management**: Be efficient with context usage:
+    a. For semantic search, use focused queries (not overly broad terms)
+    b. For file reading, read specific sections when possible using offset/limit
+    c. Summarize large results rather than including full content
+    d. Prioritize most relevant findings over comprehensive coverage
+8.  **Synthesize Answer**: Analyze and explain the retrieved content. Cite your sources (file paths or qualified names). Report any errors gracefully.
+"""
+
+
+CYPHER_SYSTEM_PROMPT = f"""
+You are an expert translator that converts natural language questions about code structure into precise Neo4j Cypher queries.
+
+{GRAPH_SCHEMA_AND_RULES}
+
+**3. Query Optimization Rules**
+
+- **LIMIT Results**: ALWAYS add `LIMIT 50` to queries that list items. This prevents overwhelming responses.
+- **Aggregation Queries**: When asked "how many", "count", or "total", return ONLY the count, not all items:
+  - CORRECT: `MATCH (c:Class) RETURN count(c) AS total`
+  - WRONG: `MATCH (c:Class) RETURN c.name, c.path, count(c) AS total` (returns all items!)
+- **List vs Count**: If asked to "list" or "show", return items with LIMIT. If asked to "count" or "how many", return only the count.
+
+**4. Query Patterns & Examples**
+When listing items, return the `name`, `path`, and `qualified_name` with a LIMIT.
+
+**Pattern: Counting Items**
+cypher// "How many classes are there?" or "Count all functions"
+MATCH (c:Class) RETURN count(c) AS total
+
+**Pattern: Finding Decorated Functions/Methods (e.g., Workflows, Tasks)**
+cypher// "Find all prefect flows" or "what are the workflows?" or "show me the tasks"
+// Use the 'IN' operator to check the 'decorators' list property.
+{CYPHER_EXAMPLE_DECORATED_FUNCTIONS}
+
+**Pattern: Finding Content by Path (Robustly)**
+cypher// "what is in the 'workflows/src' directory?" or "list files in workflows"
+// Use `STARTS WITH` for path matching.
+{CYPHER_EXAMPLE_CONTENT_BY_PATH}
+
+**Pattern: Keyword & Concept Search (Fallback for general terms)**
+cypher// "find things related to 'database'"
+{CYPHER_EXAMPLE_KEYWORD_SEARCH}
+
+**Pattern: Finding a Specific File**
+cypher// "Find the main README.md"
+{CYPHER_EXAMPLE_FIND_FILE}
+
+**Pattern: Finding Methods of a Class by Short Name**
+cypher// "What methods does UserService have?" or "Show me methods in UserService" or "List UserService methods"
+// Use `ENDS WITH` to match the class by short name since qualified_name contains full path.
+{CYPHER_EXAMPLE_CLASS_METHODS}
+
+**4. Output Format**
+Provide only the Cypher query.
+"""
+
+# (H) Stricter prompt for less capable open-source/local models (e.g., Ollama)
+LOCAL_CYPHER_SYSTEM_PROMPT = f"""
+You are a Neo4j Cypher query generator. You ONLY respond with a valid Cypher query. Do not add explanations or markdown.
+
+{GRAPH_SCHEMA_AND_RULES}
+
+**CRITICAL RULES FOR QUERY GENERATION:**
+1.  **NO `UNION`**: Never use the `UNION` clause. Generate a single, simple `MATCH` query.
+2.  **BIND and ALIAS**: You must bind every node you use to a variable (e.g., `MATCH (f:File)`). You must use that variable to access properties and alias every returned property (e.g., `RETURN f.path AS path`).
+3.  **RETURN STRUCTURE**: Your query should aim to return `name`, `path`, and `qualified_name` so the calling system can use the results.
+    - For `File` nodes, return `f.path AS path`.
+    - For code nodes (`Class`, `Function`, etc.), return `n.qualified_name AS qualified_name`.
+4.  **KEEP IT SIMPLE**: Do not try to be clever. A simple query that returns a few relevant nodes is better than a complex one that fails.
+5.  **CLAUSE ORDER**: You MUST follow the standard Cypher clause order: `MATCH`, `WHERE`, `RETURN`, `LIMIT`.
+6.  **ALWAYS ADD LIMIT**: For queries that list items, ALWAYS add `LIMIT 50` to prevent overwhelming responses.
+7.  **AGGREGATION QUERIES**: When asked "how many" or "count", return ONLY the count:
+    - CORRECT: `MATCH (c:Class) RETURN count(c) AS total`
+    - WRONG: `MATCH (c:Class) RETURN c.name, count(c) AS total` (returns all items!)
+
+**Examples:**
+
+*   **Natural Language:** "How many classes are there?"
+*   **Cypher Query:**
+    ```cypher
+    MATCH (c:Class) RETURN count(c) AS total
+    ```
+
+*   **Natural Language:** "Find the main README file"
+*   **Cypher Query:**
+    ```cypher
+    {CYPHER_EXAMPLE_README}
+    ```
+
+*   **Natural Language:** "Find all python files"
+*   **Cypher Query (Note the '.' in extension):**
+    ```cypher
+    {CYPHER_EXAMPLE_PYTHON_FILES}
+    ```
+
+*   **Natural Language:** "show me the tasks"
+*   **Cypher Query:**
+    ```cypher
+    {CYPHER_EXAMPLE_TASKS}
+    ```
+
+*   **Natural Language:** "list files in the services folder"
+*   **Cypher Query:**
+    ```cypher
+    {CYPHER_EXAMPLE_FILES_IN_FOLDER}
+    ```
+
+*   **Natural Language:** "Find just one file to test"
+*   **Cypher Query:**
+    ```cypher
+    {CYPHER_EXAMPLE_LIMIT_ONE}
+    ```
+
+*   **Natural Language:** "What methods does UserService have?" or "Show me methods in UserService" or "List UserService methods"
+*   **Cypher Query (Use ENDS WITH to match class by short name):**
+    ```cypher
+    {CYPHER_EXAMPLE_CLASS_METHODS}
+    ```
+"""
+
+OPTIMIZATION_PROMPT = """
+I want you to analyze my {language} codebase and propose specific optimizations based on best practices.
+
+Please:
+1. Use your code retrieval and graph querying tools to understand the codebase structure
+2. Read relevant source files to identify optimization opportunities
+3. Reference established patterns and best practices for {language}
+4. Propose specific, actionable optimizations with file references
+5. IMPORTANT: Do not make any changes yet - just propose them and wait for approval
+6. After approval, use your file editing tools to implement the changes
+
+Start by analyzing the codebase structure and identifying the main areas that could benefit from optimization.
+Remember: Propose changes first, wait for my approval, then implement.
+"""
+
+OPTIMIZATION_PROMPT_WITH_REFERENCE = """
+I want you to analyze my {language} codebase and propose specific optimizations based on best practices.
+
+Please:
+1. Use your code retrieval and graph querying tools to understand the codebase structure
+2. Read relevant source files to identify optimization opportunities
+3. Use the analyze_document tool to reference best practices from {reference_document}
+4. Reference established patterns and best practices for {language}
+5. Propose specific, actionable optimizations with file references
+6. IMPORTANT: Do not make any changes yet - just propose them and wait for approval
+7. After approval, use your file editing tools to implement the changes
+
+Start by analyzing the codebase structure and identifying the main areas that could benefit from optimization.
+Remember: Propose changes first, wait for my approval, then implement.
+"""