shotgun-sh 0.1.0__py3-none-any.whl

shotgun/codebase/core/nl_query.py

@@ -0,0 +1,331 @@
+"""Natural language to Cypher query conversion for code graphs."""
+
+import time
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from pydantic_ai.messages import (
+    ModelRequest,
+    SystemPromptPart,
+    TextPart,
+    UserPromptPart,
+)
+
+from shotgun.agents.config import get_provider_model
+from shotgun.agents.config.models import shotgun_model_request
+from shotgun.logging_config import get_logger
+from shotgun.prompts import PromptLoader
+
+if TYPE_CHECKING:
+    from openai import AsyncOpenAI
+
+logger = get_logger(__name__)
+
+# Global prompt loader instance
+prompt_loader = PromptLoader()
+
+
+async def llm_cypher_prompt(system_prompt: str, user_prompt: str) -> str:
+    """Generate a Cypher query from a natural language prompt using the configured LLM provider.
+
+    Args:
+        system_prompt: The system prompt defining the behavior and context for the LLM
+        user_prompt: The user's natural language query
+    Returns:
+        The generated Cypher query as a string
+    """
+    model_config = get_provider_model()
+    # Use shotgun wrapper to maximize response quality for codebase queries
+    # Limit max_tokens to 2000 for Cypher queries (they're typically 50-200 tokens)
+    # This prevents Anthropic SDK from requiring streaming for longer token limits
+    query_cypher_response = await shotgun_model_request(
+        model_config=model_config,
+        messages=[
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(content=system_prompt),
+                    UserPromptPart(content=user_prompt),
+                ]
+            ),
+        ],
+        max_tokens=2000,  # Cypher queries are short, 2000 tokens is plenty
+    )
+
+    if not query_cypher_response.parts or not query_cypher_response.parts[0]:
+        raise ValueError("Empty response from LLM")
+
+    message_part = query_cypher_response.parts[0]
+    if not isinstance(message_part, TextPart):
+        raise ValueError("Unexpected response part type from LLM")
+    cypher_query = str(message_part.content)
+    if not cypher_query:
+        raise ValueError("Empty content in LLM response")
+    return cypher_query
+
+
+async def generate_cypher(natural_language_query: str) -> str:
+    """Convert a natural language query to Cypher using Shotgun's LLM client.
+
+    Args:
+        natural_language_query: The user's query in natural language
+
+    Returns:
+        Generated Cypher query
+    """
+    # Get current time for context
+    current_timestamp = int(time.time())
+    current_datetime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    # Generate system prompt using template
+    system_prompt = prompt_loader.render("codebase/cypher_system.j2")
+
+    # Generate enhanced query using template
+    enhanced_query = prompt_loader.render(
+        "codebase/enhanced_query_context.j2",
+        current_datetime=current_datetime,
+        current_timestamp=current_timestamp,
+        natural_language_query=natural_language_query,
+    )
+
+    try:
+        cypher_query = await llm_cypher_prompt(system_prompt, enhanced_query)
+        cleaned_query = clean_cypher_response(cypher_query)
+
+        # Validate UNION ALL queries
+        is_valid, validation_error = validate_union_query(cleaned_query)
+        if not is_valid:
+            logger.warning(f"Generated query failed validation: {validation_error}")
+            logger.warning(f"Problematic query: {cleaned_query}")
+            raise ValueError(f"Generated query validation failed: {validation_error}")
+
+        return cleaned_query
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to generate Cypher query: {e}") from e
+
+
+async def generate_cypher_with_error_context(
+    natural_language_query: str, error_context: str = ""
+) -> str:
+    """Convert a natural language query to Cypher with additional error context for retry scenarios.
+
+    Args:
+        natural_language_query: The user's query in natural language
+        error_context: Additional context about previous errors to help generate better query
+
+    Returns:
+        Generated Cypher query
+    """
+    # Get current time for context
+    current_timestamp = int(time.time())
+    current_datetime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    # Generate enhanced query with error context using template
+    enhanced_query = prompt_loader.render_string(
+        """Current datetime: {{ current_datetime }} (Unix timestamp: {{ current_timestamp }})
+
+User query: {{ natural_language_query }}
+
+ERROR CONTEXT (CRITICAL - Previous attempt failed):
+{{ error_context }}
+
+IMPORTANT: All timestamps in the database are stored as Unix timestamps (INT64). When generating time-based queries:
+- For "2 minutes ago": use {{ current_timestamp - 120 }}
+- For "1 hour ago": use {{ current_timestamp - 3600 }}
+- For "today": use timestamps >= {{ current_timestamp - (current_timestamp % 86400) }}
+- For "yesterday": use timestamps between {{ current_timestamp - 86400 - (current_timestamp % 86400) }} and {{ current_timestamp - (current_timestamp % 86400) }}
+- NEVER use placeholder values like 1704067200, always calculate based on the current timestamp: {{ current_timestamp }}""",
+        current_datetime=current_datetime,
+        current_timestamp=current_timestamp,
+        natural_language_query=natural_language_query,
+        error_context=error_context,
+    )
+
+    try:
+        # Create enhanced system prompt with error recovery instructions
+        enhanced_system_prompt = prompt_loader.render_string(
+            """{{ base_system_prompt }}
+
+**CRITICAL ERROR RECOVERY INSTRUCTIONS:**
+When retrying after a UNION ALL error:
+1. Each UNION ALL branch MUST return exactly the same number of columns
+2. Column names MUST be in the same order across all branches
+3. Use explicit column aliases to ensure consistency: RETURN prop1 as name, prop2 as qualified_name, 'Type' as type
+4. If different node types have different properties, use COALESCE or NULL for missing properties
+5. Test each UNION branch separately before combining
+
+Example of CORRECT UNION ALL:
+```cypher
+MATCH (c:Class) RETURN c.name as name, c.qualified_name as qualified_name, 'Class' as type
+UNION ALL
+MATCH (f:Function) RETURN f.name as name, f.qualified_name as qualified_name, 'Function' as type
+```
+
+Example of INCORRECT UNION ALL (different column counts):
+```cypher
+MATCH (c:Class) RETURN c.name, c.qualified_name, c.docstring
+UNION ALL
+MATCH (f:Function) RETURN f.name, f.qualified_name  // WRONG: missing third column
+```""",
+            base_system_prompt=prompt_loader.render("codebase/cypher_system.j2"),
+        )
+
+        cypher_query = await llm_cypher_prompt(enhanced_system_prompt, enhanced_query)
+        cleaned_query = clean_cypher_response(cypher_query)
+
+        # Validate UNION ALL queries
+        is_valid, validation_error = validate_union_query(cleaned_query)
+        if not is_valid:
+            logger.warning(f"Retry query failed validation: {validation_error}")
+            logger.warning(f"Problematic retry query: {cleaned_query}")
+            raise ValueError(f"Retry query validation failed: {validation_error}")
+
+        return cleaned_query
+
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to generate Cypher query with error context: {e}"
+        ) from e
+
+
+async def generate_cypher_openai_async(
+    client: "AsyncOpenAI", natural_language_query: str, model: str = "gpt-4o"
+) -> str:
+    """Convert a natural language query to Cypher using async OpenAI client.
+
+    This function is for standalone usage without Shotgun's LLM infrastructure.
+
+    Args:
+        client: Async OpenAI client instance
+        natural_language_query: The user's query in natural language
+        model: OpenAI model to use (default: gpt-4o)
+
+    Returns:
+        Generated Cypher query
+    """
+    # Get current time for context
+    current_timestamp = int(time.time())
+    current_datetime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    # Generate system prompt using template
+    system_prompt = prompt_loader.render("codebase/cypher_system.j2")
+
+    # Generate enhanced query using template
+    enhanced_query = prompt_loader.render(
+        "codebase/enhanced_query_context.j2",
+        current_datetime=current_datetime,
+        current_timestamp=current_timestamp,
+        natural_language_query=natural_language_query,
+    )
+
+    try:
+        cypher_query = await llm_cypher_prompt(system_prompt, enhanced_query)
+        return clean_cypher_response(cypher_query)
+
+    except Exception as e:
+        logger.error(f"OpenAI API error: {e}")
+        raise RuntimeError(f"Failed to generate Cypher query: {e}") from e
+
+
+def validate_union_query(cypher_query: str) -> tuple[bool, str]:
+    """Validate that UNION ALL queries have matching column counts and names.
+
+    Args:
+        cypher_query: The Cypher query to validate
+
+    Returns:
+        Tuple of (is_valid, error_message)
+    """
+    query_upper = cypher_query.upper()
+    if "UNION ALL" not in query_upper:
+        return True, ""
+
+    # Split by UNION ALL and extract RETURN clauses
+    parts = query_upper.split("UNION ALL")
+    return_patterns = []
+
+    for i, part in enumerate(parts):
+        if "RETURN" not in part:
+            continue
+
+        # Extract the RETURN clause
+        return_start = part.rfind("RETURN")
+        return_clause = part[return_start + 6 :]  # Skip "RETURN "
+
+        # Stop at ORDER BY, LIMIT, or end of query
+        for stop_word in ["ORDER BY", "LIMIT", ";"]:
+            if stop_word in return_clause:
+                return_clause = return_clause.split(stop_word)[0]
+
+        # Parse columns (basic parsing - split by comma and handle AS aliases)
+        columns = []
+        for col in return_clause.split(","):
+            col = col.strip()
+            if " AS " in col:
+                # Extract the alias name after AS
+                alias = col.split(" AS ")[-1].strip()
+                columns.append(alias)
+            else:
+                # Use the column name as-is (simplified)
+                columns.append(col.strip())
+
+        return_patterns.append((i, columns))
+
+    # Check all parts have same number of columns
+    if len(return_patterns) < 2:
+        return True, ""
+
+    first_part, first_columns = return_patterns[0]
+    first_count = len(first_columns)
+
+    for part_idx, columns in return_patterns[1:]:
+        if len(columns) != first_count:
+            return (
+                False,
+                f"UNION ALL part {part_idx + 1} has {len(columns)} columns, expected {first_count}. First part columns: {first_columns}, this part: {columns}",
+            )
+
+    return True, ""
+
+
+def clean_cypher_response(response_text: str) -> str:
+    """Clean up common LLM formatting artifacts from a Cypher query.
+
+    Args:
+        response_text: Raw response from LLM
+
+    Returns:
+        Cleaned Cypher query
+    """
+    query = response_text.strip()
+
+    # Remove markdown code blocks
+    if query.startswith("```"):
+        lines = query.split("\n")
+        # Find the actual query content
+        start_idx = 0
+        end_idx = len(lines)
+
+        for i, line in enumerate(lines):
+            if line.startswith("```") and i == 0:
+                start_idx = 1
+            elif line.startswith("```") and i > 0:
+                end_idx = i
+                break
+
+        query = "\n".join(lines[start_idx:end_idx])
+
+    # Remove 'cypher' prefix if present
+    query = query.strip()
+    if query.lower().startswith("cypher"):
+        query = query[6:].strip()
+
+    # Remove backticks
+    query = query.replace("`", "")
+
+    # Ensure it ends with semicolon
+    query = query.strip()
+    if not query.endswith(";"):
+        query += ";"
+
+    return query

shotgun/codebase/core/parser_loader.py

@@ -0,0 +1,128 @@
+"""Tree-sitter parser loader for code parsing."""
+
+import sys
+from collections.abc import Callable
+from typing import Any
+
+from tree_sitter import Language, Parser
+
+from shotgun.codebase.core.language_config import LANGUAGE_CONFIGS
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+def load_parsers() -> tuple[dict[str, Parser], dict[str, Any]]:
+    """Load available Tree-sitter parsers and compile their queries.
+
+    Returns:
+        Tuple of (parsers dict, queries dict)
+    """
+    parsers: dict[str, Parser] = {}
+    queries: dict[str, Any] = {}
+    available_languages = []
+
+    # Try to import available language libraries
+    language_loaders: dict[str, Callable[[], Any]] = {}
+
+    # Try individual language imports first
+    try:
+        import tree_sitter_python
+
+        language_loaders["python"] = lambda: tree_sitter_python.language()
+        available_languages.append("python")
+    except ImportError as e:
+        logger.warning(f"Failed to import tree_sitter_python: {e}")
+
+    try:
+        import tree_sitter_javascript
+
+        language_loaders["javascript"] = lambda: tree_sitter_javascript.language()
+        available_languages.append("javascript")
+    except ImportError as e:
+        logger.warning(f"Failed to import tree_sitter_javascript: {e}")
+
+    try:
+        import tree_sitter_typescript
+
+        language_loaders["typescript"] = (
+            lambda: tree_sitter_typescript.language_typescript()
+        )
+        available_languages.append("typescript")
+    except ImportError as e:
+        logger.warning(f"Failed to import tree_sitter_typescript: {e}")
+
+    try:
+        import tree_sitter_go
+
+        language_loaders["go"] = lambda: tree_sitter_go.language()
+        available_languages.append("go")
+    except ImportError as e:
+        logger.warning(f"Failed to import tree_sitter_go: {e}")
+
+    try:
+        import tree_sitter_rust
+
+        language_loaders["rust"] = lambda: tree_sitter_rust.language()
+        available_languages.append("rust")
+    except ImportError as e:
+        logger.warning(f"Failed to import tree_sitter_rust: {e}")
+
+    logger.info(f"Available languages: {', '.join(available_languages)}")
+
+    # Create parsers for available languages
+    for lang_name, lang_loader in language_loaders.items():
+        if lang_name in LANGUAGE_CONFIGS:
+            try:
+                parser = Parser()
+                # Handle both function and direct language object
+                if callable(lang_loader):
+                    lang_obj = lang_loader()
+                else:
+                    lang_obj = lang_loader
+
+                # Create Language object if needed
+                if not isinstance(lang_obj, Language):
+                    lang_obj = Language(lang_obj)
+
+                parser.language = lang_obj
+                parsers[lang_name] = parser
+
+                # Compile queries for this language
+                config = LANGUAGE_CONFIGS[lang_name]
+                lang_queries = {}
+
+                # Compile each query type
+                for query_type in [
+                    "function_query",
+                    "class_query",
+                    "call_query",
+                    "import_query",
+                ]:
+                    query_text = getattr(config, query_type)
+                    if query_text:
+                        try:
+                            lang_queries[query_type] = lang_obj.query(query_text)
+                        except Exception as e:
+                            logger.debug(
+                                f"Failed to compile {query_type} for {lang_name}: {e}"
+                            )
+
+                if lang_queries:
+                    queries[lang_name] = lang_queries
+
+                logger.debug(f"Loaded parser for {lang_name}")
+
+            except Exception as e:
+                logger.error(f"Failed to load parser for {lang_name}: {e}")
+
+    if not parsers:
+        logger.error(
+            "No parsers could be loaded. Please install language-specific tree-sitter packages."
+        )
+        logger.error(
+            "Install with: pip install tree-sitter-python tree-sitter-javascript tree-sitter-typescript tree-sitter-go tree-sitter-rust"
+        )
+        sys.exit(1)
+
+    return parsers, queries

shotgun/codebase/models.py

@@ -0,0 +1,111 @@
+"""Data models for codebase service."""
+
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class GraphStatus(str, Enum):
+    """Status of a code knowledge graph."""
+
+    READY = "READY"  # Graph is ready for queries
+    BUILDING = "BUILDING"  # Initial build in progress
+    UPDATING = "UPDATING"  # Update in progress
+    ERROR = "ERROR"  # Last operation failed
+
+
+class QueryType(str, Enum):
+    """Type of query being executed."""
+
+    NATURAL_LANGUAGE = "natural_language"
+    CYPHER = "cypher"
+
+
+class OperationStats(BaseModel):
+    """Statistics for a graph operation (build/update)."""
+
+    operation_type: str = Field(..., description="Type of operation: build or update")
+    started_at: float = Field(..., description="Unix timestamp when operation started")
+    completed_at: float | None = Field(
+        None, description="Unix timestamp when operation completed"
+    )
+    success: bool = Field(default=False, description="Whether operation succeeded")
+    error: str | None = Field(None, description="Error message if operation failed")
+    stats: dict[str, Any] = Field(
+        default_factory=dict, description="Operation statistics"
+    )
+
+
+class CodebaseGraph(BaseModel):
+    """Represents a code knowledge graph."""
+
+    graph_id: str = Field(..., description="Unique graph ID (hash of repo path)")
+    repo_path: str = Field(..., description="Absolute path to repository")
+    graph_path: str = Field(..., description="Path to Kuzu database")
+    name: str = Field(..., description="Human-readable name for the graph")
+    created_at: float = Field(..., description="Unix timestamp of creation")
+    updated_at: float = Field(..., description="Unix timestamp of last update")
+    schema_version: str = Field(default="1.0.0", description="Graph schema version")
+    build_options: dict[str, Any] = Field(
+        default_factory=dict, description="Build configuration"
+    )
+    language_stats: dict[str, int] = Field(
+        default_factory=dict, description="File count by language"
+    )
+    node_count: int = Field(default=0, description="Total number of nodes")
+    relationship_count: int = Field(
+        default=0, description="Total number of relationships"
+    )
+    node_stats: dict[str, int] = Field(
+        default_factory=dict, description="Node counts by type"
+    )
+    relationship_stats: dict[str, int] = Field(
+        default_factory=dict, description="Relationship counts by type"
+    )
+    is_watching: bool = Field(default=False, description="Whether watcher is active")
+    status: GraphStatus = Field(
+        default=GraphStatus.READY, description="Current status of the graph"
+    )
+    last_operation: OperationStats | None = Field(
+        None, description="Statistics from the last operation"
+    )
+    current_operation_id: str | None = Field(
+        None, description="ID of current in-progress operation"
+    )
+    indexed_from_cwds: list[str] = Field(
+        default_factory=list,
+        description="List of working directories from which this graph is accessible. Empty list means globally accessible.",
+    )
+
+
+class QueryResult(BaseModel):
+    """Result of a Cypher query execution."""
+
+    query: str = Field(..., description="Original query (natural language or Cypher)")
+    cypher_query: str | None = Field(
+        None, description="Generated Cypher query if from natural language"
+    )
+    results: list[dict[str, Any]] = Field(
+        default_factory=list, description="Query results"
+    )
+    column_names: list[str] = Field(
+        default_factory=list, description="Result column names"
+    )
+    row_count: int = Field(default=0, description="Number of result rows")
+    execution_time_ms: float = Field(
+        ..., description="Query execution time in milliseconds"
+    )
+    success: bool = Field(default=True, description="Whether query succeeded")
+    error: str | None = Field(None, description="Error message if failed")
+
+
+class FileChange(BaseModel):
+    """Represents a file system change."""
+
+    event_type: str = Field(
+        ..., description="Type of change: created, modified, deleted, moved"
+    )
+    src_path: str = Field(..., description="Source file path")
+    dest_path: str | None = Field(None, description="Destination path for moves")
+    is_directory: bool = Field(default=False, description="Whether path is a directory")