kubrick-cli 0.1.4__py3-none-any.whl

kubrick_cli/providers/base.py

@@ -0,0 +1,136 @@
+"""Base provider adapter interface."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Iterator, List, Optional
+
+
+class ProviderMetadata:
+    """
+    Metadata for a provider that describes its configuration needs.
+
+    This allows providers to be self-describing and enables automatic
+    setup wizard generation and provider discovery.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        display_name: str,
+        description: str,
+        config_fields: List[Dict[str, str]],
+    ):
+        """
+        Initialize provider metadata.
+
+        Args:
+            name: Internal provider name (e.g., "triton", "openai")
+            display_name: Human-readable name (e.g., "Triton", "OpenAI")
+            description: Brief description shown in setup wizard
+            config_fields: List of configuration field definitions
+                Each field should be a dict with:
+                - key: Config key name (e.g., "openai_api_key")
+                - label: Display label (e.g., "OpenAI API key")
+                - type: Field type ("text", "password", "url")
+                - default: Default value (optional)
+                - help_text: Help text or URL (optional)
+        """
+        self.name = name
+        self.display_name = display_name
+        self.description = description
+        self.config_fields = config_fields
+
+
+class ProviderAdapter(ABC):
+    """
+    Abstract base class for LLM provider adapters.
+
+    All provider implementations must inherit from this class and implement
+    the required methods for streaming and non-streaming generation.
+
+    To create a truly plug-and-play provider:
+    1. Inherit from ProviderAdapter
+    2. Define METADATA as a class attribute with ProviderMetadata
+    3. Implement all abstract methods
+    4. Place the file in kubrick_cli/providers/ directory
+
+    The provider will be automatically discovered and available in the setup wizard.
+    """
+
+    METADATA: Optional[ProviderMetadata] = None
+
+    @abstractmethod
+    def generate_streaming(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> Iterator[str]:
+        """
+        Generate streaming response from LLM.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional streaming parameters
+
+        Yields:
+            Text chunks as they arrive
+        """
+        pass
+
+    @abstractmethod
+    def generate(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> str:
+        """
+        Generate non-streaming response from LLM.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional parameters
+
+        Returns:
+            Complete response text
+        """
+        pass
+
+    @abstractmethod
+    def is_healthy(self) -> bool:
+        """
+        Check if the provider is healthy and accessible.
+
+        Returns:
+            True if healthy, False otherwise
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def provider_name(self) -> str:
+        """
+        Get the name of this provider.
+
+        Returns:
+            Provider name (e.g., "triton", "openai", "anthropic")
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def model_name(self) -> str:
+        """
+        Get the model name being used.
+
+        Returns:
+            Model name/identifier
+        """
+        pass
+
+    def set_model(self, model_name: str):
+        """
+        Set the model name dynamically (optional, for model switching).
+
+        Args:
+            model_name: New model name
+
+        Note:
+            Default implementation does nothing. Override in subclasses
+            that support model switching.
+        """
+        pass

kubrick_cli/providers/factory.py

@@ -0,0 +1,161 @@
+"""Provider factory with automatic provider discovery."""
+
+import importlib
+import inspect
+import pkgutil
+from pathlib import Path
+from typing import Dict, List, Type
+
+from .base import ProviderAdapter, ProviderMetadata
+
+
+class ProviderFactory:
+    """
+    Factory for creating provider adapter instances with automatic discovery.
+
+    This factory automatically discovers all provider classes in the providers
+    directory that inherit from ProviderAdapter and have METADATA defined.
+    """
+
+    _provider_registry: Dict[str, Type[ProviderAdapter]] = {}
+    _discovered = False
+
+    @classmethod
+    def _discover_providers(cls):
+        """
+        Discover all provider classes in the providers directory.
+
+        This method scans the providers package for classes that:
+        1. Inherit from ProviderAdapter
+        2. Are not the base ProviderAdapter class itself
+        3. Have a METADATA attribute defined
+        """
+        if cls._discovered:
+            return
+
+        providers_dir = Path(__file__).parent
+
+        for module_info in pkgutil.iter_modules([str(providers_dir)]):
+            module_name = module_info.name
+
+            if module_name in ("base", "factory", "__init__"):
+                continue
+
+            try:
+                module = importlib.import_module(f"kubrick_cli.providers.{module_name}")
+
+                for name, obj in inspect.getmembers(module, inspect.isclass):
+                    if (
+                        issubclass(obj, ProviderAdapter)
+                        and obj is not ProviderAdapter
+                        and hasattr(obj, "METADATA")
+                        and obj.METADATA is not None
+                    ):
+                        provider_name = obj.METADATA.name.lower()
+                        cls._provider_registry[provider_name] = obj
+
+            except Exception as e:
+                print(f"Warning: Failed to load provider from {module_name}: {e}")
+                continue
+
+        cls._discovered = True
+
+    @classmethod
+    def create_provider(cls, config: Dict) -> ProviderAdapter:
+        """
+        Create a provider instance based on configuration.
+
+        Args:
+            config: Configuration dictionary containing provider settings
+
+        Returns:
+            ProviderAdapter instance
+
+        Raises:
+            ValueError: If provider is invalid or required credentials are missing
+        """
+        cls._discover_providers()
+
+        provider_name = config.get("provider", "triton").lower()
+
+        if provider_name not in cls._provider_registry:
+            available = ", ".join(cls._provider_registry.keys())
+            raise ValueError(
+                f"Unknown provider: {provider_name}. "
+                f"Available providers: {available}"
+            )
+
+        provider_class = cls._provider_registry[provider_name]
+        metadata = provider_class.METADATA
+
+        provider_config = {}
+        for field in metadata.config_fields:
+            key = field["key"]
+            value = config.get(key)
+
+            if value is None and "default" not in field:
+                raise ValueError(
+                    f"{metadata.display_name} configuration missing required field: '{key}'. "
+                    f"Please run setup wizard or add '{key}' to config."
+                )
+
+            provider_config[key] = value if value is not None else field.get("default")
+
+        init_signature = inspect.signature(provider_class.__init__)
+        init_params = {}
+
+        for param_name, param in init_signature.parameters.items():
+            if param_name == "self":
+                continue
+
+            for field in metadata.config_fields:
+                if field["key"] == param_name or field["key"].endswith(
+                    f"_{param_name}"
+                ):
+                    init_params[param_name] = provider_config[field["key"]]
+                    break
+
+        return provider_class(**init_params)
+
+    @classmethod
+    def list_available_providers(cls) -> List[ProviderMetadata]:
+        """
+        Get list of available providers with their metadata.
+
+        Returns:
+            List of ProviderMetadata objects
+        """
+        cls._discover_providers()
+
+        providers = []
+        for provider_class in cls._provider_registry.values():
+            if provider_class.METADATA:
+                providers.append(provider_class.METADATA)
+
+        def sort_key(p):
+            return (0, "") if p.name == "triton" else (1, p.name)
+
+        providers.sort(key=sort_key)
+        return providers
+
+    @classmethod
+    def get_provider_metadata(cls, provider_name: str) -> ProviderMetadata:
+        """
+        Get metadata for a specific provider.
+
+        Args:
+            provider_name: Name of the provider
+
+        Returns:
+            ProviderMetadata object
+
+        Raises:
+            ValueError: If provider not found
+        """
+        cls._discover_providers()
+
+        provider_name = provider_name.lower()
+        if provider_name not in cls._provider_registry:
+            raise ValueError(f"Provider not found: {provider_name}")
+
+        return cls._provider_registry[provider_name].METADATA

kubrick_cli/providers/openai_provider.py

@@ -0,0 +1,181 @@
+"""OpenAI provider adapter."""
+
+import http.client
+import json
+import ssl
+from typing import Dict, Iterator, List
+
+from .base import ProviderAdapter, ProviderMetadata
+
+
+class OpenAIProvider(ProviderAdapter):
+    """Provider adapter for OpenAI API."""
+
+    METADATA = ProviderMetadata(
+        name="openai",
+        display_name="OpenAI",
+        description="OpenAI API (GPT-4, GPT-3.5, etc.)",
+        config_fields=[
+            {
+                "key": "openai_api_key",
+                "label": "OpenAI API key",
+                "type": "password",
+                "help_text": "Get your API key from: https://platform.openai.com/api-keys",
+            },
+            {
+                "key": "openai_model",
+                "label": "Model name",
+                "type": "text",
+                "default": "gpt-4",
+            },
+        ],
+    )
+
+    def __init__(self, openai_api_key: str, openai_model: str = "gpt-4"):
+        """
+        Initialize OpenAI provider.
+
+        Args:
+            openai_api_key: OpenAI API key
+            openai_model: Model name (default: gpt-4)
+        """
+        self.api_key = openai_api_key
+        self._model_name = openai_model
+        self.base_url = "api.openai.com"
+        self.timeout = 600
+
+    def generate_streaming(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> Iterator[str]:
+        """
+        Generate streaming response from OpenAI.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional streaming parameters
+
+        Yields:
+            Text chunks as they arrive
+        """
+        payload = {
+            "model": self._model_name,
+            "messages": messages,
+            "stream": True,
+        }
+
+        if stream_options:
+            if "temperature" in stream_options:
+                payload["temperature"] = stream_options["temperature"]
+            if "max_tokens" in stream_options:
+                payload["max_tokens"] = stream_options["max_tokens"]
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+        body = json.dumps(payload).encode("utf-8")
+
+        context = ssl.create_default_context()
+        conn = http.client.HTTPSConnection(
+            self.base_url, 443, timeout=self.timeout, context=context
+        )
+
+        try:
+            conn.request("POST", "/v1/chat/completions", body=body, headers=headers)
+            response = conn.getresponse()
+
+            if response.status != 200:
+                error_body = response.read().decode("utf-8")
+                raise Exception(f"OpenAI API error {response.status}: {error_body}")
+
+            buffer = ""
+            while True:
+                chunk = response.read(1024)
+                if not chunk:
+                    break
+
+                if isinstance(chunk, bytes):
+                    chunk = chunk.decode("utf-8")
+
+                buffer += chunk
+
+                while "\n" in buffer:
+                    line, buffer = buffer.split("\n", 1)
+                    line = line.strip()
+
+                    if not line:
+                        continue
+
+                    if line.startswith("data: "):
+                        line = line[6:]
+
+                    if line == "[DONE]":
+                        return
+
+                    try:
+                        data = json.loads(line)
+                        if "choices" in data and len(data["choices"]) > 0:
+                            delta = data["choices"][0].get("delta", {})
+                            content = delta.get("content", "")
+                            if content:
+                                yield content
+                    except json.JSONDecodeError:
+                        continue
+
+        finally:
+            conn.close()
+
+    def generate(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> str:
+        """
+        Generate non-streaming response from OpenAI.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional parameters
+
+        Returns:
+            Complete response text
+        """
+        chunks = []
+        for chunk in self.generate_streaming(messages, stream_options):
+            chunks.append(chunk)
+        return "".join(chunks)
+
+    def is_healthy(self) -> bool:
+        """
+        Check if OpenAI API is accessible.
+
+        Returns:
+            True if healthy, False otherwise
+        """
+        try:
+            context = ssl.create_default_context()
+            conn = http.client.HTTPSConnection(
+                self.base_url, 443, timeout=10, context=context
+            )
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+            }
+            conn.request("GET", "/v1/models", headers=headers)
+            response = conn.getresponse()
+            conn.close()
+            return response.status == 200
+        except Exception:
+            return False
+
+    @property
+    def provider_name(self) -> str:
+        """Get provider name."""
+        return "openai"
+
+    @property
+    def model_name(self) -> str:
+        """Get model name."""
+        return self._model_name
+
+    def set_model(self, model_name: str):
+        """Set model name dynamically."""
+        self._model_name = model_name

kubrick_cli/providers/triton_provider.py

@@ -0,0 +1,96 @@
+"""Triton provider adapter."""
+
+from typing import Dict, Iterator, List
+
+from ..triton_client import TritonLLMClient
+from .base import ProviderAdapter, ProviderMetadata
+
+
+class TritonProvider(ProviderAdapter):
+    """Provider adapter for Triton Inference Server."""
+
+    METADATA = ProviderMetadata(
+        name="triton",
+        display_name="Triton",
+        description="Self-hosted Triton Inference Server (default)",
+        config_fields=[
+            {
+                "key": "triton_url",
+                "label": "Triton server URL",
+                "type": "url",
+                "default": "localhost:8000",
+            },
+            {
+                "key": "triton_model",
+                "label": "Triton model name",
+                "type": "text",
+                "default": "llm_decoupled",
+            },
+        ],
+    )
+
+    def __init__(self, triton_url: str = "localhost:8000", triton_model: str = "llm_decoupled"):
+        """
+        Initialize Triton provider.
+
+        Args:
+            triton_url: Triton server URL (default: localhost:8000)
+            triton_model: Triton model name (default: llm_decoupled)
+        """
+        self.client = TritonLLMClient(url=triton_url, model_name=triton_model)
+        self._url = triton_url
+        self._model_name = triton_model
+
+    def generate_streaming(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> Iterator[str]:
+        """
+        Generate streaming response from Triton.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional streaming parameters
+
+        Yields:
+            Text chunks as they arrive
+        """
+        yield from self.client.generate_streaming(messages, stream_options)
+
+    def generate(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> str:
+        """
+        Generate non-streaming response from Triton.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional parameters
+
+        Returns:
+            Complete response text
+        """
+        return self.client.generate(messages, stream_options)
+
+    def is_healthy(self) -> bool:
+        """
+        Check if Triton server is healthy.
+
+        Returns:
+            True if healthy, False otherwise
+        """
+        return self.client.is_healthy()
+
+    @property
+    def provider_name(self) -> str:
+        """Get provider name."""
+        return "triton"
+
+    @property
+    def model_name(self) -> str:
+        """Get model name."""
+        return self._model_name
+
+    @property
+    def url(self) -> str:
+        """Get Triton server URL."""
+        return self._url