mirascope 2.0.0a3__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

mirascope/llm/providers/google/model_id.py

@@ -1,20 +1,14 @@
 """Google registered LLM models."""
 
-from typing import Literal, TypeAlias
-
-GoogleModelId: TypeAlias = (
-    Literal[
-        "google/gemini-3-pro-preview",
-        "google/gemini-2.5-pro",
-        "google/gemini-2.5-flash",
-        "google/gemini-2.5-flash-lite",
-        "google/gemini-2.0-flash",
-        "google/gemini-2.0-flash-lite",
-    ]
-    | str
-)
+from typing import TypeAlias, get_args
+
+from .model_info import GoogleKnownModels
+
+GoogleModelId: TypeAlias = GoogleKnownModels | str
 """The Google model ids registered with Mirascope."""
 
+GOOGLE_KNOWN_MODELS: set[str] = set(get_args(GoogleKnownModels))
+
 
 def model_name(model_id: GoogleModelId) -> str:
     """Extract the google model name from a full model ID.

mirascope/llm/providers/google/model_info.py

@@ -0,0 +1,62 @@
+"""Google model information.
+
+This file is auto-generated by scripts/model_features/codegen_google.py
+Do not edit manually - run the codegen script to update."""
+
+from typing import Literal
+
+GoogleKnownModels = Literal[
+    "google/gemini-2.0-flash",
+    "google/gemini-2.0-flash-001",
+    "google/gemini-2.0-flash-exp",
+    "google/gemini-2.0-flash-exp-image-generation",
+    "google/gemini-2.0-flash-lite",
+    "google/gemini-2.0-flash-lite-001",
+    "google/gemini-2.0-flash-lite-preview",
+    "google/gemini-2.0-flash-lite-preview-02-05",
+    "google/gemini-2.5-flash",
+    "google/gemini-2.5-flash-image",
+    "google/gemini-2.5-flash-image-preview",
+    "google/gemini-2.5-flash-lite",
+    "google/gemini-2.5-flash-lite-preview-09-2025",
+    "google/gemini-2.5-flash-preview-09-2025",
+    "google/gemini-2.5-pro",
+    "google/gemini-3-pro-image-preview",
+    "google/gemini-3-pro-preview",
+    "google/gemini-flash-latest",
+    "google/gemini-flash-lite-latest",
+    "google/gemini-pro-latest",
+    "google/gemini-robotics-er-1.5-preview",
+    "google/gemma-3-12b-it",
+    "google/gemma-3-1b-it",
+    "google/gemma-3-27b-it",
+    "google/gemma-3-4b-it",
+    "google/gemma-3n-e2b-it",
+    "google/gemma-3n-e4b-it",
+    "google/nano-banana-pro-preview",
+]
+"""Valid Google model IDs."""
+
+
+MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT: set[str] = {
+    "gemini-2.5-flash",
+    "gemini-2.5-flash-image",
+    "gemini-2.5-flash-image-preview",
+    "gemini-2.5-flash-lite",
+    "gemini-2.5-flash-lite-preview-09-2025",
+    "gemini-2.5-flash-preview-09-2025",
+    "gemini-2.5-pro",
+    "gemini-3-pro-image-preview",
+    "gemini-flash-latest",
+    "gemini-flash-lite-latest",
+    "gemini-pro-latest",
+    "gemini-robotics-er-1.5-preview",
+    "gemma-3-12b-it",
+    "gemma-3-1b-it",
+    "gemma-3-27b-it",
+    "gemma-3-4b-it",
+    "gemma-3n-e2b-it",
+    "gemma-3n-e4b-it",
+    "nano-banana-pro-preview",
+}
+"""Models that do not support structured outputs when tools are present."""

mirascope/llm/providers/google/provider.py

@@ -81,7 +81,7 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_response = self.client.models.generate_content(**kwargs)
 
-        assistant_message, finish_reason = _utils.decode_response(
+        assistant_message, finish_reason, usage = _utils.decode_response(
             google_response, model_id
         )
 
@@ -95,6 +95,7 @@ class GoogleProvider(BaseProvider[Client]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
+            usage=usage,
             format=format,
         )
 
@@ -133,7 +134,7 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_response = self.client.models.generate_content(**kwargs)
 
-        assistant_message, finish_reason = _utils.decode_response(
+        assistant_message, finish_reason, usage = _utils.decode_response(
             google_response, model_id
         )
 
@@ -147,6 +148,7 @@ class GoogleProvider(BaseProvider[Client]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
+            usage=usage,
             format=format,
         )
 
@@ -181,7 +183,7 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_response = await self.client.aio.models.generate_content(**kwargs)
 
-        assistant_message, finish_reason = _utils.decode_response(
+        assistant_message, finish_reason, usage = _utils.decode_response(
             google_response, model_id
         )
 
@@ -195,6 +197,7 @@ class GoogleProvider(BaseProvider[Client]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
+            usage=usage,
             format=format,
         )
 
@@ -233,7 +236,7 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_response = await self.client.aio.models.generate_content(**kwargs)
 
-        assistant_message, finish_reason = _utils.decode_response(
+        assistant_message, finish_reason, usage = _utils.decode_response(
             google_response, model_id
         )
 
@@ -247,6 +250,7 @@ class GoogleProvider(BaseProvider[Client]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=finish_reason,
+            usage=usage,
             format=format,
         )
 

mirascope/llm/providers/load_provider.py

@@ -4,10 +4,12 @@ from .anthropic import AnthropicProvider
 from .base import Provider
 from .google import GoogleProvider
 from .mlx import MLXProvider
+from .ollama import OllamaProvider
 from .openai import OpenAIProvider
 from .openai.completions.provider import OpenAICompletionsProvider
 from .openai.responses.provider import OpenAIResponsesProvider
 from .provider_id import ProviderId
+from .together import TogetherProvider
 
 
 @lru_cache(maxsize=256)
@@ -32,14 +34,18 @@ def load_provider(
             return AnthropicProvider(api_key=api_key, base_url=base_url)
         case "google":
             return GoogleProvider(api_key=api_key, base_url=base_url)
+        case "mlx":  # pragma: no cover (MLX is only available on macOS)
+            return MLXProvider()
+        case "ollama":
+            return OllamaProvider(api_key=api_key, base_url=base_url)
         case "openai":
             return OpenAIProvider(api_key=api_key, base_url=base_url)
         case "openai:completions":
             return OpenAICompletionsProvider(api_key=api_key, base_url=base_url)
         case "openai:responses":
             return OpenAIResponsesProvider(api_key=api_key, base_url=base_url)
-        case "mlx":  # pragma: no cover (MLX is only available on macOS)
-            return MLXProvider()
+        case "together":
+            return TogetherProvider(api_key=api_key, base_url=base_url)
         case _:  # pragma: no cover
             raise ValueError(f"Unknown provider: '{provider_id}'")
 

mirascope/llm/providers/mlx/_utils.py

@@ -5,7 +5,7 @@ import mlx.core as mx
 from mlx_lm.generate import GenerationResponse
 from mlx_lm.sample_utils import make_sampler
 
-from ...responses import FinishReason
+from ...responses import FinishReason, Usage
 from ..base import Params, _utils as _base_utils
 
 Sampler: TypeAlias = Callable[[mx.array], mx.array]
@@ -105,3 +105,25 @@ def extract_finish_reason(response: GenerationResponse | None) -> FinishReason |
         return FinishReason.MAX_TOKENS
 
     return None
+
+
+def extract_usage(response: GenerationResponse | None) -> Usage | None:
+    """Extract usage information from an MLX generation response.
+
+    Args:
+        response: The MLX generation response to extract from.
+
+    Returns:
+        The Usage object with token counts, or None if not applicable.
+    """
+    if response is None:
+        return None
+
+    return Usage(
+        input_tokens=response.prompt_tokens,
+        output_tokens=response.generation_tokens,
+        cache_read_tokens=0,
+        cache_write_tokens=0,
+        reasoning_tokens=0,
+        raw=response,
+    )

mirascope/llm/providers/mlx/encoding/transformers.py

@@ -10,7 +10,12 @@ from transformers import PreTrainedTokenizer
 from ....content import ContentPart, TextChunk, TextEndChunk, TextStartChunk
 from ....formatting import Format, FormattableT
 from ....messages import AssistantContent, Message
-from ....responses import ChunkIterator, FinishReasonChunk, RawStreamEventChunk
+from ....responses import (
+    ChunkIterator,
+    FinishReasonChunk,
+    RawStreamEventChunk,
+    UsageDeltaChunk,
+)
 from ....tools import AnyToolSchema, BaseToolkit
 from .. import _utils
 from .base import BaseEncoder, TokenIds
@@ -129,3 +134,14 @@ class TransformersEncoder(BaseEncoder):
             yield FinishReasonChunk(finish_reason=finish_reason)
         else:
             yield TextEndChunk()
+
+        # Emit usage delta if available
+        usage = _utils.extract_usage(response)
+        if usage:
+            yield UsageDeltaChunk(
+                input_tokens=usage.input_tokens,
+                output_tokens=usage.output_tokens,
+                cache_read_tokens=usage.cache_read_tokens,
+                cache_write_tokens=usage.cache_write_tokens,
+                reasoning_tokens=usage.reasoning_tokens,
+            )

mirascope/llm/providers/mlx/provider.py

@@ -108,6 +108,7 @@ class MLXProvider(BaseProvider[None]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
             format=format,
         )
 
@@ -152,6 +153,7 @@ class MLXProvider(BaseProvider[None]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
             format=format,
         )
 
@@ -196,6 +198,7 @@ class MLXProvider(BaseProvider[None]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
             format=format,
         )
 
@@ -244,6 +247,7 @@ class MLXProvider(BaseProvider[None]):
             input_messages=input_messages,
             assistant_message=assistant_message,
             finish_reason=_utils.extract_finish_reason(response),
+            usage=_utils.extract_usage(response),
             format=format,
         )
 

mirascope/llm/providers/ollama/__init__.py

@@ -0,0 +1,19 @@
+"""Ollama provider implementation."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .provider import OllamaProvider
+else:
+    try:
+        from .provider import OllamaProvider
+    except ImportError:  # pragma: no cover
+        from .._missing_import_stubs import (
+            create_provider_stub,
+        )
+
+        OllamaProvider = create_provider_stub("openai", "OllamaProvider")
+
+__all__ = [
+    "OllamaProvider",
+]

mirascope/llm/providers/ollama/provider.py

@@ -0,0 +1,71 @@
+"""Ollama provider implementation."""
+
+import os
+from typing import ClassVar
+
+from openai import AsyncOpenAI, OpenAI
+
+from ..openai.completions.base_provider import BaseOpenAICompletionsProvider
+
+
+class OllamaProvider(BaseOpenAICompletionsProvider):
+    """Provider for Ollama's OpenAI-compatible API.
+
+    Inherits from BaseOpenAICompletionsProvider with Ollama-specific configuration:
+    - Uses Ollama's local API endpoint (default: http://localhost:11434/v1/)
+    - API key is not required (Ollama ignores API keys)
+    - Supports OLLAMA_BASE_URL environment variable
+
+    Usage:
+        Register the provider with model ID prefixes you want to use:
+
+        ```python
+        import llm
+
+        # Register for ollama models
+        llm.register_provider("ollama", "ollama/")
+
+        # Now you can use ollama models directly
+        @llm.call("ollama/llama2")
+        def my_prompt():
+            return [llm.messages.user("Hello!")]
+        ```
+    """
+
+    id: ClassVar[str] = "ollama"
+    default_scope: ClassVar[str | list[str]] = "ollama/"
+    default_base_url: ClassVar[str | None] = "http://localhost:11434/v1/"
+    api_key_env_var: ClassVar[str] = "OLLAMA_API_KEY"
+    api_key_required: ClassVar[bool] = False
+    provider_name: ClassVar[str | None] = "Ollama"
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+    ) -> None:
+        """Initialize the Ollama provider.
+
+        Args:
+            api_key: API key (optional). Defaults to OLLAMA_API_KEY env var or 'ollama'.
+            base_url: Custom base URL. Defaults to OLLAMA_BASE_URL env var
+                or http://localhost:11434/v1/.
+        """
+        resolved_api_key = api_key or os.environ.get(self.api_key_env_var) or "ollama"
+        resolved_base_url = (
+            base_url or os.environ.get("OLLAMA_BASE_URL") or self.default_base_url
+        )
+
+        self.client = OpenAI(
+            api_key=resolved_api_key,
+            base_url=resolved_base_url,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=resolved_api_key,
+            base_url=resolved_base_url,
+        )
+
+    def _model_name(self, model_id: str) -> str:
+        """Strip 'ollama/' prefix from model ID for Ollama API."""
+        return model_id.removeprefix("ollama/")

mirascope/llm/providers/openai/completions/__init__.py

@@ -1,20 +1,25 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from .base_provider import BaseOpenAICompletionsProvider
     from .provider import OpenAICompletionsProvider
 else:
     try:
+        from .base_provider import BaseOpenAICompletionsProvider
         from .provider import OpenAICompletionsProvider
     except ImportError:  # pragma: no cover
         from ..._missing_import_stubs import (
-            create_import_error_stub,
             create_provider_stub,
         )
 
+        BaseOpenAICompletionsProvider = create_provider_stub(
+            "openai", "BaseOpenAICompletionsProvider"
+        )
         OpenAICompletionsProvider = create_provider_stub(
             "openai", "OpenAICompletionsProvider"
         )
 
 __all__ = [
+    "BaseOpenAICompletionsProvider",
     "OpenAICompletionsProvider",
 ]

mirascope/llm/providers/openai/completions/_utils/decode.py

@@ -4,6 +4,7 @@ from typing import Literal
 
 from openai import AsyncStream, Stream
 from openai.types import chat as openai_types
+from openai.types.completion_usage import CompletionUsage
 
 from .....content import (
     AssistantContentPart,
@@ -23,6 +24,8 @@ from .....responses import (
     FinishReason,
     FinishReasonChunk,
     RawStreamEventChunk,
+    Usage,
+    UsageDeltaChunk,
 )
 from ...model_id import OpenAIModelId, model_name
 
@@ -32,12 +35,40 @@ OPENAI_FINISH_REASON_MAP = {
 }
 
 
+def _decode_usage(
+    usage: CompletionUsage | None,
+) -> Usage | None:
+    """Convert OpenAI CompletionUsage to Mirascope Usage."""
+    if usage is None:  # pragma: no cover
+        return None
+
+    return Usage(
+        input_tokens=usage.prompt_tokens,
+        output_tokens=usage.completion_tokens,
+        cache_read_tokens=(
+            usage.prompt_tokens_details.cached_tokens
+            if usage.prompt_tokens_details
+            else None
+        )
+        or 0,
+        cache_write_tokens=0,
+        reasoning_tokens=(
+            usage.completion_tokens_details.reasoning_tokens
+            if usage.completion_tokens_details
+            else None
+        )
+        or 0,
+        raw=usage,
+    )
+
+
 def decode_response(
     response: openai_types.ChatCompletion,
     model_id: OpenAIModelId,
-    provider_id: Literal["openai", "openai:completions"],
-) -> tuple[AssistantMessage, FinishReason | None]:
-    """Convert OpenAI ChatCompletion to mirascope AssistantMessage."""
+    provider_id: str,
+    provider_model_name: str | None = None,
+) -> tuple[AssistantMessage, FinishReason | None, Usage | None]:
+    """Convert OpenAI ChatCompletion to mirascope AssistantMessage and usage."""
     choice = response.choices[0]
     message = choice.message
     refused = False
@@ -72,11 +103,12 @@ def decode_response(
         content=parts,
         provider_id=provider_id,
         model_id=model_id,
-        provider_model_name=model_name(model_id, "completions"),
+        provider_model_name=provider_model_name or model_name(model_id, "completions"),
         raw_message=message.model_dump(exclude_none=True),
     )
 
-    return assistant_message, finish_reason
+    usage = _decode_usage(response.usage)
+    return assistant_message, finish_reason, usage
 
 
 class _OpenAIChunkProcessor:
@@ -91,6 +123,26 @@ class _OpenAIChunkProcessor:
         """Process a single OpenAI chunk and yield the appropriate content chunks."""
         yield RawStreamEventChunk(raw_stream_event=chunk)
 
+        if chunk.usage:
+            usage = chunk.usage
+            yield UsageDeltaChunk(
+                input_tokens=usage.prompt_tokens,
+                output_tokens=usage.completion_tokens,
+                cache_read_tokens=(
+                    usage.prompt_tokens_details.cached_tokens
+                    if usage.prompt_tokens_details
+                    else None
+                )
+                or 0,
+                cache_write_tokens=0,
+                reasoning_tokens=(
+                    usage.completion_tokens_details.reasoning_tokens
+                    if usage.completion_tokens_details
+                    else None
+                )
+                or 0,
+            )
+
         choice = chunk.choices[0] if chunk.choices else None
         if not choice:
             return  # pragma: no cover

mirascope/llm/providers/openai/completions/_utils/encode.py

@@ -22,8 +22,11 @@ from .....messages import AssistantMessage, Message, UserMessage
 from .....tools import FORMAT_TOOL_NAME, AnyToolSchema, BaseToolkit
 from ....base import Params, _utils as _base_utils
 from ...model_id import OpenAIModelId, model_name
-from ...model_info import MODELS_WITHOUT_AUDIO_SUPPORT
-from ...shared import _utils as _shared_utils
+from ...model_info import (
+    MODELS_WITHOUT_AUDIO_SUPPORT,
+    MODELS_WITHOUT_JSON_OBJECT_SUPPORT,
+    MODELS_WITHOUT_JSON_SCHEMA_SUPPORT,
+)
 
 
 class ChatCompletionCreateKwargs(TypedDict, total=False):
@@ -233,7 +236,7 @@ def _convert_tool_to_tool_param(
     """Convert a single Mirascope `Tool` to OpenAI ChatCompletionToolParam with caching."""
     schema_dict = tool.parameters.model_dump(by_alias=True, exclude_none=True)
     schema_dict["type"] = "object"
-    _shared_utils.ensure_additional_properties_false(schema_dict)
+    _base_utils.ensure_additional_properties_false(schema_dict)
     return openai_types.ChatCompletionToolParam(
         type="function",
         function={
@@ -258,7 +261,7 @@ def _create_strict_response_format(
     """
     schema = format.schema.copy()
 
-    _shared_utils.ensure_additional_properties_false(schema)
+    _base_utils.ensure_additional_properties_false(schema)
 
     json_schema = JSONSchema(
         name=format.name,
@@ -321,9 +324,7 @@ def encode_request(
 
     openai_tools = [_convert_tool_to_tool_param(tool) for tool in tools]
 
-    model_supports_strict = (
-        base_model_name not in _shared_utils.MODELS_WITHOUT_JSON_SCHEMA_SUPPORT
-    )
+    model_supports_strict = base_model_name not in MODELS_WITHOUT_JSON_SCHEMA_SUPPORT
     default_mode = "strict" if model_supports_strict else "tool"
     format = resolve_format(format, default_mode=default_mode)
     if format is not None:
@@ -348,7 +349,7 @@ def encode_request(
             openai_tools.append(_convert_tool_to_tool_param(format_tool_schema))
         elif (
             format.mode == "json"
-            and base_model_name not in _shared_utils.MODELS_WITHOUT_JSON_OBJECT_SUPPORT
+            and base_model_name not in MODELS_WITHOUT_JSON_OBJECT_SUPPORT
         ):
             kwargs["response_format"] = {"type": "json_object"}