mirascope 2.0.0a6__py3-none-any.whl → 2.0.1__py3-none-any.whl

mirascope/llm/providers/google/provider.py

@@ -1,13 +1,16 @@
 """Google provider implementation."""
 
+from __future__ import annotations
+
 from collections.abc import Sequence
+from typing import TYPE_CHECKING
 from typing_extensions import Unpack
 
 from google.genai import Client
 from google.genai.types import HttpOptions
 
 from ...context import Context, DepsT
-from ...formatting import Format, FormattableT
+from ...formatting import Format, FormattableT, OutputParser
 from ...messages import Message
 from ...responses import (
     AsyncContextResponse,
@@ -29,10 +32,13 @@ from ...tools import (
     Tool,
     Toolkit,
 )
-from ..base import BaseProvider, Params
+from ..base import BaseProvider
 from . import _utils
 from .model_id import GoogleModelId, model_name
 
+if TYPE_CHECKING:
+    from ...models import Params
+
 
 class GoogleProvider(BaseProvider[Client]):
     """The client for the Google LLM model."""
@@ -64,7 +70,10 @@ class GoogleProvider(BaseProvider[Client]):
         model_id: GoogleModelId,
         messages: Sequence[Message],
         tools: Sequence[Tool] | Toolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> Response | Response[FormattableT]:
         """Generate an `llm.Response` by synchronously calling the Google GenAI API.
@@ -88,8 +97,9 @@ class GoogleProvider(BaseProvider[Client]):
         )
         google_response = self.client.models.generate_content(**kwargs)
 
+        include_thoughts = _utils.get_include_thoughts(params)
         assistant_message, finish_reason, usage = _utils.decode_response(
-            google_response, model_id
+            google_response, model_id, include_thoughts=include_thoughts
         )
 
         return Response(
@@ -115,7 +125,10 @@ class GoogleProvider(BaseProvider[Client]):
         tools: Sequence[Tool | ContextTool[DepsT]]
         | ContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
         """Generate an `llm.ContextResponse` by synchronously calling the Google GenAI API.
@@ -140,8 +153,9 @@ class GoogleProvider(BaseProvider[Client]):
         )
         google_response = self.client.models.generate_content(**kwargs)
 
+        include_thoughts = _utils.get_include_thoughts(params)
         assistant_message, finish_reason, usage = _utils.decode_response(
-            google_response, model_id
+            google_response, model_id, include_thoughts=include_thoughts
         )
 
         return ContextResponse(
@@ -164,7 +178,10 @@ class GoogleProvider(BaseProvider[Client]):
         model_id: GoogleModelId,
         messages: Sequence[Message],
         tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> AsyncResponse | AsyncResponse[FormattableT]:
         """Generate an `llm.AsyncResponse` by asynchronously calling the Google GenAI API.
@@ -188,8 +205,9 @@ class GoogleProvider(BaseProvider[Client]):
         )
         google_response = await self.client.aio.models.generate_content(**kwargs)
 
+        include_thoughts = _utils.get_include_thoughts(params)
         assistant_message, finish_reason, usage = _utils.decode_response(
-            google_response, model_id
+            google_response, model_id, include_thoughts=include_thoughts
         )
 
         return AsyncResponse(
@@ -215,7 +233,10 @@ class GoogleProvider(BaseProvider[Client]):
         tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
         | AsyncContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
         """Generate an `llm.AsyncContextResponse` by asynchronously calling the Google GenAI API.
@@ -240,8 +261,9 @@ class GoogleProvider(BaseProvider[Client]):
         )
         google_response = await self.client.aio.models.generate_content(**kwargs)
 
+        include_thoughts = _utils.get_include_thoughts(params)
         assistant_message, finish_reason, usage = _utils.decode_response(
-            google_response, model_id
+            google_response, model_id, include_thoughts=include_thoughts
         )
 
         return AsyncContextResponse(
@@ -264,7 +286,10 @@ class GoogleProvider(BaseProvider[Client]):
         model_id: GoogleModelId,
         messages: Sequence[Message],
         tools: Sequence[Tool] | Toolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> StreamResponse | StreamResponse[FormattableT]:
         """Generate an `llm.StreamResponse` by synchronously streaming from the Google GenAI API.
@@ -289,7 +314,10 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_stream = self.client.models.generate_content_stream(**kwargs)
 
-        chunk_iterator = _utils.decode_stream(google_stream)
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
 
         return StreamResponse(
             provider_id="google",
@@ -311,7 +339,10 @@ class GoogleProvider(BaseProvider[Client]):
         tools: Sequence[Tool | ContextTool[DepsT]]
         | ContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> ContextStreamResponse[DepsT] | ContextStreamResponse[DepsT, FormattableT]:
         """Generate an `llm.ContextStreamResponse` by synchronously streaming from the Google GenAI API.
@@ -337,7 +368,10 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_stream = self.client.models.generate_content_stream(**kwargs)
 
-        chunk_iterator = _utils.decode_stream(google_stream)
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
 
         return ContextStreamResponse(
             provider_id="google",
@@ -356,7 +390,10 @@ class GoogleProvider(BaseProvider[Client]):
         model_id: GoogleModelId,
         messages: Sequence[Message],
         tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> AsyncStreamResponse | AsyncStreamResponse[FormattableT]:
         """Generate an `llm.AsyncStreamResponse` by asynchronously streaming from the Google GenAI API.
@@ -381,7 +418,10 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_stream = await self.client.aio.models.generate_content_stream(**kwargs)
 
-        chunk_iterator = _utils.decode_async_stream(google_stream)
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_async_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
 
         return AsyncStreamResponse(
             provider_id="google",
@@ -403,7 +443,10 @@ class GoogleProvider(BaseProvider[Client]):
         tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
         | AsyncContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> (
         AsyncContextStreamResponse[DepsT]
@@ -432,7 +475,10 @@ class GoogleProvider(BaseProvider[Client]):
 
         google_stream = await self.client.aio.models.generate_content_stream(**kwargs)
 
-        chunk_iterator = _utils.decode_async_stream(google_stream)
+        include_thoughts = _utils.get_include_thoughts(params)
+        chunk_iterator = _utils.decode_async_stream(
+            google_stream, include_thoughts=include_thoughts
+        )
 
         return AsyncContextStreamResponse(
             provider_id="google",

mirascope/llm/providers/mirascope/_utils.py

@@ -7,15 +7,15 @@ from ..base import Provider
 from ..provider_id import ProviderId
 
 
-def extract_provider_prefix(model_id: str) -> str | None:
-    """Extract provider prefix from model ID.
+def extract_model_scope(model_id: str) -> str | None:
+    """Extract model scope from model ID.
 
     Args:
-        model_id: Model identifier in the format "provider/model-name"
+        model_id: Model identifier in the format "scope/model-name"
                   e.g., "openai/gpt-4", "anthropic/claude-3", "google/gemini-pro"
 
     Returns:
-        The provider prefix (e.g., "openai", "anthropic", "google") or None if invalid format.
+        The model scope (e.g., "openai", "anthropic", "google") or None if invalid format.
     """
     if "/" not in model_id:
         return None
@@ -29,12 +29,12 @@ def get_default_router_base_url() -> str:
         The router base URL (without trailing provider path).
     """
     return os.environ.get(
-        "MIRASCOPE_ROUTER_BASE_URL", "https://mirascope.com/router/v0"
+        "MIRASCOPE_ROUTER_BASE_URL", "https://mirascope.com/router/v2"
     )
 
 
 def create_underlying_provider(
-    provider_prefix: str, api_key: str, router_base_url: str
+    model_scope: str, api_key: str, router_base_url: str
 ) -> Provider:
     """Create and cache an underlying provider instance using provider_singleton.
 
@@ -42,10 +42,9 @@ def create_underlying_provider(
     delegates to provider_singleton for caching and instantiation.
 
     Args:
-        provider_prefix: The provider name (e.g., "openai", "anthropic", "google",
-                         "openai:completions", "openai:responses")
+        model_scope: The model scope (e.g., "openai", "anthropic", "google")
         api_key: The API key to use for authentication
-        router_base_url: The base router URL (e.g., "http://mirascope.com/router/v0")
+        router_base_url: The base router URL (e.g., "http://mirascope.com/router/v2")
 
     Returns:
         A cached provider instance configured for the Mirascope Router.
@@ -53,17 +52,14 @@ def create_underlying_provider(
     Raises:
         ValueError: If the provider is unsupported.
     """
-    # Extract base provider name (handles variants like "openai:completions")
-    base_provider = provider_prefix.split(":")[0]
-
-    if base_provider not in ["anthropic", "google", "openai"]:
+    if model_scope not in ["anthropic", "google", "openai"]:
         raise ValueError(
-            f"Unsupported provider: {provider_prefix}. "
+            f"Unsupported provider: {model_scope}. "
             f"Mirascope Router currently supports: anthropic, google, openai"
         )
 
-    base_url = f"{router_base_url}/{base_provider}"
-    if base_provider == "openai":  # OpenAI expects /v1, which their SDK doesn't add
+    base_url = f"{router_base_url}/{model_scope}"
+    if model_scope == "openai":  # OpenAI expects /v1, which their SDK doesn't add
         base_url = f"{base_url}/v1"
 
     # Lazy import to avoid circular dependencies
@@ -71,7 +67,7 @@ def create_underlying_provider(
 
     # Use provider_singleton which provides caching
     return provider_singleton(
-        cast(ProviderId, provider_prefix),
+        cast(ProviderId, model_scope),
         api_key=api_key,
         base_url=base_url,
     )

mirascope/llm/providers/mirascope/provider.py

@@ -1,11 +1,14 @@
 """Mirascope Router provider that routes requests through the Mirascope Router API."""
 
+from __future__ import annotations
+
 import os
 from collections.abc import Sequence
+from typing import TYPE_CHECKING
 from typing_extensions import Unpack
 
 from ...context import Context, DepsT
-from ...formatting import Format, FormattableT
+from ...formatting import Format, FormattableT, OutputParser
 from ...messages import Message
 from ...responses import (
     AsyncContextResponse,
@@ -27,9 +30,12 @@ from ...tools import (
     Tool,
     Toolkit,
 )
-from ..base import BaseProvider, Params, Provider
+from ..base import BaseProvider, Provider
 from . import _utils
 
+if TYPE_CHECKING:
+    from ...models import Params
+
 
 class MirascopeProvider(BaseProvider[None]):
     """Provider that routes LLM requests through the Mirascope Router API.
@@ -38,14 +44,15 @@ class MirascopeProvider(BaseProvider[None]):
     (Anthropic, Google, OpenAI) with usage tracking and cost calculation.
 
     This provider:
-    - Takes model IDs in the format "provider/model-name" (e.g., "openai/gpt-4")
+    - Takes model IDs in the format "scope/model-name" (e.g., "openai/gpt-4")
     - Routes requests to the Mirascope Router endpoint
     - Delegates to the appropriate underlying provider (Anthropic, Google, or OpenAI)
     - Uses MIRASCOPE_API_KEY for authentication
 
     Environment Variables:
         MIRASCOPE_API_KEY: Required API key for Mirascope Router authentication
-        MIRASCOPE_ROUTER_BASE_URL: Optional base URL override (default: https://mirascope.com/router/v0)
+        MIRASCOPE_ROUTER_BASE_URL: Optional base URL override
+            (default: https://mirascope.com/router/v2)
 
     Example:
         ```python
@@ -90,7 +97,7 @@ class MirascopeProvider(BaseProvider[None]):
                 environment variable.
             base_url: Optional base URL override for the Mirascope Router. If not
                 provided, reads from MIRASCOPE_ROUTER_BASE_URL environment variable
-                or defaults to https://mirascope.com/router/v0
+                or defaults to https://mirascope.com/router/v2
         """
         api_key = api_key or os.environ.get("MIRASCOPE_API_KEY")
         if not api_key:
@@ -122,7 +129,7 @@ class MirascopeProvider(BaseProvider[None]):
         """Get the underlying provider for a model ID.
 
         Args:
-            model_id: Model identifier in format "provider/model-name"
+            model_id: Model identifier in format "scope/model-name"
 
         Returns:
             The appropriate cached provider instance (Anthropic, Google, or OpenAI)
@@ -130,16 +137,16 @@ class MirascopeProvider(BaseProvider[None]):
         Raises:
             ValueError: If the model ID format is invalid or provider is unsupported
         """
-        provider_prefix = _utils.extract_provider_prefix(model_id)
-        if not provider_prefix:
+        model_scope = _utils.extract_model_scope(model_id)
+        if not model_scope:
             raise ValueError(
                 f"Invalid model ID format: {model_id}. "
-                f"Expected format 'provider/model-name' (e.g., 'openai/gpt-4')"
+                f"Expected format 'scope/model-name' (e.g., 'openai/gpt-4')"
             )
 
         # Use the cached function to get/create the provider
         return _utils.create_underlying_provider(
-            provider_prefix=provider_prefix,
+            model_scope=model_scope,
             api_key=self.api_key,
             router_base_url=self.router_base_url,
         )
@@ -150,7 +157,10 @@ class MirascopeProvider(BaseProvider[None]):
         model_id: str,
         messages: Sequence[Message],
         tools: Sequence[Tool] | Toolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> Response | Response[FormattableT]:
         """Generate an `llm.Response` by calling through the Mirascope Router."""
@@ -172,7 +182,10 @@ class MirascopeProvider(BaseProvider[None]):
         tools: Sequence[Tool | ContextTool[DepsT]]
         | ContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
         """Generate an `llm.ContextResponse` by calling through the Mirascope Router."""
@@ -192,7 +205,10 @@ class MirascopeProvider(BaseProvider[None]):
         model_id: str,
         messages: Sequence[Message],
         tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> AsyncResponse | AsyncResponse[FormattableT]:
         """Generate an `llm.AsyncResponse` by calling through the Mirascope Router."""
@@ -214,7 +230,10 @@ class MirascopeProvider(BaseProvider[None]):
         tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
         | AsyncContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
         """Generate an `llm.AsyncContextResponse` by calling through the Mirascope Router."""
@@ -234,7 +253,10 @@ class MirascopeProvider(BaseProvider[None]):
         model_id: str,
         messages: Sequence[Message],
         tools: Sequence[Tool] | Toolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> StreamResponse | StreamResponse[FormattableT]:
         """Stream an `llm.StreamResponse` by calling through the Mirascope Router."""
@@ -256,7 +278,10 @@ class MirascopeProvider(BaseProvider[None]):
         tools: Sequence[Tool | ContextTool[DepsT]]
         | ContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> (
         ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]
@@ -278,7 +303,10 @@ class MirascopeProvider(BaseProvider[None]):
         model_id: str,
         messages: Sequence[Message],
         tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> AsyncStreamResponse | AsyncStreamResponse[FormattableT]:
         """Stream an `llm.AsyncStreamResponse` by calling through the Mirascope Router."""
@@ -300,7 +328,10 @@ class MirascopeProvider(BaseProvider[None]):
         tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
         | AsyncContextToolkit[DepsT]
         | None = None,
-        format: type[FormattableT] | Format[FormattableT] | None = None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None = None,
         **params: Unpack[Params],
     ) -> (
         AsyncContextStreamResponse[DepsT, None]

mirascope/llm/providers/mlx/_utils.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from collections.abc import Callable
-from typing import TypeAlias, TypedDict
+from typing import TYPE_CHECKING, TypeAlias, TypedDict
 
 import mlx.core as mx
 from huggingface_hub.errors import LocalEntryNotFoundError
@@ -8,7 +10,10 @@ from mlx_lm.sample_utils import make_sampler
 
 from ...exceptions import NotFoundError
 from ...responses import FinishReason, Usage
-from ..base import Params, ProviderErrorMap, _utils as _base_utils
+from ..base import ProviderErrorMap, _utils as _base_utils
+
+if TYPE_CHECKING:
+    from ...models import Params
 
 Sampler: TypeAlias = Callable[[mx.array], mx.array]
 

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

@@ -6,7 +6,7 @@ from typing import TypeAlias
 
 from mlx_lm.generate import GenerationResponse
 
-from ....formatting import Format, FormattableT
+from ....formatting import Format, FormattableT, OutputParser
 from ....messages import AssistantContent, Message
 from ....responses import ChunkIterator
 from ....tools import AnyToolSchema, BaseToolkit
@@ -22,7 +22,10 @@ class BaseEncoder(abc.ABC):
         self,
         messages: Sequence[Message],
         tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT] | Format[FormattableT] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
     ) -> tuple[Sequence[Message], Format[FormattableT] | None, TokenIds]:
         """Encode the request messages into a format suitable for the model.
 

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

@@ -8,7 +8,7 @@ from mlx_lm.generate import GenerationResponse
 from transformers import PreTrainedTokenizer
 
 from ....content import ContentPart, TextChunk, TextEndChunk, TextStartChunk
-from ....formatting import Format, FormattableT
+from ....formatting import Format, FormattableT, OutputParser
 from ....messages import AssistantContent, Message
 from ....responses import (
     ChunkIterator,
@@ -81,7 +81,10 @@ class TransformersEncoder(BaseEncoder):
         self,
         messages: Sequence[Message],
         tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT] | Format[FormattableT] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
     ) -> tuple[Sequence[Message], Format[FormattableT] | None, TokenIds]:
         """Encode a request into a format suitable for the model."""
         tool_schemas = tools.tools if isinstance(tools, BaseToolkit) else tools or []

mirascope/llm/providers/mlx/mlx.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import asyncio
 import threading
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 from typing_extensions import Unpack
 
 import mlx.core as mx
@@ -10,15 +13,17 @@ from mlx_lm import stream_generate  # type: ignore[reportPrivateImportUsage]
 from mlx_lm.generate import GenerationResponse
 from transformers import PreTrainedTokenizer
 
-from ...formatting import Format, FormattableT
+from ...formatting import Format, FormattableT, OutputParser
 from ...messages import AssistantMessage, Message, assistant
 from ...responses import AsyncChunkIterator, ChunkIterator, StreamResponseChunk
 from ...tools import AnyToolSchema, BaseToolkit
-from ..base import Params
 from . import _utils
 from .encoding import BaseEncoder, TokenIds
 from .model_id import MLXModelId
 
+if TYPE_CHECKING:
+    from ...models import Params
+
 
 def _consume_sync_stream_into_queue(
     generation_stream: ChunkIterator,
@@ -133,7 +138,10 @@ class MLX:
         self,
         messages: Sequence[Message],
         tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT] | Format[FormattableT] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
         params: Params,
     ) -> tuple[Sequence[Message], Format[FormattableT] | None, ChunkIterator]:
         """Stream response chunks synchronously.
@@ -156,7 +164,10 @@ class MLX:
         self,
         messages: Sequence[Message],
         tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT] | Format[FormattableT] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
         params: Params,
     ) -> tuple[Sequence[Message], Format[FormattableT] | None, AsyncChunkIterator]:
         """Stream response chunks asynchronously.
@@ -180,7 +191,10 @@ class MLX:
         self,
         messages: Sequence[Message],
         tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT] | Format[FormattableT] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
         params: Params,
     ) -> tuple[
         Sequence[Message],
@@ -216,7 +230,10 @@ class MLX:
         self,
         messages: Sequence[Message],
         tools: Sequence[AnyToolSchema] | BaseToolkit[AnyToolSchema] | None,
-        format: type[FormattableT] | Format[FormattableT] | None,
+        format: type[FormattableT]
+        | Format[FormattableT]
+        | OutputParser[FormattableT]
+        | None,
         params: Params,
     ) -> tuple[
         Sequence[Message],