cli-modelarium 0.1.3__py3-none-any.whl

cli_modelarium/providers/local_provider.py

@@ -0,0 +1,94 @@
+"""Local model provider for Ollama, LM Studio, vLLM, llama.cpp, etc.
+
+All major local LLM tools expose an OpenAI-compatible REST API. This
+provider inherits from `OpenAIProvider` and overrides only:
+
+    * the base URL (default Ollama at localhost:11434)
+    * `_transform_model()` to strip the `local/` prefix before sending to
+      the SDK (`local/llama-3.3-70b` -> `llama-3.3-70b` on the wire)
+    * a localhost-only URL guard - this is a SECURITY feature, not a
+      convenience. The CLI treats "local" as meaning local; a typo'd remote
+      URL would silently send local prompts to a third party. Refused.
+
+Streaming, system prompts, and on_chunk callbacks all work without further
+overrides via the inherited OpenAIProvider behaviour.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import urlparse
+
+import httpx
+
+from cli_modelarium.exceptions import LocalURLError
+from cli_modelarium.providers.openai_provider import OpenAIProvider
+
+
+class LocalProvider(OpenAIProvider):
+    """Provider for any OpenAI-compatible LLM server running on localhost."""
+
+    name: str = "local"
+    DEFAULT_URL: str = "http://localhost:11434/v1"
+    LOCAL_HOSTS: frozenset[str] = frozenset({"localhost", "127.0.0.1", "::1", "0.0.0.0"})
+    DUMMY_API_KEY: str = "not-required"
+    DISCOVERY_TIMEOUT_SECONDS: float = 5.0
+
+    def __init__(self, base_url: str | None = None) -> None:
+        url = base_url or self.DEFAULT_URL
+        validated = self._validate_local_url(url)
+        super().__init__(api_key=self.DUMMY_API_KEY, base_url=validated)
+
+    def _transform_model(self, model: str) -> str:
+        """Strip the `local/` routing prefix before sending to the SDK."""
+        return model.removeprefix("local/")
+
+    @classmethod
+    def _validate_local_url(cls, url: str) -> str:
+        """Return `url` if it parses as an http(s) URL targeting a localhost host.
+
+        Raises LocalURLError otherwise with an actionable message. This is the
+        security boundary - a misconfigured URL must never silently leak local
+        data to a remote endpoint.
+        """
+        parsed = urlparse(url)
+        if parsed.scheme not in ("http", "https"):
+            raise LocalURLError(
+                f"Local URL must use http or https (got scheme {parsed.scheme!r} in {url!r}).\n"
+                f"  Example: http://localhost:11434/v1"
+            )
+        if parsed.hostname not in cls.LOCAL_HOSTS:
+            allowed = ", ".join(sorted(cls.LOCAL_HOSTS))
+            raise LocalURLError(
+                f"Local URL must point to localhost "
+                f"(got hostname {parsed.hostname!r} in {url!r}).\n"
+                f"  Allowed hostnames: {allowed}\n"
+                f"  For remote OpenAI-compatible endpoints, use a real "
+                f"provider integration,\n"
+                f"  not the local provider - this guard prevents accidental "
+                f"data leakage."
+            )
+        return url
+
+    @classmethod
+    async def discover_models(cls, base_url: str | None = None) -> list[dict[str, Any]]:
+        """Query the local server's `/models` endpoint and return the entries.
+
+        Returns the `data` array from an OpenAI-compatible `/models` response:
+            [{"id": "...", "created": ..., "owned_by": "...", ...}, ...]
+
+        URL validation runs first - a non-localhost URL raises LocalURLError
+        before any network I/O happens. Network errors (connection refused,
+        timeout, malformed response) propagate so callers can render a friendly
+        "no local server" message.
+        """
+        url = cls._validate_local_url(base_url or cls.DEFAULT_URL)
+        models_url = url.rstrip("/") + "/models"
+        async with httpx.AsyncClient(timeout=cls.DISCOVERY_TIMEOUT_SECONDS) as client:
+            response = await client.get(models_url)
+            response.raise_for_status()
+            payload = response.json()
+        data = payload.get("data", [])
+        if not isinstance(data, list):
+            return []
+        return data

cli_modelarium/providers/mistral_provider.py

@@ -0,0 +1,172 @@
+"""Mistral provider using the official mistralai SDK.
+
+Notes on the SDK layout in mistralai 2.4.7:
+
+    The spec's `from mistralai import Mistral` import path does not exist
+    in 2.4.7 - the top-level package is now a namespace package with
+    subpackages for each cloud host (azure, gcp, generic). The canonical
+    client lives at `mistralai.client.Mistral`.
+
+    Stream events wrap the chunk in a `CompletionEvent` whose `.data`
+    attribute is the OpenAI-style `CompletionChunk` with the usual
+    `choices[0].delta.content` shape.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from mistralai.client import Mistral
+from mistralai.client.errors.sdkerror import SDKError
+
+from cli_modelarium.exceptions import (
+    AuthenticationError,
+    ProviderError,
+    RateLimitError,
+)
+from cli_modelarium.pricing import calculate_cost
+from cli_modelarium.providers._utils import extract_retry_after
+from cli_modelarium.providers.base import BaseProvider, CompletionResult, OnChunk
+from cli_modelarium.security import redact_secrets
+
+
+class MistralProvider(BaseProvider):
+    """Provider using the official Mistral Python SDK."""
+
+    name: str = "mistral"
+
+    def __init__(self, api_key: str) -> None:
+        self.client = Mistral(api_key=api_key)
+
+    def _build_messages(self, prompt: str, system_prompt: str | None) -> list[dict[str, str]]:
+        messages: list[dict[str, str]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+        return messages
+
+    async def stream(
+        self,
+        prompt: str,
+        model: str,
+        temperature: float,
+        system_prompt: str | None = None,
+    ) -> AsyncIterator[str]:
+        messages = self._build_messages(prompt, system_prompt)
+        try:
+            response_stream = await self.client.chat.stream_async(
+                model=model,
+                messages=messages,
+                temperature=temperature,
+            )
+            async for event in response_stream:
+                content = _extract_chunk_text(event)
+                if content:
+                    yield content
+        except SDKError as e:
+            self._reraise(e)
+
+    async def complete(
+        self,
+        prompt: str,
+        model: str,
+        temperature: float,
+        system_prompt: str | None = None,
+        *,
+        on_chunk: OnChunk | None = None,
+    ) -> CompletionResult:
+        messages = self._build_messages(prompt, system_prompt)
+
+        start = time.monotonic()
+        ttft_ms: float | None = None
+        chunks: list[str] = []
+        input_tokens = 0
+        output_tokens = 0
+
+        try:
+            response_stream = await self.client.chat.stream_async(
+                model=model,
+                messages=messages,
+                temperature=temperature,
+            )
+            async for event in response_stream:
+                content = _extract_chunk_text(event)
+                if content:
+                    if ttft_ms is None:
+                        ttft_ms = (time.monotonic() - start) * 1000
+                    if on_chunk is not None:
+                        on_chunk(content)
+                    chunks.append(content)
+                usage = _extract_chunk_usage(event)
+                if usage is not None:
+                    input_tokens, output_tokens = usage
+        except SDKError as e:
+            self._reraise(e)
+
+        latency_ms = (time.monotonic() - start) * 1000
+
+        try:
+            cost = calculate_cost(model, input_tokens, output_tokens)
+        except Exception:
+            cost = 0.0
+
+        return CompletionResult(
+            output="".join(chunks),
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cost_usd=cost,
+            latency_ms=latency_ms,
+            ttft_ms=ttft_ms,
+            model=model,
+            provider=self.name,
+            temperature=temperature,
+        )
+
+    def _reraise(self, error: SDKError) -> None:
+        """Translate a Mistral SDK error into a redacted Cli Modelarium exception."""
+        message = redact_secrets(str(error))
+        status_code = _extract_status_code(error)
+
+        if status_code in (401, 403):
+            raise AuthenticationError(message, provider=self.name) from None
+        if status_code == 429:
+            retry_after = extract_retry_after(error)
+            raise RateLimitError(message, provider=self.name, retry_after=retry_after) from None
+        raise ProviderError(message, provider=self.name) from None
+
+
+def _extract_chunk_text(event: Any) -> str | None:
+    """Pull `data.choices[0].delta.content` out of a streamed CompletionEvent."""
+    data = getattr(event, "data", None)
+    if data is None:
+        return None
+    choices = getattr(data, "choices", None)
+    if not choices:
+        return None
+    delta = getattr(choices[0], "delta", None)
+    if delta is None:
+        return None
+    return getattr(delta, "content", None)
+
+
+def _extract_chunk_usage(event: Any) -> tuple[int, int] | None:
+    """Pull (input_tokens, output_tokens) out of a streamed CompletionEvent's usage payload."""
+    data = getattr(event, "data", None)
+    if data is None:
+        return None
+    usage = getattr(data, "usage", None)
+    if usage is None:
+        return None
+    input_tokens = getattr(usage, "prompt_tokens", 0) or 0
+    output_tokens = getattr(usage, "completion_tokens", 0) or 0
+    return input_tokens, output_tokens
+
+
+def _extract_status_code(error: SDKError) -> int | None:
+    """Pull the HTTP status code from an SDKError's wrapped httpx response."""
+    response = getattr(error, "raw_response", None)
+    if response is None:
+        return None
+    return getattr(response, "status_code", None)

cli_modelarium/providers/openai_provider.py

@@ -0,0 +1,163 @@
+"""OpenAI provider implementation.
+
+Designed to be subclassed by the OpenAI-compatible providers (xAI, DeepSeek,
+Groq, OpenRouter, Local) by overriding `name`, the constructor (to pass a
+different base_url), and optionally `_transform_model()`.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+import openai
+from openai import AsyncOpenAI
+
+from cli_modelarium.exceptions import (
+    AuthenticationError,
+    ProviderError,
+    ProviderOverloadedError,
+    RateLimitError,
+)
+from cli_modelarium.pricing import calculate_cost
+from cli_modelarium.providers._utils import extract_retry_after
+from cli_modelarium.providers.base import BaseProvider, CompletionResult, OnChunk
+from cli_modelarium.security import redact_secrets
+
+
+class OpenAIProvider(BaseProvider):
+    """Provider using the official OpenAI Python SDK."""
+
+    name: str = "openai"
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str | None = None,
+        default_headers: dict[str, str] | None = None,
+    ) -> None:
+        client_kwargs: dict[str, Any] = {"api_key": api_key}
+        if base_url is not None:
+            client_kwargs["base_url"] = base_url
+        if default_headers is not None:
+            client_kwargs["default_headers"] = default_headers
+        self.client = AsyncOpenAI(**client_kwargs)
+
+    def _transform_model(self, model: str) -> str:
+        """Subclass hook for providers that need to rewrite the model ID (e.g. LocalProvider)."""
+        return model
+
+    def _build_messages(self, prompt: str, system_prompt: str | None) -> list[dict[str, str]]:
+        messages: list[dict[str, str]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+        return messages
+
+    async def stream(
+        self,
+        prompt: str,
+        model: str,
+        temperature: float,
+        system_prompt: str | None = None,
+    ) -> AsyncIterator[str]:
+        messages = self._build_messages(prompt, system_prompt)
+        actual_model = self._transform_model(model)
+        try:
+            response = await self.client.chat.completions.create(
+                model=actual_model,
+                messages=messages,
+                temperature=temperature,
+                stream=True,
+            )
+            async for chunk in response:
+                if chunk.choices and chunk.choices[0].delta.content:
+                    yield chunk.choices[0].delta.content
+        except openai.APIError as e:
+            self._reraise(e)
+
+    async def complete(
+        self,
+        prompt: str,
+        model: str,
+        temperature: float,
+        system_prompt: str | None = None,
+        *,
+        on_chunk: OnChunk | None = None,
+    ) -> CompletionResult:
+        """Run a completion via streaming with usage included.
+
+        Streaming with `stream_options={"include_usage": True}` gives us both
+        token-by-token output (for TTFT) and final usage numbers in one call.
+
+        If `on_chunk` is provided it is called with each text chunk as it
+        arrives, so the streaming orchestrator can surface output live.
+        """
+        messages = self._build_messages(prompt, system_prompt)
+        actual_model = self._transform_model(model)
+
+        start = time.monotonic()
+        ttft_ms: float | None = None
+        chunks: list[str] = []
+        input_tokens = 0
+        output_tokens = 0
+        cached_tokens = 0
+
+        try:
+            response = await self.client.chat.completions.create(
+                model=actual_model,
+                messages=messages,
+                temperature=temperature,
+                stream=True,
+                stream_options={"include_usage": True},
+            )
+            async for chunk in response:
+                if chunk.choices and chunk.choices[0].delta.content:
+                    text = chunk.choices[0].delta.content
+                    if ttft_ms is None:
+                        ttft_ms = (time.monotonic() - start) * 1000
+                    if on_chunk is not None:
+                        on_chunk(text)
+                    chunks.append(text)
+                if getattr(chunk, "usage", None) is not None:
+                    input_tokens = chunk.usage.prompt_tokens or 0
+                    output_tokens = chunk.usage.completion_tokens or 0
+                    details = getattr(chunk.usage, "prompt_tokens_details", None)
+                    if details is not None:
+                        cached_tokens = getattr(details, "cached_tokens", 0) or 0
+        except openai.APIError as e:
+            self._reraise(e)
+
+        latency_ms = (time.monotonic() - start) * 1000
+
+        try:
+            cost = calculate_cost(model, input_tokens, output_tokens, cached_tokens)
+        except Exception:
+            cost = 0.0
+
+        return CompletionResult(
+            output="".join(chunks),
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cached_tokens=cached_tokens,
+            cost_usd=cost,
+            latency_ms=latency_ms,
+            ttft_ms=ttft_ms,
+            model=model,
+            provider=self.name,
+            temperature=temperature,
+        )
+
+    def _reraise(self, error: openai.APIError) -> None:
+        """Translate an OpenAI SDK error into a redacted Cli Modelarium exception."""
+        message = redact_secrets(str(error))
+
+        if isinstance(error, openai.AuthenticationError):
+            raise AuthenticationError(message, provider=self.name) from None
+        if isinstance(error, openai.RateLimitError):
+            retry_after = extract_retry_after(error)
+            raise RateLimitError(message, provider=self.name, retry_after=retry_after) from None
+        if isinstance(error, openai.APIStatusError) and getattr(error, "status_code", None) == 529:
+            raise ProviderOverloadedError(message, provider=self.name) from None
+        raise ProviderError(message, provider=self.name) from None

cli_modelarium/providers/openrouter_provider.py

@@ -0,0 +1,33 @@
+"""OpenRouter provider - OpenAI-compatible meta-provider with custom headers.
+
+OpenRouter aggregates models from many vendors behind one OpenAI-compatible
+API. They expect two extra headers identifying the calling application:
+
+    HTTP-Referer  - URL of the calling app (for analytics/leaderboard)
+    X-Title       - human-readable app name
+
+We pass these as `default_headers` on the AsyncOpenAI client so they go on
+every request without per-call wiring.
+"""
+
+from __future__ import annotations
+
+from cli_modelarium.providers.openai_provider import OpenAIProvider
+
+
+class OpenRouterProvider(OpenAIProvider):
+    """OpenRouter models via the OpenAI-compatible endpoint at openrouter.ai."""
+
+    name: str = "openrouter"
+    BASE_URL: str = "https://openrouter.ai/api/v1"
+    DEFAULT_HEADERS: dict[str, str] = {
+        "HTTP-Referer": "https://github.com/lavellehatcherjr/cli-modelarium",
+        "X-Title": "Cli Modelarium",
+    }
+
+    def __init__(self, api_key: str) -> None:
+        super().__init__(
+            api_key=api_key,
+            base_url=self.BASE_URL,
+            default_headers=self.DEFAULT_HEADERS,
+        )

cli_modelarium/providers/xai_provider.py

@@ -0,0 +1,15 @@
+"""xAI (Grok) provider - uses the OpenAI SDK with a different base URL."""
+
+from __future__ import annotations
+
+from cli_modelarium.providers.openai_provider import OpenAIProvider
+
+
+class XAIProvider(OpenAIProvider):
+    """xAI Grok models via the OpenAI-compatible endpoint at api.x.ai."""
+
+    name: str = "xai"
+    BASE_URL: str = "https://api.x.ai/v1"
+
+    def __init__(self, api_key: str) -> None:
+        super().__init__(api_key=api_key, base_url=self.BASE_URL)