Merge pull request #20 from nick-galluzzo/refactor/custom-embeddings-cli

refactor: Restructure embedding CLI

Author: AstraBert

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "notebookllama"
-version = "0.2.2.post1"
+version = "0.2.3"
 description = "An OSS and LlamaCloud-backed alternative to NotebookLM"
 readme = "README.md"
 requires-python = ">=3.13"

tools/cli/config/__init__.py

@@ -0,0 +1,3 @@
+from .models import EmbeddingConfig
+
+__all__ = ["EmbeddingConfig"]

tools/cli/config/models.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class EmbeddingConfig:
+    provider: str
+    api_key: Optional[str] = None
+    model: Optional[str] = None
+    region: Optional[str] = None
+    key_id: Optional[str] = None
+    embedding_config: Optional[object] = None

tools/cli/embedding_app.py

@@ -0,0 +1,40 @@
+import os
+from textual.app import App
+
+from .config import EmbeddingConfig
+from .screens import InitialScreen
+
+
+class EmbeddingSetupApp(App):
+    """Main application for embedding configuration setup."""
+
+    CSS_PATH = "stylesheets/base.tcss"
+
+    def __init__(self):
+        super().__init__()
+        self.config = EmbeddingConfig(provider="")
+
+    def on_mount(self) -> None:
+        self.push_screen(InitialScreen())
+
+    def handle_completion(self, config: EmbeddingConfig) -> None:
+        self.exit(config)
+
+    def handle_default_setup(self) -> None:
+        from llama_index.embeddings.openai import OpenAIEmbedding
+        from llama_cloud import PipelineCreateEmbeddingConfig_OpenaiEmbedding
+
+        self.config.provider = "OpenAI"
+        self.config.api_key = os.getenv("OPENAI_API_KEY")
+        self.config.model = "text-embedding-3-small"
+
+        embed_model = OpenAIEmbedding(
+            model=self.config.model, api_key=self.config.api_key
+        )
+        embedding_config = PipelineCreateEmbeddingConfig_OpenaiEmbedding(
+            type="OPENAI_EMBEDDING",
+            component=embed_model,
+        )
+        self.config = embedding_config
+
+        self.handle_completion(self.config)

tools/cli/screens/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseScreen, ConfigurationScreen
+from .initial import InitialScreen
+from .embedding_provider import ProviderSelectScreen
+
+__all__ = ["BaseScreen", "ConfigurationScreen", "InitialScreen", "ProviderSelectScreen"]

tools/cli/screens/base.py

@@ -0,0 +1,60 @@
+from textual.app import ComposeResult
+from textual.containers import Container
+from textual.screen import Screen
+from textual.widgets import Label, Footer
+from textual.binding import Binding
+from textual.widgets import Input
+
+
+class BaseScreen(Screen):
+    """Base screen with common functionality for all screens."""
+
+    BINDINGS = [
+        Binding("ctrl+q", "quit", "Exit", key_display="ctrl+q"),
+        Binding("ctrl+d", "toggle_dark", "Toggle Dark Theme", key_display="ctrl+d"),
+    ]
+
+    def action_toggle_dark(self) -> None:
+        self.app.theme = (
+            "textual-dark" if self.app.theme == "textual-light" else "textual-light"
+        )
+
+    def action_quit(self) -> None:
+        self.app.exit()
+
+    def compose(self) -> ComposeResult:
+        yield Container(
+            Label(self.get_title(), classes="form-title"),
+            *self.get_form_elements(),
+            classes="form-container",
+        )
+        yield Footer()
+
+    def get_title(self) -> str:
+        return "Base Screen"
+
+    def get_form_elements(self) -> list[ComposeResult]:
+        return []
+
+
+class ConfigurationScreen(BaseScreen):
+    """Base screen provider configuration with submit functionality."""
+
+    BINDINGS = BaseScreen.BINDINGS + [
+        Binding("shift+enter", "submit", "Submit"),
+    ]
+
+    def on_input_submitted(self, event: Input.Submitted) -> None:
+        """Catches the Enter key press and delegates the work."""
+        self.process_submission()
+
+    def process_submission(self) -> None:
+        """
+        To be implemented by each specific provider screen.
+        This method contains the unique logic for validating and creating
+        the embedding configuration.
+        """
+
+        raise NotImplementedError(
+            "Each configuration screen must implement 'process_submission'"
+        )

tools/cli/screens/embedding_provider.py

@@ -0,0 +1,62 @@
+from textual import on
+from textual.widgets import Select
+
+from .base import BaseScreen
+from .embedding_providers import (
+    OpenAIEmbeddingScreen,
+    BedrockEmbeddingScreen,
+    AzureEmbeddingScreen,
+    GeminiEmbeddingScreen,
+    CohereEmbeddingScreen,
+    HuggingFaceEmbeddingScreen,
+)
+
+
+class ProviderSelectScreen(BaseScreen):
+    """Screen for selecting embedding provider."""
+
+    def get_title(self) -> str:
+        return "Select an embedding provider"
+
+    def get_form_elements(self) -> list:
+        return [
+            Select(
+                options=[
+                    ("OpenAI", "OpenAI"),
+                    ("Cohere", "Cohere"),
+                    ("Bedrock", "Bedrock"),
+                    ("HuggingFace", "HuggingFace"),
+                    ("Azure", "Azure"),
+                    ("Gemini", "Gemini"),
+                ],
+                prompt="Please select an embedding provider",
+                classes="form-control",
+                id="provider_select",
+            )
+        ]
+
+    @on(Select.Changed, "#provider_select")
+    def handle_selection(self, event: Select.Changed) -> None:
+        from ..embedding_app import EmbeddingSetupApp
+
+        app = self.app
+        if isinstance(app, EmbeddingSetupApp):
+            app.config.provider = event.value
+            self.handle_next()
+
+    def handle_next(self) -> None:
+        from ..embedding_app import EmbeddingSetupApp
+
+        app = self.app
+        if isinstance(app, EmbeddingSetupApp):
+            provider_screens = {
+                "OpenAI": OpenAIEmbeddingScreen,
+                "Bedrock": BedrockEmbeddingScreen,
+                "Azure": AzureEmbeddingScreen,
+                "Gemini": GeminiEmbeddingScreen,
+                "Cohere": CohereEmbeddingScreen,
+                "HuggingFace": HuggingFaceEmbeddingScreen,
+            }
+            screen_class = provider_screens.get(app.config.provider)
+            if screen_class:
+                app.push_screen(screen_class())

tools/cli/screens/embedding_providers/__init__.py

@@ -0,0 +1,15 @@
+from .openai import OpenAIEmbeddingScreen
+from .bedrock import BedrockEmbeddingScreen
+from .azure import AzureEmbeddingScreen
+from .gemini import GeminiEmbeddingScreen
+from .cohere import CohereEmbeddingScreen
+from .huggingface import HuggingFaceEmbeddingScreen
+
+__all__ = [
+    "OpenAIEmbeddingScreen",
+    "BedrockEmbeddingScreen",
+    "AzureEmbeddingScreen",
+    "GeminiEmbeddingScreen",
+    "CohereEmbeddingScreen",
+    "HuggingFaceEmbeddingScreen",
+]

tools/cli/screens/embedding_providers/azure.py

@@ -0,0 +1,44 @@
+from textual.app import ComposeResult
+from textual.widgets import Input
+
+from llama_index.embeddings.azure_inference import AzureAIEmbeddingsModel
+from llama_cloud import PipelineCreateEmbeddingConfig_AzureEmbedding
+
+from ..base import ConfigurationScreen
+
+
+class AzureEmbeddingScreen(ConfigurationScreen):
+    """Configuration screen for Azure embeddings."""
+
+    def get_title(self) -> str:
+        return "Azure Embedding Configuration"
+
+    def get_form_elements(self) -> list[ComposeResult]:
+        return [
+            Input(
+                placeholder="API Key",
+                password=True,
+                id="api_key",
+                classes="form-control",
+            ),
+            Input(placeholder="Endpoint URL", id="endpoint", classes="form-control"),
+        ]
+
+    def process_submission(self) -> None:
+        api_key = self.query_one("#api_key", Input).value
+        endpoint = self.query_one("#endpoint", Input).value
+
+        if not all([api_key, endpoint]):
+            self.notify("All fields are required.", severity="error")
+            return
+
+        try:
+            embed_model = AzureAIEmbeddingsModel(credential=api_key, endpoint=endpoint)
+            embedding_config = PipelineCreateEmbeddingConfig_AzureEmbedding(
+                type="AZURE_EMBEDDING", component=embed_model
+            )
+
+            self.app.config = embedding_config
+            self.app.handle_completion(self.app.config)
+        except Exception as e:
+            self.notify(f"Error: {e}", severity="error")

tools/cli/screens/embedding_providers/bedrock.py

@@ -0,0 +1,79 @@
+from textual.app import ComposeResult
+from textual.widgets import Input, Select
+
+from llama_index.embeddings.bedrock import BedrockEmbedding
+from llama_cloud import PipelineCreateEmbeddingConfig_BedrockEmbedding
+
+from ..base import ConfigurationScreen
+
+
+class BedrockEmbeddingScreen(ConfigurationScreen):
+    """Configuration screen for Bedrock embeddings."""
+
+    def get_title(self) -> str:
+        return "Bedrock Embedding Configuration"
+
+    def get_form_elements(self) -> list[ComposeResult]:
+        model_options = []
+        try:
+            supported_models = BedrockEmbedding.list_supported_models()
+            model_options = [
+                (f"{provider.title()}: {model_id.split('.')[-1]}", model_id)
+                for provider, models in supported_models.items()
+                for model_id in models
+            ]
+        except Exception as e:
+            self.notify(
+                f"Could not fetch Bedrock models: {e}", severity="error", timeout=10
+            )
+
+        return [
+            Select(
+                options=model_options,
+                prompt="Select Bedrock Model",
+                id="model",
+                classes="form-control",
+            ),
+            Input(
+                placeholder="Region (e.g., us-east-1)",
+                id="region",
+                classes="form-control",
+            ),
+            Input(
+                placeholder="Access Key ID (Optional)",
+                id="access_key_id",
+                classes="form-control",
+            ),
+            Input(
+                placeholder="Secret Access Key (Optional)",
+                password=True,
+                id="secret_access_key",
+                classes="form-control",
+            ),
+        ]
+
+    def process_submission(self) -> None:
+        model = self.query_one("#model", Select).value
+        region = self.query_one("#region", Input).value
+        access_key_id = self.query_one("#access_key_id", Input).value
+        secret_access_key = self.query_one("#secret_access_key", Input).value
+
+        if not all([model, region]):
+            self.notify("All fields are required.", severity="error")
+            return
+
+        try:
+            embed_model = BedrockEmbedding(
+                model_name=model,
+                region_name=region,
+                aws_access_key_id=access_key_id,
+                aws_secret_access_key=secret_access_key,
+            )
+            embedding_config = PipelineCreateEmbeddingConfig_BedrockEmbedding(
+                type="BEDROCK_EMBEDDING", component=embed_model
+            )
+
+            self.app.config = embedding_config
+            self.app.handle_completion(self.app.config)
+        except Exception as e:
+            self.notify(f"Error: {e}", severity="error")