shotgun-sh 0.1.0.dev31__tar.gz → 0.1.1__tar.gz

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shotgun-sh
-Version: 0.1.0.dev31
+Version: 0.1.1
 Summary: AI-powered research, planning, and task management CLI tool
 Project-URL: Homepage, https://shotgun.sh/
 Project-URL: Repository, https://github.com/shotgun-sh/shotgun

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "shotgun-sh"
-version = "0.1.0.dev31"
+version = "0.1.1"
 description = "AI-powered research, planning, and task management CLI tool"
 readme = "README.md"
 license = { text = "MIT" }

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/agents/agent_manager.py

@@ -76,20 +76,25 @@ class MessageHistoryUpdated(Message):
 class PartialResponseMessage(Message):
     """Event posted when a partial response is received."""
 
-    def __init__(self, message: ModelResponse | None, is_last: bool) -> None:
+    def __init__(
+        self,
+        message: ModelResponse | None,
+        messages: list[ModelMessage],
+        is_last: bool,
+    ) -> None:
         """Initialize the partial response message."""
         super().__init__()
         self.message = message
+        self.messages = messages
         self.is_last = is_last
 
 
 @dataclass(slots=True)
 class _PartialStreamState:
-    """Tracks partial response parts while streaming a single agent run."""
+    """Tracks streamed messages while handling a single agent run."""
 
-    parts: list[ModelResponsePart | ToolCallPartDelta] = field(default_factory=list)
-    latest_partial: ModelResponse | None = None
-    final_sent: bool = False
+    messages: list[ModelRequest | ModelResponse] = field(default_factory=list)
+    current_response: ModelResponse | None = None
 
 
 class AgentManager(Widget):
@@ -272,6 +277,7 @@ class AgentManager(Widget):
 
         # Clear file tracker before each run to track only this run's operations
         deps.file_tracker.clear()
+        original_messages = self.ui_message_history.copy()
 
         if prompt:
             self.ui_message_history.append(ModelRequest.user_text_prompt(prompt))
@@ -356,16 +362,11 @@ class AgentManager(Widget):
                 **kwargs,
             )
         finally:
-            # If the stream ended unexpectedly without a final result, clear accumulated state.
-            if self._stream_state is not None and not self._stream_state.final_sent:
-                partial_message = self._build_partial_response(self._stream_state.parts)
-                if partial_message is not None:
-                    self._post_partial_message(partial_message, True)
             self._stream_state = None
 
-        self.ui_message_history = self.ui_message_history + [
-            mes for mes in result.new_messages() if not isinstance(mes, ModelRequest)
-        ]
+        self.ui_message_history = original_messages + cast(
+            list[ModelRequest | ModelResponse | HintMessage], result.new_messages()
+        )
 
         # Apply compaction to persistent message history to prevent cascading growth
         all_messages = result.all_messages()
@@ -390,7 +391,13 @@ class AgentManager(Widget):
         if state is None:
             state = self._stream_state = _PartialStreamState()
 
-        partial_parts = state.parts
+        if state.current_response is not None:
+            partial_parts: list[ModelResponsePart | ToolCallPartDelta] = list(
+                state.current_response.parts
+                # cast(Sequence[ModelResponsePart], state.current_response.parts)
+            )
+        else:
+            partial_parts = []
 
         async for event in stream:
             try:
@@ -409,8 +416,8 @@ class AgentManager(Widget):
 
                     partial_message = self._build_partial_response(partial_parts)
                     if partial_message is not None:
-                        state.latest_partial = partial_message
-                        self._post_partial_message(partial_message, False)
+                        state.current_response = partial_message
+                        self._post_partial_message(False)
 
                 elif isinstance(event, PartDeltaEvent):
                     index = event.index
@@ -435,8 +442,8 @@ class AgentManager(Widget):
 
                     partial_message = self._build_partial_response(partial_parts)
                     if partial_message is not None:
-                        state.latest_partial = partial_message
-                        self._post_partial_message(partial_message, False)
+                        state.current_response = partial_message
+                        self._post_partial_message(False)
 
                 elif isinstance(event, FunctionToolCallEvent):
                     existing_call_idx = next(
@@ -448,29 +455,54 @@ class AgentManager(Widget):
                         ),
                         None,
                     )
+
                     if existing_call_idx is not None:
                         partial_parts[existing_call_idx] = event.part
+                    elif state.messages:
+                        existing_call_idx = next(
+                            (
+                                i
+                                for i, part in enumerate(state.messages[-1].parts)
+                                if isinstance(part, ToolCallPart)
+                                and part.tool_call_id == event.part.tool_call_id
+                            ),
+                            None,
+                        )
                     else:
                         partial_parts.append(event.part)
                     partial_message = self._build_partial_response(partial_parts)
                     if partial_message is not None:
-                        state.latest_partial = partial_message
-                        self._post_partial_message(partial_message, False)
+                        state.current_response = partial_message
+                        self._post_partial_message(False)
                 elif isinstance(event, FunctionToolResultEvent):
-                    self.ui_message_history.append(ModelRequest(parts=[event.result]))
-                    self._post_messages_updated()  ## this is what the user responded with
-                elif isinstance(event, FinalResultEvent):
-                    final_message = (
-                        state.latest_partial
-                        or self._build_partial_response(partial_parts)
-                    )
-                    self._post_partial_message(final_message, False)
+                    request_message = ModelRequest(parts=[event.result])
+                    state.messages.append(request_message)
+                    if (
+                        event.result.tool_name == "ask_user"
+                    ):  # special handling to ask_user, because deferred tool results mean we missed the user response
+                        self.ui_message_history.append(request_message)
+                        self._post_messages_updated()
+                    ## this is what the user responded with
+                    self._post_partial_message(is_last=False)
 
+                elif isinstance(event, FinalResultEvent):
+                    pass
             except Exception:  # pragma: no cover - defensive logging
                 logger.exception(
                     "Error while handling agent stream event", extra={"event": event}
                 )
 
+        final_message = state.current_response or self._build_partial_response(
+            partial_parts
+        )
+        if final_message is not None:
+            state.current_response = final_message
+            if final_message not in state.messages:
+                state.messages.append(final_message)
+            state.current_response = None
+            self._post_partial_message(True)
+        state.current_response = None
+
     def _build_partial_response(
         self, parts: list[ModelResponsePart | ToolCallPartDelta]
     ) -> ModelResponse | None:
@@ -483,11 +515,20 @@ class AgentManager(Widget):
             return None
         return ModelResponse(parts=list(completed_parts))
 
-    def _post_partial_message(
-        self, message: ModelResponse | None, is_last: bool
-    ) -> None:
+    def _post_partial_message(self, is_last: bool) -> None:
         """Post a partial message to the UI."""
-        self.post_message(PartialResponseMessage(message, is_last))
+        if self._stream_state is None:
+            return
+        self.post_message(
+            PartialResponseMessage(
+                self._stream_state.current_response
+                if self._stream_state.current_response
+                not in self._stream_state.messages
+                else None,
+                self._stream_state.messages,
+                is_last,
+            )
+        )
 
     def _post_messages_updated(
         self, file_operations: list[FileOperation] | None = None

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/agents/config/constants.py

@@ -2,7 +2,6 @@
 
 # Field names
 API_KEY_FIELD = "api_key"
-MODEL_NAME_FIELD = "model_name"
 DEFAULT_PROVIDER_FIELD = "default_provider"
 USER_ID_FIELD = "user_id"
 CONFIG_VERSION_FIELD = "config_version"

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/agents/config/models.py

@@ -117,21 +117,18 @@ class OpenAIConfig(BaseModel):
     """Configuration for OpenAI provider."""
 
     api_key: SecretStr | None = None
-    model_name: str = "gpt-5"
 
 
 class AnthropicConfig(BaseModel):
     """Configuration for Anthropic provider."""
 
     api_key: SecretStr | None = None
-    model_name: str = "claude-opus-4-1"
 
 
 class GoogleConfig(BaseModel):
     """Configuration for Google provider."""
 
     api_key: SecretStr | None = None
-    model_name: str = "gemini-2.5-pro"
 
 
 class ShotgunConfig(BaseModel):

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/agents/config/provider.py

@@ -127,8 +127,8 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
                 f"OpenAI API key not configured. Set via environment variable {OPENAI_API_KEY_ENV} or config."
             )
 
-        # Get model spec
-        model_name = config.openai.model_name
+        # Get model spec - hardcoded to gpt-5
+        model_name = "gpt-5"
         if model_name not in MODEL_SPECS:
             raise ValueError(f"Model '{model_name}' not found")
         spec = MODEL_SPECS[model_name]
@@ -149,8 +149,8 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
                 f"Anthropic API key not configured. Set via environment variable {ANTHROPIC_API_KEY_ENV} or config."
             )
 
-        # Get model spec
-        model_name = config.anthropic.model_name
+        # Get model spec - hardcoded to claude-opus-4-1
+        model_name = "claude-opus-4-1"
         if model_name not in MODEL_SPECS:
             raise ValueError(f"Model '{model_name}' not found")
         spec = MODEL_SPECS[model_name]
@@ -171,8 +171,8 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
                 f"Gemini API key not configured. Set via environment variable {GEMINI_API_KEY_ENV} or config."
             )
 
-        # Get model spec
-        model_name = config.google.model_name
+        # Get model spec - hardcoded to gemini-2.5-pro
+        model_name = "gemini-2.5-pro"
         if model_name not in MODEL_SPECS:
             raise ValueError(f"Model '{model_name}' not found")
         spec = MODEL_SPECS[model_name]

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/build_constants.py

@@ -12,8 +12,8 @@ POSTHOG_API_KEY = ''
 POSTHOG_PROJECT_ID = '191396'
 
 # Logfire configuration embedded at build time (only for dev builds)
-LOGFIRE_ENABLED = 'true'
-LOGFIRE_TOKEN = 'pylf_v1_us_KZ5NM1pP3NwgJkbBJt6Ftdzk8mMhmrXcGJHQQgDJ1LfK'
+LOGFIRE_ENABLED = ''
+LOGFIRE_TOKEN = ''
 
 # Build metadata
 BUILD_TIME_ENV = "production" if SENTRY_DSN else "development"

shotgun_sh-0.1.1/src/shotgun/codebase/core/cypher_models.py

@@ -0,0 +1,46 @@
+"""Pydantic models and exceptions for Cypher query generation."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class CypherGenerationResponse(BaseModel):
+    """Structured response from LLM for Cypher query generation.
+
+    This model ensures the LLM explicitly indicates whether it can generate
+    a valid Cypher query and provides a reason if it cannot.
+    """
+
+    cypher_query: str | None = Field(
+        default=None,
+        description="The generated Cypher query, or None if generation not possible",
+    )
+    can_generate_valid_cypher: bool = Field(
+        description="Whether a valid Cypher query can be generated for this request"
+    )
+    reason_cannot_generate: str | None = Field(
+        default=None,
+        description="Explanation why query cannot be generated (if applicable)",
+    )
+
+    def model_post_init(self, __context: Any) -> None:
+        """Validate that reason is provided when query cannot be generated."""
+        if not self.can_generate_valid_cypher and not self.reason_cannot_generate:
+            self.reason_cannot_generate = "No reason provided"
+        if self.can_generate_valid_cypher and not self.cypher_query:
+            raise ValueError(
+                "cypher_query must be provided when can_generate_valid_cypher is True"
+            )
+
+
+class CypherGenerationNotPossibleError(Exception):
+    """Raised when LLM cannot generate valid Cypher for the query.
+
+    This typically happens when the query is conceptual rather than structural,
+    or when it requires interpretation beyond what can be expressed in Cypher.
+    """
+
+    def __init__(self, reason: str):
+        self.reason = reason
+        super().__init__(f"Cannot generate Cypher query: {reason}")

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/codebase/core/nl_query.py

@@ -4,15 +4,13 @@ import time
 from datetime import datetime
 from typing import TYPE_CHECKING
 
-from pydantic_ai.messages import (
-    ModelRequest,
-    SystemPromptPart,
-    TextPart,
-    UserPromptPart,
-)
+from pydantic_ai import Agent
 
 from shotgun.agents.config import get_provider_model
-from shotgun.agents.config.models import shotgun_model_request
+from shotgun.codebase.core.cypher_models import (
+    CypherGenerationNotPossibleError,
+    CypherGenerationResponse,
+)
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
 
@@ -25,42 +23,52 @@ logger = get_logger(__name__)
 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.
+async def llm_cypher_prompt(
+    system_prompt: str, user_prompt: str
+) -> CypherGenerationResponse:
+    """Generate a Cypher query from a natural language prompt using structured output.
 
     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
+        CypherGenerationResponse with cypher_query, can_generate flag, and reason if not
     """
     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
+
+    # Create an agent with structured output for Cypher generation
+    cypher_agent = Agent(
+        model=model_config.model_instance,
+        output_type=CypherGenerationResponse,
+        retries=2,
     )
 
-    if not query_cypher_response.parts or not query_cypher_response.parts[0]:
-        raise ValueError("Empty response from LLM")
+    # Combine system and user prompts
+    combined_prompt = f"{system_prompt}\n\nUser Query: {user_prompt}"
+
+    try:
+        # Run the agent to get structured response
+        result = await cypher_agent.run(combined_prompt)
+        response = result.output
+
+        # Log the structured response for debugging
+        logger.debug(
+            "Cypher generation response - can_generate: %s, query: %s, reason: %s",
+            response.can_generate_valid_cypher,
+            response.cypher_query[:50] if response.cypher_query else None,
+            response.reason_cannot_generate,
+        )
+
+        return response
 
-    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
+    except Exception as e:
+        logger.error("Failed to generate Cypher query with structured output: %s", e)
+        # Return a failure response
+        return CypherGenerationResponse(
+            cypher_query=None,
+            can_generate_valid_cypher=False,
+            reason_cannot_generate=f"LLM error: {str(e)}",
+        )
 
 
 async def generate_cypher(natural_language_query: str) -> str:
@@ -71,6 +79,10 @@ async def generate_cypher(natural_language_query: str) -> str:
 
     Returns:
         Generated Cypher query
+
+    Raises:
+        CypherGenerationNotPossibleError: If the query cannot be converted to Cypher
+        RuntimeError: If there's an error during generation
     """
     # Get current time for context
     current_timestamp = int(time.time())
@@ -88,8 +100,30 @@ async def generate_cypher(natural_language_query: str) -> str:
     )
 
     try:
-        cypher_query = await llm_cypher_prompt(system_prompt, enhanced_query)
-        cleaned_query = clean_cypher_response(cypher_query)
+        response = await llm_cypher_prompt(system_prompt, enhanced_query)
+
+        # Check if the LLM could generate a valid Cypher query
+        if not response.can_generate_valid_cypher:
+            logger.info(
+                "Cannot generate Cypher for query '%s': %s",
+                natural_language_query,
+                response.reason_cannot_generate,
+            )
+            raise CypherGenerationNotPossibleError(
+                response.reason_cannot_generate or "Query cannot be converted to Cypher"
+            )
+
+        if not response.cypher_query:
+            raise ValueError("LLM indicated success but provided no query")
+
+        cleaned_query = clean_cypher_response(response.cypher_query)
+
+        # Validate Cypher keywords
+        is_valid, validation_error = validate_cypher_keywords(cleaned_query)
+        if not is_valid:
+            logger.warning(f"Generated query has invalid syntax: {validation_error}")
+            logger.warning(f"Problematic query: {cleaned_query}")
+            raise ValueError(f"Generated query validation failed: {validation_error}")
 
         # Validate UNION ALL queries
         is_valid, validation_error = validate_union_query(cleaned_query)
@@ -100,6 +134,8 @@ async def generate_cypher(natural_language_query: str) -> str:
 
         return cleaned_query
 
+    except CypherGenerationNotPossibleError:
+        raise  # Re-raise as-is
     except Exception as e:
         raise RuntimeError(f"Failed to generate Cypher query: {e}") from e
 
@@ -170,8 +206,31 @@ MATCH (f:Function) RETURN f.name, f.qualified_name  // WRONG: missing third colu
             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)
+        response = await llm_cypher_prompt(enhanced_system_prompt, enhanced_query)
+
+        # Check if the LLM could generate a valid Cypher query
+        if not response.can_generate_valid_cypher:
+            logger.info(
+                "Cannot generate Cypher for retry query '%s': %s",
+                natural_language_query,
+                response.reason_cannot_generate,
+            )
+            raise CypherGenerationNotPossibleError(
+                response.reason_cannot_generate
+                or "Query cannot be converted to Cypher even with error context"
+            )
+
+        if not response.cypher_query:
+            raise ValueError("LLM indicated success but provided no query on retry")
+
+        cleaned_query = clean_cypher_response(response.cypher_query)
+
+        # Validate Cypher keywords
+        is_valid, validation_error = validate_cypher_keywords(cleaned_query)
+        if not is_valid:
+            logger.warning(f"Generated query has invalid syntax: {validation_error}")
+            logger.warning(f"Problematic query: {cleaned_query}")
+            raise ValueError(f"Generated query validation failed: {validation_error}")
 
         # Validate UNION ALL queries
         is_valid, validation_error = validate_union_query(cleaned_query)
@@ -182,6 +241,8 @@ MATCH (f:Function) RETURN f.name, f.qualified_name  // WRONG: missing third colu
 
         return cleaned_query
 
+    except CypherGenerationNotPossibleError:
+        raise  # Re-raise as-is
     except Exception as e:
         raise RuntimeError(
             f"Failed to generate Cypher query with error context: {e}"
@@ -202,6 +263,10 @@ async def generate_cypher_openai_async(
 
     Returns:
         Generated Cypher query
+
+    Raises:
+        CypherGenerationNotPossibleError: If the query cannot be converted to Cypher
+        RuntimeError: If there's an error during generation
     """
     # Get current time for context
     current_timestamp = int(time.time())
@@ -219,9 +284,26 @@ async def generate_cypher_openai_async(
     )
 
     try:
-        cypher_query = await llm_cypher_prompt(system_prompt, enhanced_query)
-        return clean_cypher_response(cypher_query)
+        response = await llm_cypher_prompt(system_prompt, enhanced_query)
+
+        # Check if the LLM could generate a valid Cypher query
+        if not response.can_generate_valid_cypher:
+            logger.info(
+                "Cannot generate Cypher for query '%s': %s",
+                natural_language_query,
+                response.reason_cannot_generate,
+            )
+            raise CypherGenerationNotPossibleError(
+                response.reason_cannot_generate or "Query cannot be converted to Cypher"
+            )
 
+        if not response.cypher_query:
+            raise ValueError("LLM indicated success but provided no query")
+
+        return clean_cypher_response(response.cypher_query)
+
+    except CypherGenerationNotPossibleError:
+        raise  # Re-raise as-is
     except Exception as e:
         logger.error(f"OpenAI API error: {e}")
         raise RuntimeError(f"Failed to generate Cypher query: {e}") from e
@@ -288,6 +370,65 @@ def validate_union_query(cypher_query: str) -> tuple[bool, str]:
     return True, ""
 
 
+def validate_cypher_keywords(query: str) -> tuple[bool, str]:
+    """Validate that a query starts with valid Kuzu Cypher keywords.
+
+    Args:
+        query: The Cypher query to validate
+
+    Returns:
+        Tuple of (is_valid, error_message)
+    """
+    # Valid Kuzu Cypher starting keywords based on parser expectations
+    valid_cypher_keywords = {
+        "ALTER",
+        "ATTACH",
+        "BEGIN",
+        "CALL",
+        "CHECKPOINT",
+        "COMMENT",
+        "COMMIT",
+        "COPY",
+        "CREATE",
+        "DELETE",
+        "DETACH",
+        "DROP",
+        "EXPLAIN",
+        "EXPORT",
+        "FORCE",
+        "IMPORT",
+        "INSTALL",
+        "LOAD",
+        "MATCH",
+        "MERGE",
+        "OPTIONAL",
+        "PROFILE",
+        "RETURN",
+        "ROLLBACK",
+        "SET",
+        "UNWIND",
+        "UNINSTALL",
+        "UPDATE",
+        "USE",
+        "WITH",
+    }
+
+    query = query.strip()
+    if not query:
+        return False, "Empty query"
+
+    # Get the first word
+    first_word = query.upper().split()[0] if query else ""
+
+    if first_word not in valid_cypher_keywords:
+        return (
+            False,
+            f"Query doesn't start with valid Cypher keyword. Found: '{first_word}'",
+        )
+
+    return True, ""
+
+
 def clean_cypher_response(response_text: str) -> str:
     """Clean up common LLM formatting artifacts from a Cypher query.
 

{shotgun_sh-0.1.0.dev31 → shotgun_sh-0.1.1}/src/shotgun/codebase/service.py

@@ -4,6 +4,7 @@ import time
 from pathlib import Path
 from typing import Any
 
+from shotgun.codebase.core.cypher_models import CypherGenerationNotPossibleError
 from shotgun.codebase.core.manager import CodebaseGraphManager
 from shotgun.codebase.core.nl_query import generate_cypher
 from shotgun.codebase.models import CodebaseGraph, QueryResult, QueryType
@@ -190,6 +191,22 @@ class CodebaseService:
                 error=None,
             )
 
+        except CypherGenerationNotPossibleError as e:
+            # Handle queries that cannot be converted to Cypher
+            execution_time = (time.time() - start_time) * 1000
+            logger.info(f"Query cannot be converted to Cypher: {e.reason}")
+
+            return QueryResult(
+                query=query,
+                cypher_query=None,
+                results=[],
+                column_names=[],
+                row_count=0,
+                execution_time_ms=execution_time,
+                success=False,
+                error=f"This query cannot be converted to Cypher: {e.reason}",
+            )
+
         except Exception as e:
             execution_time = (time.time() - start_time) * 1000
             logger.error(f"Query execution failed: {e}")