mirascope 1.22.4__py3-none-any.whl → 2.1.1__py3-none-any.whl

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/__init__.py

@@ -0,0 +1,15 @@
+"""OpenAI client implementation."""
+
+from .completions.base_provider import BaseOpenAICompletionsProvider
+from .completions.provider import OpenAICompletionsProvider
+from .model_id import OpenAIModelId
+from .provider import OpenAIProvider
+from .responses.provider import OpenAIResponsesProvider
+
+__all__ = [
+    "BaseOpenAICompletionsProvider",
+    "OpenAICompletionsProvider",
+    "OpenAIModelId",
+    "OpenAIProvider",
+    "OpenAIResponsesProvider",
+]

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

@@ -0,0 +1,5 @@
+"""Shared OpenAI utilities."""
+
+from .errors import OPENAI_ERROR_MAP
+
+__all__ = ["OPENAI_ERROR_MAP"]

mirascope/llm/providers/openai/_utils/errors.py

@@ -0,0 +1,46 @@
+"""OpenAI error handling utilities."""
+
+from openai import (
+    APIConnectionError as OpenAIAPIConnectionError,
+    APIResponseValidationError as OpenAIAPIResponseValidationError,
+    APITimeoutError as OpenAIAPITimeoutError,
+    AuthenticationError as OpenAIAuthenticationError,
+    BadRequestError as OpenAIBadRequestError,
+    ConflictError as OpenAIConflictError,
+    InternalServerError as OpenAIInternalServerError,
+    NotFoundError as OpenAINotFoundError,
+    OpenAIError,
+    PermissionDeniedError as OpenAIPermissionDeniedError,
+    RateLimitError as OpenAIRateLimitError,
+    UnprocessableEntityError as OpenAIUnprocessableEntityError,
+)
+
+from ....exceptions import (
+    AuthenticationError,
+    BadRequestError,
+    ConnectionError,
+    NotFoundError,
+    PermissionError,
+    ProviderError,
+    RateLimitError,
+    ResponseValidationError,
+    ServerError,
+    TimeoutError,
+)
+from ...base import ProviderErrorMap
+
+# Shared error mapping used by OpenAI Responses and Completions providers
+OPENAI_ERROR_MAP: ProviderErrorMap = {
+    OpenAIAuthenticationError: AuthenticationError,
+    OpenAIPermissionDeniedError: PermissionError,
+    OpenAINotFoundError: NotFoundError,
+    OpenAIBadRequestError: BadRequestError,
+    OpenAIUnprocessableEntityError: BadRequestError,
+    OpenAIConflictError: BadRequestError,
+    OpenAIRateLimitError: RateLimitError,
+    OpenAIInternalServerError: ServerError,
+    OpenAIAPITimeoutError: TimeoutError,
+    OpenAIAPIConnectionError: ConnectionError,
+    OpenAIAPIResponseValidationError: ResponseValidationError,
+    OpenAIError: ProviderError,  # Catch-all for unknown OpenAI errors
+}

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

@@ -0,0 +1,7 @@
+from .base_provider import BaseOpenAICompletionsProvider
+from .provider import OpenAICompletionsProvider
+
+__all__ = [
+    "BaseOpenAICompletionsProvider",
+    "OpenAICompletionsProvider",
+]

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

@@ -0,0 +1,18 @@
+from .decode import (
+    decode_async_stream,
+    decode_response,
+    decode_stream,
+    model_name,
+)
+from .encode import encode_request
+from .feature_info import CompletionsModelFeatureInfo, feature_info_for_openai_model
+
+__all__ = [
+    "CompletionsModelFeatureInfo",
+    "decode_async_stream",
+    "decode_response",
+    "decode_stream",
+    "encode_request",
+    "feature_info_for_openai_model",
+    "model_name",
+]

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

@@ -0,0 +1,252 @@
+"""OpenAI completions response decoding."""
+
+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,
+    Text,
+    TextChunk,
+    TextEndChunk,
+    TextStartChunk,
+    ToolCall,
+    ToolCallChunk,
+    ToolCallEndChunk,
+    ToolCallStartChunk,
+)
+from .....messages import AssistantMessage
+from .....responses import (
+    AsyncChunkIterator,
+    ChunkIterator,
+    FinishReason,
+    FinishReasonChunk,
+    RawStreamEventChunk,
+    Usage,
+    UsageDeltaChunk,
+)
+from ...model_id import OpenAIModelId, model_name
+
+OPENAI_FINISH_REASON_MAP = {
+    "length": FinishReason.MAX_TOKENS,
+    "content_filter": FinishReason.REFUSAL,
+}
+
+
+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: 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
+
+    parts: list[AssistantContentPart] = []
+    if message.content:
+        parts.append(Text(text=message.content))
+    if message.refusal:
+        parts.append(Text(text=message.refusal))
+        refused = True
+    if message.tool_calls:
+        for tool_call in message.tool_calls:
+            if tool_call.type == "custom":
+                # This should never happen, because we never create "custom" tools
+                # https://platform.openai.com/docs/guides/function-calling#custom-tools
+                raise NotImplementedError("OpenAI custom tools are not supported.")
+            parts.append(
+                ToolCall(
+                    id=tool_call.id,
+                    name=tool_call.function.name,
+                    args=tool_call.function.arguments,
+                )
+            )
+
+    finish_reason = (
+        FinishReason.REFUSAL
+        if refused
+        else OPENAI_FINISH_REASON_MAP.get(choice.finish_reason)
+    )
+
+    assistant_message = AssistantMessage(
+        content=parts,
+        provider_id=provider_id,
+        model_id=model_id,
+        provider_model_name=provider_model_name or model_name(model_id, "completions"),
+        raw_message=message.model_dump(exclude_none=True),
+    )
+
+    usage = _decode_usage(response.usage)
+    return assistant_message, finish_reason, usage
+
+
+class _OpenAIChunkProcessor:
+    """Processes OpenAI chat completion chunks and maintains state across chunks."""
+
+    def __init__(self) -> None:
+        self.current_content_type: Literal["text", "tool_call"] | None = None
+        self.current_tool_index: int | None = None
+        self.current_tool_id: str | None = None
+        self.refusal_encountered = False
+
+    def process_chunk(self, chunk: openai_types.ChatCompletionChunk) -> ChunkIterator:
+        """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
+
+        delta = choice.delta
+
+        content = delta.content or delta.refusal
+        if delta.refusal:
+            self.refusal_encountered = True
+        if content is not None:
+            if self.current_content_type is None:
+                yield TextStartChunk()
+                self.current_content_type = "text"
+            yield TextChunk(delta=content)
+
+        if delta.tool_calls:
+            if self.current_content_type == "text":
+                # In testing, I can't get OpenAI to emit text and tool calls in the same chunk
+                # But we handle this defensively.
+                yield TextEndChunk()  # pragma: no cover
+            self.current_content_type = "tool_call"
+
+            for tool_call_delta in delta.tool_calls:
+                index = tool_call_delta.index
+
+                if (
+                    self.current_tool_index is not None
+                    and self.current_tool_index > index
+                ):
+                    raise RuntimeError(
+                        f"Received tool data for already-finished tool at index {index}"
+                    )  # pragma: no cover
+
+                if (
+                    self.current_tool_index is not None
+                    and self.current_tool_index < index
+                ):
+                    if self.current_tool_id is None:  # pragma: no cover
+                        raise RuntimeError("No current_tool_id for ToolCallChunk")
+                    yield ToolCallEndChunk(id=self.current_tool_id)
+                    self.current_tool_index = None
+
+                if self.current_tool_index is None:
+                    if not tool_call_delta.function or not (
+                        name := tool_call_delta.function.name
+                    ):
+                        raise RuntimeError(
+                            f"Missing name for tool call at index {index}"
+                        )  # pragma: no cover
+
+                    self.current_tool_index = index
+                    if not (tool_id := tool_call_delta.id):
+                        raise RuntimeError(
+                            f"Missing id for tool call at index {index}"
+                        )  # pragma: no cover
+
+                    yield ToolCallStartChunk(
+                        id=tool_id,
+                        name=name,
+                    )
+                    self.current_tool_id = tool_id
+
+                if tool_call_delta.function and tool_call_delta.function.arguments:
+                    if self.current_tool_id is None:  # pragma: no cover
+                        raise RuntimeError("No current_tool_id for ToolCallChunk")
+                    yield ToolCallChunk(
+                        id=self.current_tool_id,
+                        delta=tool_call_delta.function.arguments,
+                    )
+
+        if choice.finish_reason:
+            if self.current_content_type == "text":
+                yield TextEndChunk()
+            elif self.current_content_type == "tool_call":
+                if self.current_tool_id is None:  # pragma: no cover
+                    raise RuntimeError("No current_tool_id for ToolCallChunk")
+                yield ToolCallEndChunk(id=self.current_tool_id)
+            elif self.current_content_type is not None:  # pragma: no cover
+                raise NotImplementedError()
+
+            finish_reason = (
+                FinishReason.REFUSAL
+                if self.refusal_encountered
+                else OPENAI_FINISH_REASON_MAP.get(choice.finish_reason)
+            )
+            if finish_reason is not None:
+                yield FinishReasonChunk(finish_reason=finish_reason)
+
+
+def decode_stream(
+    openai_stream: Stream[openai_types.ChatCompletionChunk],
+) -> ChunkIterator:
+    """Returns a ChunkIterator converted from an OpenAI Stream[ChatCompletionChunk]"""
+    processor = _OpenAIChunkProcessor()
+    for chunk in openai_stream:
+        yield from processor.process_chunk(chunk)
+
+
+async def decode_async_stream(
+    openai_stream: AsyncStream[openai_types.ChatCompletionChunk],
+) -> AsyncChunkIterator:
+    """Returns an AsyncChunkIterator converted from an OpenAI AsyncStream[ChatCompletionChunk]"""
+    processor = _OpenAIChunkProcessor()
+    async for chunk in openai_stream:
+        for item in processor.process_chunk(chunk):
+            yield item