shotgun-sh 0.1.1__tar.gz → 0.1.2__tar.gz

{shotgun_sh-0.1.1 → shotgun_sh-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shotgun-sh
-Version: 0.1.1
+Version: 0.1.2
 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.1 → shotgun_sh-0.1.2}/pyproject.toml

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

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

@@ -362,6 +362,17 @@ class AgentManager(Widget):
                 **kwargs,
             )
         finally:
+            # If the stream ended unexpectedly without a final result, clear accumulated state.
+            # state = self._stream_state
+            # if state is not None:
+            #     pending_response = state.current_response
+            #     if pending_response is not None:
+            #         already_recorded = (
+            #             bool(state.messages) and state.messages[-1] is pending_response
+            #         )
+            #         if not already_recorded:
+            #             self._post_partial_message(pending_response, True)
+            #             state.messages.append(pending_response)
             self._stream_state = None
 
         self.ui_message_history = original_messages + cast(

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

@@ -4,13 +4,15 @@ import time
 from datetime import datetime
 from typing import TYPE_CHECKING
 
-from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    ModelRequest,
+    SystemPromptPart,
+    TextPart,
+    UserPromptPart,
+)
 
 from shotgun.agents.config import get_provider_model
-from shotgun.codebase.core.cypher_models import (
-    CypherGenerationNotPossibleError,
-    CypherGenerationResponse,
-)
+from shotgun.agents.config.models import shotgun_model_request
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
 
@@ -23,52 +25,42 @@ logger = get_logger(__name__)
 prompt_loader = PromptLoader()
 
 
-async def llm_cypher_prompt(
-    system_prompt: str, user_prompt: str
-) -> CypherGenerationResponse:
-    """Generate a Cypher query from a natural language prompt using structured output.
+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:
-        CypherGenerationResponse with cypher_query, can_generate flag, and reason if not
+        The generated Cypher query as a string
     """
     model_config = get_provider_model()
-
-    # Create an agent with structured output for Cypher generation
-    cypher_agent = Agent(
-        model=model_config.model_instance,
-        output_type=CypherGenerationResponse,
-        retries=2,
+    # 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
     )
 
-    # 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
+    if not query_cypher_response.parts or not query_cypher_response.parts[0]:
+        raise ValueError("Empty response from LLM")
 
-    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)}",
-        )
+    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:
@@ -79,10 +71,6 @@ 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())
@@ -100,30 +88,8 @@ async def generate_cypher(natural_language_query: str) -> str:
     )
 
     try:
-        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}")
+        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)
@@ -134,8 +100,6 @@ 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
 
@@ -206,31 +170,8 @@ MATCH (f:Function) RETURN f.name, f.qualified_name  // WRONG: missing third colu
             base_system_prompt=prompt_loader.render("codebase/cypher_system.j2"),
         )
 
-        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}")
+        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)
@@ -241,8 +182,6 @@ 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}"
@@ -263,10 +202,6 @@ 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())
@@ -284,26 +219,9 @@ async def generate_cypher_openai_async(
     )
 
     try:
-        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"
-            )
+        cypher_query = await llm_cypher_prompt(system_prompt, enhanced_query)
+        return clean_cypher_response(cypher_query)
 
-        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
@@ -370,65 +288,6 @@ 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.1 → shotgun_sh-0.1.2}/src/shotgun/codebase/service.py

@@ -4,7 +4,6 @@ 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
@@ -191,22 +190,6 @@ 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}")

{shotgun_sh-0.1.1 → shotgun_sh-0.1.2}/src/shotgun/main.py

@@ -22,7 +22,7 @@ from shotgun.posthog_telemetry import setup_posthog_observability
 from shotgun.sentry_telemetry import setup_sentry_observability
 from shotgun.telemetry import setup_logfire_observability
 from shotgun.tui import app as tui_app
-from shotgun.utils.update_checker import check_for_updates_async
+from shotgun.utils.update_checker import check_and_install_updates_async
 
 # Load environment variables from .env file
 load_dotenv()
@@ -52,6 +52,7 @@ logger.debug("PostHog analytics enabled: %s", _posthog_enabled)
 
 # Global variable to store update notification
 _update_notification: str | None = None
+_update_progress: str | None = None
 
 
 def _update_callback(notification: str) -> None:
@@ -60,6 +61,13 @@ def _update_callback(notification: str) -> None:
     _update_notification = notification
 
 
+def _update_progress_callback(progress: str) -> None:
+    """Callback to store update progress."""
+    global _update_progress
+    _update_progress = progress
+    logger.debug(f"Update progress: {progress}")
+
+
 app = typer.Typer(
     name="shotgun",
     help="Shotgun - AI-powered CLI tool for research, planning, and task management",
@@ -121,10 +129,12 @@ def main(
     """Shotgun - AI-powered CLI tool."""
     logger.debug("Starting shotgun CLI application")
 
-    # Start async update check (non-blocking)
+    # Start async update check and install (non-blocking)
     if not ctx.resilient_parsing:
-        check_for_updates_async(
-            callback=_update_callback, no_update_check=no_update_check
+        check_and_install_updates_async(
+            callback=_update_callback,
+            no_update_check=no_update_check,
+            progress_callback=_update_progress_callback,
         )
 
     if ctx.invoked_subcommand is None and not ctx.resilient_parsing:

{shotgun_sh-0.1.1 → shotgun_sh-0.1.2}/src/shotgun/prompts/codebase/cypher_system.j2

@@ -25,18 +25,4 @@ Your goal is to return appropriate properties for each node type. Common propert
 {% include 'codebase/partials/temporal_context.j2' %}
 
 **6. Output Format**
-You must return a structured JSON response with the following fields:
-- `cypher_query`: The generated Cypher query string (or null if not possible)
-- `can_generate_valid_cypher`: Boolean indicating if a valid Cypher query can be generated
-- `reason_cannot_generate`: String explaining why generation isn't possible (or null if successful)
-
-**IMPORTANT:** Some queries cannot be expressed in Cypher:
-- Conceptual questions requiring interpretation (e.g., "What is the main purpose of this codebase?")
-- Questions about code quality or best practices
-- Questions requiring semantic understanding beyond structure
-
-For these, set `can_generate_valid_cypher` to false and provide a clear explanation in `reason_cannot_generate`.
-
-Examples:
-- Query: "Show all classes" → can_generate_valid_cypher: true, cypher_query: "MATCH (c:Class) RETURN c.name, c.qualified_name;"
-- Query: "What is the main purpose of this codebase?" → can_generate_valid_cypher: false, reason_cannot_generate: "This is a conceptual question requiring interpretation and analysis of the code's overall design and intent, rather than a structural query about specific code elements."
+Provide only the Cypher query.

{shotgun_sh-0.1.1 → shotgun_sh-0.1.2}/src/shotgun/tui/app.py

@@ -9,7 +9,7 @@ from shotgun.agents.config import ConfigManager, get_config_manager
 from shotgun.logging_config import get_logger
 from shotgun.tui.screens.splash import SplashScreen
 from shotgun.utils.file_system_utils import get_shotgun_base_path
-from shotgun.utils.update_checker import check_for_updates_async
+from shotgun.utils.update_checker import check_and_install_updates_async
 
 from .screens.chat import ChatScreen
 from .screens.directory_setup import DirectorySetupScreen
@@ -37,16 +37,26 @@ class ShotgunApp(App[None]):
         self.no_update_check = no_update_check
         self.continue_session = continue_session
         self.update_notification: str | None = None
+        self.update_progress: str | None = None
 
-        # Start async update check
+        # Start async update check and install
         if not no_update_check:
-            check_for_updates_async(callback=self._update_callback)
+            check_and_install_updates_async(
+                callback=self._update_callback,
+                no_update_check=no_update_check,
+                progress_callback=self._update_progress_callback,
+            )
 
     def _update_callback(self, notification: str) -> None:
         """Store update notification to show later."""
         self.update_notification = notification
         logger.debug(f"Update notification received: {notification}")
 
+    def _update_progress_callback(self, progress: str) -> None:
+        """Store update progress."""
+        self.update_progress = progress
+        logger.debug(f"Update progress: {progress}")
+
     def on_mount(self) -> None:
         self.theme = "gruvbox"
         # Track TUI startup

{shotgun_sh-0.1.1 → shotgun_sh-0.1.2}/src/shotgun/tui/screens/chat.py

@@ -411,7 +411,7 @@ class ChatScreen(Screen[None]):
             await self.codebase_sdk.list_codebases_for_directory()
         ).graphs
         if accessible_graphs:
-            self.mount_hint(help_text_with_codebase(already_indexed=True))
+            self.mount_hint(help_text_with_codebase())
             return
 
         should_index = await self.app.push_screen_wait(CodebaseIndexPromptScreen())
@@ -419,8 +419,6 @@ class ChatScreen(Screen[None]):
             self.mount_hint(help_text_empty_dir())
             return
 
-        self.mount_hint(help_text_with_codebase(already_indexed=False))
-
         self.index_codebase_command()
 
     def watch_mode(self, new_mode: AgentType) -> None:
@@ -694,6 +692,7 @@ class ChatScreen(Screen[None]):
                 timeout=8,
             )
 
+            self.mount_hint(codebase_indexed_hint(selection.name))
         except CodebaseAlreadyIndexedError as exc:
             logger.warning(f"Codebase already indexed: {exc}")
             self.notify(str(exc), severity="warning")
@@ -797,10 +796,17 @@ class ChatScreen(Screen[None]):
         self.mode = AgentType(conversation.last_agent_model)
 
 
-def help_text_with_codebase(already_indexed: bool = False) -> str:
+def codebase_indexed_hint(codebase_name: str) -> str:
+    return (
+        f"Codebase **{codebase_name}** indexed successfully. You can now use it in your chat.\n\n"
+        + help_text_with_codebase()
+    )
+
+
+def help_text_with_codebase() -> str:
     return (
         "Howdy! Welcome to Shotgun - the context tool for software engineering. \n\nYou can research, build specs, plan, create tasks, and export context to your favorite code-gen agents.\n\n"
-        f"{'' if already_indexed else 'Once your codebase is indexed, '}I can help with:\n\n"
+        "I can help with:\n\n"
         "- Speccing out a new feature\n"
         "- Onboarding you onto this project\n"
         "- Helping with a refactor spec\n"

{shotgun_sh-0.1.1 → shotgun_sh-0.1.2}/src/shotgun/tui/screens/chat_screen/history.py

@@ -136,8 +136,7 @@ class UserQuestionWidget(Widget):
                 if part.tool_name == "ask_user" and isinstance(part.content, dict):
                     acc += f"**>** {part.content['answer']}\n\n"
                 else:
-                    # acc += "  ∟ finished\n\n"  # let's not show anything yet
-                    pass
+                    acc += "  ∟ finished\n\n"  # let's not show anything yet
             elif isinstance(part, UserPromptPart):
                 acc += f"**>** {part.content}\n\n"
         return acc
@@ -153,7 +152,7 @@ class AgentResponseWidget(Widget):
         if self.item is None:
             yield Markdown(markdown="")
         else:
-            yield Markdown(markdown=self.compute_output())
+            yield Markdown(markdown=f"**⏺** {self.compute_output()}")
 
     def compute_output(self) -> str:
         acc = ""
@@ -161,10 +160,10 @@ class AgentResponseWidget(Widget):
             return ""
         for idx, part in enumerate(self.item.parts):
             if isinstance(part, TextPart):
-                acc += f"**⏺** {part.content}\n\n"
+                acc += f"{part.content}\n\n"
             elif isinstance(part, ToolCallPart):
                 parts_str = self._format_tool_call_part(part)
-                acc += parts_str + "\n\n"
+                acc += f"{part.tool_name}: " + parts_str + "\n\n"
             elif isinstance(part, BuiltinToolCallPart):
                 acc += f"{part.tool_name}({part.args})\n\n"
             elif isinstance(part, BuiltinToolReturnPart):