open-data-sci 0.1.0__py3-none-any.whl

opendatasci/_tui/app.py

@@ -0,0 +1,429 @@
+import argparse
+import importlib.metadata
+import logging
+import uuid
+from pathlib import Path
+
+from dotenv import load_dotenv
+from rich.console import Console
+from rich.table import Table
+from textual import events, on, work
+from textual.app import App, ComposeResult
+from textual.binding import Binding
+from textual.containers import Horizontal
+from textual.timer import Timer
+from textual.widgets import Footer, Input
+
+from opendatasci.configs import DEFAULT_MODEL, DEFAULT_SECONDARY_MODEL, OpenDataSciConfig
+from opendatasci.models.providers import Provider
+
+from . import theme as _theme
+from .controller import CLIController, UIAdapter
+from .widgets import (
+    AppHeader,
+    ChatPane,
+    CompletionPopup,
+    MessageBubble,
+    SmartInput,
+    ThinkingBlock,
+    ToolCallBlock,
+    TurnStatusBar,
+)
+
+
+def _print_providers() -> None:
+    table = Table(title=None, show_header=True, header_style="bold")
+    table.add_column("Provider")
+    table.add_column("Default model")
+    for provider, model in DEFAULT_MODEL.items():
+        table.add_row(provider, model)
+    Console().print(table)
+
+
+def _get_version() -> str:
+    try:
+        return importlib.metadata.version("open-data-sci")
+    except importlib.metadata.PackageNotFoundError:
+        logging.getLogger(__name__).warning(
+            "open-data-sci package not found; falling back to hardcoded version '0.1.0'"
+        )
+        return "0.1.0"
+
+
+class OpenDataSciApp(App[None]):
+    """OpenDataSci — full TUI for AI-powered data science."""
+
+    CSS_PATH = "styles.tcss"
+
+    BINDINGS = [
+        Binding("ctrl+c", "request_quit", "Quit"),
+        Binding("ctrl+d", "quit", "Quit", show=False),
+        Binding("ctrl+r", "reset", "Reset"),
+        Binding("ctrl+l", "clear_conv", "Clear", show=False),
+        Binding("escape", "focus_input", "Focus", show=False),
+    ]
+
+    def __init__(
+        self,
+        workspace_path: str,
+        session_id: str,
+        datasci_config: OpenDataSciConfig,
+        theme: str = "default",
+    ) -> None:
+        palette = _theme.THEMES.get(theme, _theme.DARK)
+        _theme.active.update(palette)
+        _theme.active_name = theme if theme in _theme.THEMES else "default"
+        if theme == "accessible":
+            self.CSS_PATH = str(Path(__file__).parent / "styles_visible.tcss")  # type: ignore[misc]
+        super().__init__()
+        self._controller = CLIController(
+            ui=self,  # type: ignore[arg-type]
+            workspace_path=workspace_path,
+            datasci_config=datasci_config,
+            session_id=session_id,
+        )
+
+    def compose(self) -> ComposeResult:
+        yield AppHeader(
+            version=_get_version(),
+            provider=self._controller.provider,
+            model=self._controller.model,
+            workspace=str(Path(self._controller._workspace_path).resolve()),
+        )
+        with Horizontal(id="main"):
+            yield ChatPane()
+        yield Footer()
+
+    def on_mount(self) -> None:
+        self._quit_requested = False
+        self._quit_timer: Timer | None = None
+        self.query_one("#user-input", Input).focus()
+        self._boot()
+
+    async def on_unmount(self) -> None:
+        await self._controller.close()
+
+    # ── UIAdapter implementation ──────────────────────────────────────────────
+
+    def add_message(self, role: str, content: str = "") -> MessageBubble:
+        return self.query_one(ChatPane).add_message(role, content)
+
+    def add_divider(self) -> None:
+        self.query_one(ChatPane).add_divider()
+
+    def add_turn_status_bar(self) -> TurnStatusBar:
+        return self.query_one(ChatPane).add_turn_status_bar()
+
+    def add_ephemeral_block(self, communication: str, label: str, summary: str) -> ToolCallBlock:
+        return self.query_one(ChatPane).add_ephemeral_block(communication, label, summary)
+
+    def add_worker_block(self, communication: str, worker_summaries: list[str]) -> ToolCallBlock:
+        return self.query_one(ChatPane).add_worker_block(communication, worker_summaries)
+
+    def add_thinking_block(self) -> ThinkingBlock:
+        return self.query_one(ChatPane).add_thinking_block()
+
+    def clear_messages(self) -> None:
+        self.query_one(ChatPane).clear_messages()
+
+    def set_workspace(self, name: str) -> None:
+        self.query_one(AppHeader).set_workspace(name)
+
+    def set_file_count(self, description: str) -> None:
+        self.query_one(AppHeader).set_file_count(description)
+
+    def show_workspace_panel(self, files: list[str]) -> None:
+        self.query_one(ChatPane).show_workspace_panel(files)
+
+    def show_attachment(self, label: str) -> None:
+        self.query_one(ChatPane).show_attachment(label)
+
+    def hide_attachment(self) -> None:
+        self.query_one(ChatPane).hide_attachment()
+
+    def stop_agent(self) -> None:
+        self.workers.cancel_group(self, "agent")
+
+    def set_input_placeholder(self, text: str) -> None:
+        self.query_one("#user-input", Input).placeholder = text
+
+    def add_input_class(self, cls: str) -> None:
+        self.query_one("#user-input", Input).add_class(cls)
+
+    def remove_input_class(self, cls: str) -> None:
+        self.query_one("#user-input", Input).remove_class(cls)
+
+    def set_input_value(self, value: str, cursor: int | None = None) -> None:
+        inp = self.query_one("#user-input", Input)
+        inp.value = value
+        if cursor is not None:
+            inp.cursor_position = cursor
+
+    def show_completion(self, matches: list[str], selected: int) -> None:
+        self.query_one(CompletionPopup).show_matches(matches, selected)
+
+    def hide_completion(self) -> None:
+        self.query_one(CompletionPopup).hide()
+
+    # ── Event handlers ────────────────────────────────────────────────────────
+
+    @on(SmartInput.Pasted)
+    def on_paste_attachment(self, event: SmartInput.Pasted) -> None:
+        self._controller.on_paste(event._text)
+
+    @on(Input.Changed, "#user-input")
+    def on_input_changed(self, event: Input.Changed) -> None:
+        self._controller.on_input_changed(event.value)
+
+    @on(Input.Submitted, "#user-input")
+    async def on_submit(self, event: Input.Submitted) -> None:
+        raw = event.value.strip()
+        if raw:
+            self.query_one("#user-input", SmartInput).push_history(raw)
+        self.query_one("#user-input", Input).value = ""
+        action, query = await self._controller.on_submit(raw)
+        if action == "run":
+            self._run_agent(query)
+        elif action == "quit":
+            self.exit()
+
+    @on(events.Key)
+    def on_input_key(self, event: events.Key) -> None:
+        if event.key not in {"up", "down"}:
+            return
+        inp = self.query_one("#user-input", Input)
+        if self.focused is not inp:
+            return
+        direction = 1 if event.key == "down" else -1
+        if self._controller.has_completion_matches:
+            if self._controller.cycle_completion(inp.value, direction=direction):
+                event.stop()
+                event.prevent_default()
+        else:
+            self._controller._completing = True  # suppress Input.Changed fired by value update
+            if self.query_one("#user-input", SmartInput).navigate_history(direction):
+                event.stop()
+                event.prevent_default()
+            else:
+                self._controller._completing = False
+
+    # ── @work wrappers ────────────────────────────────────────────────────────
+
+    @work
+    async def _boot(self) -> None:
+        await self._controller.boot()
+
+    @work(exclusive=True, group="agent", exit_on_error=False)
+    async def _run_agent(self, query: str) -> None:
+        await self._controller.run_agent(query)
+
+    @work
+    async def _compact(self) -> None:
+        await self._controller.compact()
+
+    # ── Action handlers ───────────────────────────────────────────────────────
+
+    async def action_quit(self) -> None:
+        self.exit()
+
+    def action_request_quit(self) -> None:
+        if self._quit_requested:
+            self.exit()
+            return
+        self._quit_requested = True
+        self.notify("Press Ctrl+C again to quit", severity="warning", timeout=3)
+        if self._quit_timer is not None:
+            self._quit_timer.stop()
+        self._quit_timer = self.set_timer(3, self._reset_quit_request)
+
+    def _reset_quit_request(self) -> None:
+        self._quit_requested = False
+        self._quit_timer = None
+
+    async def action_reset(self) -> None:
+        await self._controller.reset()
+
+    async def action_clear_conv(self) -> None:
+        await self._controller.clear_conv()
+
+    def action_compact(self) -> None:
+        self._compact()
+
+    def action_focus_input(self) -> None:
+        self._controller.hide_completion()
+        self._controller.clear_paste_attachment()
+        if self._controller.awaiting_choice:
+            resume_input = self._controller.cancel_choice()
+            if resume_input is not None:
+                self._run_agent(resume_input)
+        self.query_one("#user-input", Input).focus()
+
+    @on(SmartInput.TabComplete)
+    def on_smart_input_tab_complete(self, event: SmartInput.TabComplete) -> None:
+        inp = self.query_one("#user-input", Input)
+        if not self._controller.cycle_completion(inp.value, direction=event._direction):
+            self.action_focus_next()
+
+
+# Maps a provider name to the OpenDataSciConfig field that holds its API key.
+# Providers that use cloud-native auth (bedrock, vertexai, ollama) have no key field.
+_PROVIDER_KEY_FIELD: dict[Provider, str | None] = {
+    Provider.ANTHROPIC: "anthropic_api_key",
+    Provider.OPENAI: "openai_api_key",
+    Provider.GEMINI: "google_api_key",
+    Provider.AZURE: "azure_api_key",
+    Provider.OPENAI_COMPATIBLE_SERVER: "openai_api_key",
+    Provider.BEDROCK: None,
+    Provider.VERTEXAI: None,
+    Provider.OLLAMA: None,
+}
+
+
+def main() -> None:
+    load_dotenv()
+
+    parser = argparse.ArgumentParser(
+        description="OpenDataSci — AI-powered data analytics",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  opendatasci data.xlsx
+  opendatasci data.csv --provider bedrock
+  opendatasci ./data_folder --provider openai --model gpt-4o
+  opendatasci data.csv --secondary-provider openai --secondary-model gpt-4o-mini
+  opendatasci data.csv --config path/to/datasci_config.yaml
+        """,
+    )
+    parser.add_argument(
+        "workspace_or_file",
+        nargs="?",
+        default=None,
+        help="Data file or directory containing data files to work with",
+    )
+    parser.add_argument(
+        "--provider",
+        default=None,
+        choices=list(Provider),
+        help="LLM provider for the primary model (default: anthropic)",
+    )
+    parser.add_argument(
+        "--model",
+        dest="model",
+        default=None,
+        help="Primary model name (provider-specific)",
+    )
+    parser.add_argument(
+        "--secondary-provider",
+        dest="secondary_provider",
+        default=None,
+        choices=list(Provider),
+        help="LLM provider for the secondary (auxiliary) model — may differ from --provider",
+    )
+    parser.add_argument(
+        "--secondary-model",
+        dest="secondary_model",
+        default=None,
+        help="Secondary model name (resolved against --secondary-provider or --provider)",
+    )
+    parser.add_argument(
+        "--api-key",
+        dest="api_key",
+        default=None,
+        help="API key for the primary provider (or set via environment variable)",
+    )
+    parser.add_argument(
+        "--theme",
+        choices=list(_theme.THEMES.keys()),
+        default="default",
+        help=(
+            "Colour palette. Choices: "
+            + ", ".join(_theme.THEMES.keys())
+            + ". Run `/themes` inside the TUI for descriptions."
+        ),
+    )
+    parser.add_argument(
+        "--config",
+        default=None,
+        metavar="FILE",
+        help=(
+            "Path to a YAML file containing OpenDataSciConfig fields. "
+            "Explicit TUI flags take precedence over values in the file."
+        ),
+    )
+    parser.add_argument(
+        "--list-providers",
+        action="store_true",
+        help="List supported providers and their default models, then exit",
+    )
+    parser.add_argument("--version", action="version", version=f"OpenDataSci {_get_version()}")
+    args = parser.parse_args()
+
+    if args.list_providers:
+        _print_providers()
+        return
+
+    if args.workspace_or_file is None:
+        parser.error("the following arguments are required: path")
+
+    # Build OpenDataSciConfig: YAML file provides the base; explicit TUI flags override.
+    if args.config:
+        datasci_config = OpenDataSciConfig.from_yaml(args.config)
+        overrides: dict[str, object] = {}
+        if args.provider is not None:
+            overrides["provider"] = args.provider
+        if args.model is not None:
+            overrides["model"] = args.model
+        if args.secondary_provider is not None:
+            overrides["secondary_provider"] = args.secondary_provider
+        if args.secondary_model is not None:
+            overrides["secondary_model"] = args.secondary_model
+        if args.api_key is not None:
+            effective_provider = str(args.provider or datasci_config.provider)
+            key_field = _PROVIDER_KEY_FIELD.get(Provider(effective_provider))
+            if key_field:
+                overrides[key_field] = args.api_key
+            else:
+                parser.error(
+                    f"--api-key is not supported for provider '{effective_provider}' "
+                    f"(uses cloud-native authentication)"
+                )
+        if overrides:
+            datasci_config = datasci_config.model_copy(update=overrides)
+    else:
+        provider: Provider = args.provider or Provider.ANTHROPIC
+        resolved_secondary_provider: Provider = args.secondary_provider or provider
+        kwargs: dict[str, object] = {
+            "provider": provider,
+            "model": args.model or DEFAULT_MODEL[provider],
+            "secondary_provider": resolved_secondary_provider,
+            "secondary_model": args.secondary_model
+            or DEFAULT_SECONDARY_MODEL[resolved_secondary_provider],
+        }
+        if args.api_key is not None:
+            key_field = _PROVIDER_KEY_FIELD.get(provider)
+            if key_field:
+                kwargs[key_field] = args.api_key
+            else:
+                parser.error(
+                    f"--api-key is not supported for provider '{provider}' "
+                    f"(uses cloud-native authentication)"
+                )
+        datasci_config = OpenDataSciConfig(**kwargs)  # type: ignore[arg-type]
+
+    session_id = uuid.uuid4().hex
+
+    OpenDataSciApp(
+        workspace_path=args.workspace_or_file,
+        session_id=session_id,
+        datasci_config=datasci_config,
+        theme=args.theme,
+    ).run()
+
+
+# Register OpenDataSciApp as a virtual subclass of UIAdapter to avoid the metaclass
+# conflict between Textual's _MessagePumpMeta and ABCMeta.
+UIAdapter.register(OpenDataSciApp)
+
+
+if __name__ == "__main__":
+    main()

opendatasci/_tui/commands.py

@@ -0,0 +1,95 @@
+"""Slash command registry and display-text formatters for the OpenDataSci TUI."""
+
+import re
+
+from opendatasci.models.providers import Provider
+
+SLASH_COMMANDS: list[str] = [
+    "/clear",
+    "/compact",
+    "/exit",
+    "/help",
+    "/ls-workspace",
+    "/models",
+    "/reset",
+    "/stop",
+    "/themes",
+]
+
+SLASH_COMMAND_DESCRIPTIONS: dict[str, str] = {
+    "/clear": "Clear conversation context",
+    "/compact": "Summarize conversation history",
+    "/exit": "Exit OpenDataSci",
+    "/help": "Show all commands",
+    "/ls-workspace": "List workspace files",
+    "/models": "Show models in use",
+    "/reset": "Reset agent session",
+    "/stop": "Stop the running agent",
+    "/themes": "List available colour themes",
+}
+
+_PROVIDER_DISPLAY: dict[Provider, str] = {
+    Provider.ANTHROPIC: "Anthropic",
+    Provider.OPENAI: "OpenAI",
+    Provider.BEDROCK: "AWS Bedrock",
+    Provider.GEMINI: "Google",
+    Provider.VERTEXAI: "Google Vertex AI",
+    Provider.AZURE: "Azure OpenAI",
+    Provider.OLLAMA: "Ollama",
+    Provider.OPENAI_COMPATIBLE_SERVER: "OpenAI-compatible server",
+}
+
+
+def _fmt_model(provider: str, model_id: str) -> str:
+    try:
+        provider_label = _PROVIDER_DISPLAY[Provider(provider)]
+    except (KeyError, ValueError):
+        provider_label = provider.title()
+    m = re.search(r"claude-([a-z]+)-(\d+)-(\d+)", model_id)
+    if m:
+        variant, major, minor = m.groups()
+        return f"{provider_label} Claude {variant.title()} {major}.{minor}"
+    return f"{provider_label} {model_id}"
+
+
+def format_models_message(
+    primary_provider: str, model: str, secondary_provider: str, secondary_model: str
+) -> str:
+    """Return the Markdown text shown by the /models command."""
+    lines = [
+        "## Models\n",
+        f"- **Primary Model**   : {_fmt_model(primary_provider, model)}",
+        f"- **Secondary Model** : {_fmt_model(secondary_provider, secondary_model)}",
+    ]
+    return "\n".join(lines)
+
+
+def format_help_message() -> str:
+    """Return the Markdown text shown by the /help command."""
+    lines = [
+        "## Available Commands\n",
+        "- **/clear** — Clear the conversation (preserves session variables)",
+        "- **/compact** — Summarize and compress the conversation history",
+        "- **/exit** — Exit OpenDataSci",
+        "- **/help** — Show this help message",
+        "- **/ls-workspace** — List files in the workspace",
+        "- **/models** — Show the primary and secondary model in use",
+        "- **/reset** — Reset the agent session and reload data from disk",
+        "- **/stop** — Stop the running agent (future messages pick up where it left off)",
+        "- **/themes** — List available colour themes (selected at launch with `--theme`)",
+    ]
+    lines.append("\n**Tip:** Type `/` to see commands via autocomplete, or `@` to attach a file.")
+    return "\n".join(lines)
+
+
+def format_themes_message(active_name: str, themes: dict[str, str]) -> str:
+    """Return the Markdown text shown by the /themes command.
+
+    `themes` is a mapping of theme name to a one-line description.
+    """
+    lines = ["## Colour Themes\n"]
+    for name, description in themes.items():
+        marker = " *(active)*" if name == active_name else ""
+        lines.append(f"- **{name}**{marker} — {description}")
+    lines.append("\nSwitch themes by relaunching with `--theme <name>`.")
+    return "\n".join(lines)

opendatasci/_tui/completion.py

@@ -0,0 +1,139 @@
+"""Tab-completion state machine for the OpenDataSci TUI input bar.
+
+Manages slash-command and @file-path completion independently of Textual
+widgets, so it can be unit-tested without a running app.
+"""
+
+import logging
+
+from .adapter import UIAdapter
+from .commands import SLASH_COMMAND_DESCRIPTIONS, SLASH_COMMANDS
+from .file_refs import _discover_files, _find_at_fragment, _find_slash_fragment
+
+logger = logging.getLogger(__name__)
+
+
+class CompletionState:
+    """Encapsulates all mutable state for the input-bar completion popup.
+
+    Owned by ``CLIController``; delegates UI updates through ``UIAdapter``
+    so the logic remains fully testable without a Textual widget tree.
+    """
+
+    def __init__(self, extra_commands: list[str] | None = None) -> None:
+        self._matches: list[str] = []
+        self._displays: list[str] = []
+        self._idx: int = -1
+        self._at_pos: int = -1
+        self._mode: str = "file"
+        # Set to True by cycle() before it changes the input value, so that
+        # the resulting on_input_changed callback knows to ignore the event.
+        self._completing: bool = False
+        # Cache the last @-fragment so _discover_files is not called on every
+        # keystroke when the fragment hasn't changed.
+        self._last_at_fragment: str | None = None
+        self._cached_at_matches: list[str] = []
+        self._all_commands: list[str] = SLASH_COMMANDS + list(extra_commands or [])
+
+    @property
+    def has_matches(self) -> bool:
+        """True when the completion popup currently has items to navigate."""
+        return bool(self._matches)
+
+    def on_input_changed(self, value: str, ui: UIAdapter) -> bool:
+        """Handle an input-text change.
+
+        Returns ``True`` when the change was a programmatic tab-completion
+        update (the caller should skip further processing).
+        """
+        if self._completing:
+            self._completing = False
+            return True
+
+        slash_frag = _find_slash_fragment(value)
+        if slash_frag is not None:
+            matches = [cmd for cmd in self._all_commands if cmd.startswith(slash_frag)]
+            if matches and not (len(matches) == 1 and slash_frag == matches[0]):
+                self._matches = matches
+                self._displays = [
+                    f"{cmd}  {SLASH_COMMAND_DESCRIPTIONS.get(cmd, '')}" for cmd in matches
+                ]
+                self._idx = -1
+                self._at_pos = -1
+                self._mode = "slash"
+                ui.show_completion(self._displays, self._idx)
+            else:
+                self.hide(ui)
+            return False
+
+        result = _find_at_fragment(value)
+        if result is None:
+            self.hide(ui)
+            return False
+
+        fragment, at_pos = result
+        if fragment != self._last_at_fragment:
+            self._cached_at_matches = _discover_files(fragment)
+            self._last_at_fragment = fragment
+        matches = self._cached_at_matches
+        if not matches:
+            self.hide(ui)
+            return False
+
+        self._matches = matches
+        self._idx = -1
+        self._at_pos = at_pos
+        self._mode = "file"
+        ui.show_completion(matches, self._idx)
+        return False
+
+    def cycle(self, current_value: str, direction: int, ui: UIAdapter) -> bool:
+        """Cycle the completion selection by ``direction`` (+1 down, -1 up).
+
+        Returns ``True`` if a completion was applied (caller should NOT call
+        focus_next); ``False`` if there are no completions to cycle.
+        """
+        if not self._matches:
+            return False
+
+        if direction < 0 and self._idx == -1:
+            self._idx = len(self._matches) - 1
+        else:
+            self._idx = (self._idx + direction) % len(self._matches)
+
+        match = self._matches[self._idx]
+
+        if self._mode == "slash":
+            self._completing = True
+            ui.set_input_value(match, len(match))
+            ui.show_completion(self._displays, self._idx)
+            return True
+
+        after = current_value[self._at_pos + 1 :]
+        space_pos = after.find(" ")
+        rest = after[space_pos:] if space_pos != -1 else ""
+        new_value = current_value[: self._at_pos + 1] + match + rest
+        cursor = self._at_pos + 1 + len(match)
+        self._completing = True
+        ui.set_input_value(new_value, cursor)
+        ui.show_completion(self._matches, self._idx)
+        return True
+
+    def hide(self, ui: UIAdapter) -> None:
+        """Clear completion state and hide the popup (no-op if already hidden)."""
+        was_showing = bool(self._matches)
+        self._matches = []
+        self._displays = []
+        self._idx = -1
+        self._at_pos = -1
+        self._mode = "file"
+        self._last_at_fragment = None
+        self._cached_at_matches = []
+        # Only update the UI when the popup was actually visible, to avoid
+        # spurious re-renders on every normal keystroke.
+        if not was_showing:
+            return
+        try:
+            ui.hide_completion()
+        except Exception:
+            logger.exception("Failed to hide completion popup")