shotgun-sh 0.1.0.dev15__py3-none-any.whl → 0.1.0.dev16__py3-none-any.whl

shotgun/tui/screens/chat.py

@@ -1,27 +1,31 @@
-from collections.abc import AsyncGenerator
-from typing import cast
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
 
 from pydantic_ai import DeferredToolResults, RunContext
 from pydantic_ai.messages import (
-    BuiltinToolCallPart,
-    BuiltinToolReturnPart,
     ModelMessage,
     ModelRequest,
     ModelResponse,
     TextPart,
-    ThinkingPart,
-    ToolCallPart,
+    UserPromptPart,
 )
 from textual import on, work
 from textual.app import ComposeResult
-from textual.command import DiscoveryHit, Hit, Provider
-from textual.containers import Container
+from textual.command import CommandPalette
+from textual.containers import Container, Grid
 from textual.reactive import reactive
-from textual.screen import Screen
+from textual.screen import ModalScreen, Screen
 from textual.widget import Widget
-from textual.widgets import Markdown
+from textual.widgets import Button, DirectoryTree, Input, Label, Markdown, Static
 
-from shotgun.agents.agent_manager import AgentManager, AgentType, MessageHistoryUpdated
+from shotgun.agents.agent_manager import (
+    AgentManager,
+    AgentType,
+    MessageHistoryUpdated,
+    PartialResponseMessage,
+)
 from shotgun.agents.config import get_provider_model
 from shotgun.agents.models import (
     AgentDeps,
@@ -29,16 +33,22 @@ from shotgun.agents.models import (
     UserAnswer,
     UserQuestion,
 )
+from shotgun.sdk.codebase import CodebaseSDK
+from shotgun.sdk.exceptions import CodebaseNotFoundError, InvalidPathError
 from shotgun.sdk.services import get_artifact_service, get_codebase_service
+from shotgun.tui.commands import CommandHandler
+from shotgun.tui.screens.chat_screen.history import ChatHistory
 
 from ..components.prompt_input import PromptInput
 from ..components.spinner import Spinner
-from ..components.vertical_tail import VerticalTail
-
+from .chat_screen.command_providers import (
+    AgentModeProvider,
+    CodebaseCommandProvider,
+    DeleteCodebasePaletteProvider,
+    ProviderSetupProvider,
+)
 
-def _dummy_system_prompt_fn(ctx: RunContext[AgentDeps]) -> str:
-    """Dummy system prompt function for TUI chat interface."""
-    return "You are a helpful AI assistant."
+logger = logging.getLogger(__name__)
 
 
 class PromptHistory:
@@ -67,91 +77,19 @@ class PromptHistory:
         self.curr = None
 
 
-class ChatHistory(Widget):
-    DEFAULT_CSS = """
-        VerticalTail {
-            align: left bottom;
-
-        }
-        VerticalTail > * {
-            height: auto;
-        }
-
-        Horizontal {
-            height: auto;
-            background: $secondary-muted;
-        }
-
-        Markdown {
-            height: auto;
-        }
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.items: list[ModelMessage] = []
-        self.vertical_tail: VerticalTail | None = None
-
-    def compose(self) -> ComposeResult:
-        self.vertical_tail = VerticalTail()
-        yield self.vertical_tail
-
-    def update_messages(self, messages: list[ModelMessage]) -> None:
-        """Update the displayed messages without recomposing."""
-        if not self.vertical_tail:
-            return
-
-        # Clear existing widgets
-        self.vertical_tail.remove_children()
-
-        # Add new message widgets
-        for item in messages:
-            if isinstance(item, ModelRequest):
-                self.vertical_tail.mount(UserQuestionWidget(item))
-            elif isinstance(item, ModelResponse):
-                self.vertical_tail.mount(AgentResponseWidget(item))
+@dataclass
+class CodebaseIndexSelection:
+    """User-selected repository path and name for indexing."""
 
-        self.items = messages
-
-
-class UserQuestionWidget(Widget):
-    def __init__(self, item: ModelRequest) -> None:
-        super().__init__()
-        self.item = item
-
-    def compose(self) -> ComposeResult:
-        prompt = "".join(str(part.content) for part in self.item.parts if part.content)
-        yield Markdown(markdown=f"**>** {prompt}")
-
-
-class AgentResponseWidget(Widget):
-    def __init__(self, item: ModelResponse) -> None:
-        super().__init__()
-        self.item = item
-
-    def compose(self) -> ComposeResult:
-        yield Markdown(markdown=f"**⏺** {self.compute_output()}")
-
-    def compute_output(self) -> str:
-        acc = ""
-        for part in self.item.parts:  # TextPart | ToolCallPart | BuiltinToolCallPart | BuiltinToolReturnPart | ThinkingPart
-            if isinstance(part, TextPart):
-                acc += part.content
-            elif isinstance(part, ToolCallPart):
-                acc += f"{part.tool_name}({part.args})\n"
-            elif isinstance(part, BuiltinToolCallPart):
-                acc += f"{part.tool_name}({part.args})\n"
-            elif isinstance(part, BuiltinToolReturnPart):
-                acc += f"{part.tool_name}()\n"
-            elif isinstance(part, ThinkingPart):
-                acc += f"Thinking: {part.content}\n"
-        return acc
+    repo_path: Path
+    name: str
 
 
 class StatusBar(Widget):
     DEFAULT_CSS = """
         StatusBar {
             text-wrap: wrap;
+            padding-left: 1;
         }
     """
 
@@ -165,6 +103,7 @@ class ModeIndicator(Widget):
     DEFAULT_CSS = """
         ModeIndicator {
             text-wrap: wrap;
+            padding-left: 1;
         }
     """
 
@@ -183,11 +122,13 @@ class ModeIndicator(Widget):
             AgentType.RESEARCH: "Research",
             AgentType.PLAN: "Planning",
             AgentType.TASKS: "Tasks",
+            AgentType.SPECIFY: "Specify",
         }
         mode_description = {
             AgentType.RESEARCH: "Research topics with web search and synthesize findings",
             AgentType.PLAN: "Create comprehensive, actionable plans with milestones",
             AgentType.TASKS: "Generate specific, actionable tasks from research and plans",
+            AgentType.SPECIFY: "Create detailed specifications and requirements documents",
         }
 
         mode_title = mode_display.get(self.mode, self.mode.value.title())
@@ -196,97 +137,154 @@ class ModeIndicator(Widget):
         return f"[bold $text-accent]{mode_title} mode[/][$foreground-muted] ({description})[/]"
 
 
-class AgentModeProvider(Provider):
-    """Command provider for agent mode switching."""
+class FilteredDirectoryTree(DirectoryTree):
+    def filter_paths(self, paths: Iterable[Path]) -> Iterable[Path]:
+        return [path for path in paths if path.is_dir()]
 
-    @property
-    def chat_screen(self) -> "ChatScreen":
-        return cast(ChatScreen, self.screen)
 
-    def set_mode(self, mode: AgentType) -> None:
-        """Switch to research mode."""
-        self.chat_screen.mode = mode
+class CodebaseIndexPromptScreen(ModalScreen[bool]):
+    """Modal dialog asking whether to index the detected codebase."""
 
-    async def discover(self) -> AsyncGenerator[DiscoveryHit, None]:
-        """Provide default mode switching commands when palette opens."""
-        yield DiscoveryHit(
-            "Switch to Research Mode",
-            lambda: self.set_mode(AgentType.RESEARCH),
-            help="🔬 Research topics with web search and synthesize findings",
-        )
-        yield DiscoveryHit(
-            "Switch to Plan Mode",
-            lambda: self.set_mode(AgentType.PLAN),
-            help="📋 Create comprehensive, actionable plans with milestones",
-        )
-        yield DiscoveryHit(
-            "Switch to Tasks Mode",
-            lambda: self.set_mode(AgentType.TASKS),
-            help="✅ Generate specific, actionable tasks from research and plans",
-        )
+    DEFAULT_CSS = """
+        CodebaseIndexPromptScreen {
+            align: center middle;
+            background: rgba(0, 0, 0, 0.0);
+        }
 
-    async def search(self, query: str) -> AsyncGenerator[Hit, None]:
-        """Search for mode commands."""
-        matcher = self.matcher(query)
+        CodebaseIndexPromptScreen > #index-prompt-dialog {
+            width: 60%;
+            max-width: 60;
+            height: auto;
+            border: wide $primary;
+            padding: 1 2;
+            layout: vertical;
+            background: $surface;
+            height: auto;
+        }
 
-        commands = [
-            (
-                "Switch to Research Mode",
-                "🔬 Research topics with web search and synthesize findings",
-                lambda: self.set_mode(AgentType.RESEARCH),
-                AgentType.RESEARCH,
-            ),
-            (
-                "Switch to Plan Mode",
-                "📋 Create comprehensive, actionable plans with milestones",
-                lambda: self.set_mode(AgentType.PLAN),
-                AgentType.PLAN,
-            ),
-            (
-                "Switch to Tasks Mode",
-                "✅ Generate specific, actionable tasks from research and plans",
-                lambda: self.set_mode(AgentType.TASKS),
-                AgentType.TASKS,
-            ),
-        ]
+        #index-prompt-buttons {
+            layout: horizontal;
+            align-horizontal: right;
+            height: auto;
+        }
+    """
 
-        for title, help_text, callback, mode in commands:
-            if self.chat_screen.mode == mode:
-                continue
-            score = matcher.match(title)
-            if score > 0:
-                yield Hit(score, matcher.highlight(title), callback, help=help_text)
+    def compose(self) -> ComposeResult:
+        with Container(id="index-prompt-dialog"):
+            yield Label("Index your codebase?", id="index-prompt-title")
+            yield Static(
+                "We found project files but no index yet. Indexing enables smarter chat."
+            )
+            with Container(id="index-prompt-buttons"):
+                yield Button(
+                    "Index now",
+                    id="index-prompt-confirm",
+                    variant="primary",
+                )
+                yield Button("Not now", id="index-prompt-cancel")
 
+    @on(Button.Pressed, "#index-prompt-cancel")
+    def handle_cancel(self, event: Button.Pressed) -> None:
+        event.stop()
+        self.dismiss(False)
 
-class ProviderSetupProvider(Provider):
-    """Command palette entries for provider configuration."""
+    @on(Button.Pressed, "#index-prompt-confirm")
+    def handle_confirm(self, event: Button.Pressed) -> None:
+        event.stop()
+        self.dismiss(True)
 
-    @property
-    def chat_screen(self) -> "ChatScreen":
-        return cast(ChatScreen, self.screen)
 
-    def open_provider_config(self) -> None:
-        """Show the provider configuration screen."""
-        self.chat_screen.app.push_screen("provider_config")
+class CodebaseIndexScreen(ModalScreen[CodebaseIndexSelection | None]):
+    """Modal dialog for choosing a repository and name to index."""
 
-    async def discover(self) -> AsyncGenerator[DiscoveryHit, None]:
-        yield DiscoveryHit(
-            "Open Provider Setup",
-            self.open_provider_config,
-            help="⚙️ Manage API keys for available providers",
-        )
+    DEFAULT_CSS = """
+        CodebaseIndexScreen {
+            align: center middle;
+            background: rgba(0, 0, 0, 0.0);
+        }
+        CodebaseIndexScreen > #index-dialog {
+            width: 80%;
+            max-width: 80;
+            height: 80%;
+            max-height: 40;
+            border: wide $primary;
+            padding: 1;
+            layout: vertical;
+            background: $surface;
+        }
+
+        #index-dialog DirectoryTree {
+            height: 1fr;
+            border: solid $accent;
+            overflow: auto;
+        }
+
+        #index-dialog-controls {
+            layout: horizontal;
+            align-horizontal: right;
+            padding-top: 1;
+        }
+    """
 
-    async def search(self, query: str) -> AsyncGenerator[Hit, None]:
-        matcher = self.matcher(query)
-        title = "Open Provider Setup"
-        score = matcher.match(title)
-        if score > 0:
-            yield Hit(
-                score,
-                matcher.highlight(title),
-                self.open_provider_config,
-                help="⚙️ Manage API keys for available providers",
+    def __init__(self, start_path: Path | None = None) -> None:
+        super().__init__()
+        self.start_path = Path(start_path or Path.cwd())
+        self.selected_path: Path | None = self.start_path
+
+    def compose(self) -> ComposeResult:
+        with Container(id="index-dialog"):
+            yield Label("Index a codebase", id="index-dialog-title")
+            yield FilteredDirectoryTree(self.start_path, id="index-directory-tree")
+            yield Input(
+                placeholder="Enter a name for the codebase",
+                id="index-codebase-name",
             )
+            with Container(id="index-dialog-controls"):
+                yield Button("Cancel", id="index-cancel")
+                yield Button(
+                    "Index",
+                    id="index-confirm",
+                    variant="primary",
+                    disabled=True,
+                )
+
+    def _update_confirm(self) -> None:
+        confirm = self.query_one("#index-confirm", Button)
+        name_input = self.query_one("#index-codebase-name", Input)
+        confirm.disabled = not (self.selected_path and name_input.value.strip())
+
+    @on(DirectoryTree.DirectorySelected, "#index-directory-tree")
+    def handle_directory_selected(self, event: DirectoryTree.DirectorySelected) -> None:
+        event.stop()
+        selected = event.path if event.path.is_dir() else event.path.parent
+        self.selected_path = selected
+        name_input = self.query_one("#index-codebase-name", Input)
+        if not name_input.value:
+            name_input.value = selected.name
+        self._update_confirm()
+
+    @on(Input.Changed, "#index-codebase-name")
+    def handle_name_changed(self, event: Input.Changed) -> None:
+        event.stop()
+        self._update_confirm()
+
+    @on(Button.Pressed, "#index-cancel")
+    def handle_cancel(self, event: Button.Pressed) -> None:
+        event.stop()
+        self.dismiss(None)
+
+    @on(Button.Pressed, "#index-confirm")
+    def handle_confirm(self, event: Button.Pressed) -> None:
+        event.stop()
+        name_input = self.query_one("#index-codebase-name", Input)
+        if not self.selected_path:
+            self.dismiss(None)
+            return
+        selection = CodebaseIndexSelection(
+            repo_path=self.selected_path,
+            name=name_input.value.strip(),
+        )
+        self.dismiss(selection)
 
 
 class ChatScreen(Screen[None]):
@@ -297,7 +295,7 @@ class ChatScreen(Screen[None]):
         ("shift+tab", "toggle_mode", "Toggle mode"),
     ]
 
-    COMMANDS = {AgentModeProvider, ProviderSetupProvider}
+    COMMANDS = {AgentModeProvider, ProviderSetupProvider, CodebaseCommandProvider}
 
     _PLACEHOLDER_BY_MODE: dict[AgentType, str] = {
         AgentType.RESEARCH: (
@@ -309,6 +307,9 @@ class ChatScreen(Screen[None]):
         AgentType.TASKS: (
             "Request actionable work, e.g. break down tasks to wire OpenTelemetry into the API"
         ),
+        AgentType.SPECIFY: (
+            "Request detailed specifications, e.g. create a comprehensive spec for user authentication system"
+        ),
     }
 
     value = reactive("")
@@ -317,6 +318,8 @@ class ChatScreen(Screen[None]):
     messages = reactive(list[ModelMessage]())
     working = reactive(False)
     question: reactive[UserQuestion | None] = reactive(None)
+    indexing_job: reactive[CodebaseIndexSelection | None] = reactive(None)
+    partial_message: reactive[ModelMessage | None] = reactive(None)
 
     def __init__(self) -> None:
         super().__init__()
@@ -324,19 +327,62 @@ class ChatScreen(Screen[None]):
         model_config = get_provider_model()
         codebase_service = get_codebase_service()
         artifact_service = get_artifact_service()
+        self.codebase_sdk = CodebaseSDK()
+
+        # Create shared deps without system_prompt_fn (agents provide their own)
+        # We need a placeholder system_prompt_fn to satisfy the field requirement
+        def _placeholder_system_prompt_fn(ctx: RunContext[AgentDeps]) -> str:
+            raise RuntimeError(
+                "This should not be called - agents provide their own system_prompt_fn"
+            )
+
         self.deps = AgentDeps(
             interactive_mode=True,
             llm_model=model_config,
             codebase_service=codebase_service,
             artifact_service=artifact_service,
-            system_prompt_fn=_dummy_system_prompt_fn,
+            system_prompt_fn=_placeholder_system_prompt_fn,
         )
         self.agent_manager = AgentManager(deps=self.deps, initial_type=self.mode)
+        self.command_handler = CommandHandler()
 
     def on_mount(self) -> None:
         self.query_one(PromptInput).focus(scroll_visible=True)
         # Hide spinner initially
         self.query_one("#spinner").display = False
+        self.call_later(self.check_if_codebase_is_indexed)
+        # Start the question listener worker to handle ask_user interactions
+        self.call_later(self.add_question_listener)
+
+    @work
+    async def check_if_codebase_is_indexed(self) -> None:
+        cur_dir = Path.cwd().resolve()
+        is_empty = all(
+            dir.is_dir() and dir.name in ["__pycache__", ".git", ".shotgun"]
+            for dir in cur_dir.iterdir()
+        )
+        if is_empty:
+            return
+
+        # find at least one codebase that is indexed in the current directory
+        directory_indexed = next(
+            (
+                dir
+                for dir in (await self.codebase_sdk.list_codebases()).graphs
+                if cur_dir.is_relative_to(Path(dir.repo_path).resolve())
+            ),
+            None,
+        )
+        if directory_indexed:
+            self.mount_hint(help_text_with_codebase())
+            return
+
+        should_index = await self.app.push_screen_wait(CodebaseIndexPromptScreen())
+        if not should_index:
+            self.mount_hint(help_text_empty_dir())
+            return
+
+        self.index_codebase_command()
 
     def watch_mode(self, new_mode: AgentType) -> None:
         """React to mode changes by updating the agent manager."""
@@ -377,7 +423,7 @@ class ChatScreen(Screen[None]):
                 question_display.display = False
 
     def action_toggle_mode(self) -> None:
-        modes = [AgentType.RESEARCH, AgentType.PLAN, AgentType.TASKS]
+        modes = [AgentType.RESEARCH, AgentType.PLAN, AgentType.TASKS, AgentType.SPECIFY]
         self.mode = modes[(modes.index(self.mode) + 1) % len(modes)]
         self.agent_manager.set_agent(self.mode)
         # whoops it actually changes focus. Let's be brutal for now
@@ -394,9 +440,9 @@ class ChatScreen(Screen[None]):
     def compose(self) -> ComposeResult:
         """Create child widgets for the app."""
         with Container(id="window"):
+            yield self.agent_manager
             yield ChatHistory()
             yield Markdown(markdown="", id="question-display")
-            yield self.agent_manager
             with Container(id="footer"):
                 yield Spinner(
                     text="Processing...",
@@ -410,7 +456,24 @@ class ChatScreen(Screen[None]):
                     id="prompt-input",
                     placeholder=self._placeholder_for_mode(self.mode),
                 )
-                yield ModeIndicator(mode=self.mode)
+                with Grid():
+                    yield ModeIndicator(mode=self.mode)
+                    yield Static("", id="indexing-job-display")
+
+    def mount_hint(self, markdown: str) -> None:
+        chat_history = self.query_one(ChatHistory)
+        if not chat_history.vertical_tail:
+            return
+        chat_history.vertical_tail.mount(Markdown(markdown))
+
+    @on(PartialResponseMessage)
+    def handle_partial_response(self, event: PartialResponseMessage) -> None:
+        self.partial_message = event.message
+
+        partial_response_widget = self.query_one(ChatHistory)
+        partial_response_widget.partial_response = self.partial_message
+        if event.is_last:
+            partial_response_widget.partial_response = None
 
     @on(MessageHistoryUpdated)
     def handle_message_history_updated(self, event: MessageHistoryUpdated) -> None:
@@ -451,11 +514,42 @@ class ChatScreen(Screen[None]):
 
     @on(PromptInput.Submitted)
     async def handle_submit(self, message: PromptInput.Submitted) -> None:
+        text = message.text.strip()
+
+        # If empty text, just clear input and return
+        if not text:
+            prompt_input = self.query_one(PromptInput)
+            prompt_input.clear()
+            self.value = ""
+            return
+
+        # Check if it's a command
+        if self.command_handler.is_command(text):
+            success, response = self.command_handler.handle_command(text)
+
+            # Add the command to history
+            self.history.append(message.text)
+
+            # Display the command in chat history
+            user_message = ModelRequest(parts=[UserPromptPart(content=text)])
+            self.messages = self.messages + [user_message]
+
+            # Display the response (help text or error message)
+            response_message = ModelResponse(parts=[TextPart(content=response)])
+            self.messages = self.messages + [response_message]
+
+            # Clear the input
+            prompt_input = self.query_one(PromptInput)
+            prompt_input.clear()
+            self.value = ""
+            return
+
+        # Not a command, process as normal
         self.history.append(message.text)
 
         # Clear the input
         self.value = ""
-        self.run_agent(message.text)
+        self.run_agent(text)  # Use stripped text
 
         prompt_input = self.query_one(PromptInput)
         prompt_input.clear()
@@ -464,6 +558,69 @@ class ChatScreen(Screen[None]):
         """Return the placeholder text appropriate for the current mode."""
         return self._PLACEHOLDER_BY_MODE.get(mode, "Type your message")
 
+    def index_codebase_command(self) -> None:
+        start_path = Path.cwd()
+
+        def handle_result(result: CodebaseIndexSelection | None) -> None:
+            if result:
+                self.call_later(lambda: self.index_codebase(result))
+
+        self.app.push_screen(
+            CodebaseIndexScreen(start_path=start_path),
+            handle_result,
+        )
+
+    def delete_codebase_command(self) -> None:
+        self.app.push_screen(
+            CommandPalette(
+                providers=[DeleteCodebasePaletteProvider],
+                placeholder="Select a codebase to delete…",
+            )
+        )
+
+    def delete_codebase_from_palette(self, graph_id: str) -> None:
+        stack = getattr(self.app, "screen_stack", None)
+        if stack and isinstance(stack[-1], CommandPalette):
+            self.app.pop_screen()
+
+        self.call_later(lambda: self.delete_codebase(graph_id))
+
+    @work
+    async def delete_codebase(self, graph_id: str) -> None:
+        try:
+            await self.codebase_sdk.delete_codebase(graph_id)
+            self.notify(f"Deleted codebase: {graph_id}", severity="information")
+        except CodebaseNotFoundError as exc:
+            self.notify(str(exc), severity="error")
+        except Exception as exc:  # pragma: no cover - defensive UI path
+            self.notify(f"Failed to delete codebase: {exc}", severity="error")
+
+    @work
+    async def index_codebase(self, selection: CodebaseIndexSelection) -> None:
+        label = self.query_one("#indexing-job-display", Static)
+        label.update(
+            f"[$foreground-muted]Indexing [bold $text-accent]{selection.name}[/]...[/]"
+        )
+        label.refresh()
+        try:
+            result = await self.codebase_sdk.index_codebase(
+                selection.repo_path, selection.name
+            )
+            self.notify(
+                f"Indexed codebase '{result.name}' (ID: {result.graph_id})",
+                severity="information",
+                timeout=8,
+            )
+        except InvalidPathError as exc:
+            self.notify(str(exc), severity="error")
+
+        except Exception as exc:  # pragma: no cover - defensive UI path
+            self.notify(f"Failed to index codebase: {exc}", severity="error")
+        finally:
+            label.update("")
+            label.refresh()
+            self.mount_hint(codebase_indexed_hint(selection.name))
+
     @work
     async def run_agent(self, message: str) -> None:
         deferred_tool_results = None
@@ -495,3 +652,29 @@ class ChatScreen(Screen[None]):
 
         prompt_input = self.query_one(PromptInput)
         prompt_input.focus()
+
+
+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 (
+        "I can help with:\n\n"
+        "- Speccing out a new feature\n"
+        "- Onboarding you onto this project\n"
+        "- Helping with a refactor spec\n"
+        "- Creating AGENTS.md file for this project\n"
+    )
+
+
+def help_text_empty_dir() -> str:
+    return (
+        "What would you like to build? Here are some examples:\n\n"
+        "- Research FastAPI vs Django\n"
+        "- Plan my new web app using React\n"
+        "- Create PRD for my planned product\n"
+    )