mirascope 1.18.3__py3-none-any.whl → 1.19.0__py3-none-any.whl

mirascope/core/xai/_utils/_calculate_cost.py

@@ -0,0 +1,104 @@
+"""Calculate the cost of a Grok API call."""
+
+
+def calculate_cost(
+    input_tokens: int | float | None,
+    cached_tokens: int | float | None,
+    output_tokens: int | float | None,
+    model: str,
+) -> float | None:
+    """Calculate the cost of an xAI Grok API call.
+
+    https://docs.x.ai
+
+    Pricing (per 1M tokens):
+
+    Model                     Input              Cached             Output
+    grok-3                    $3.50              $0.875             $10.50
+    grok-3-latest             $3.50              $0.875             $10.50
+    grok-2                    $2.00              $0.50              $6.00
+    grok-2-latest             $2.00              $0.50              $6.00
+    grok-2-1212               $2.00              $0.50              $6.00
+    grok-2-mini               $0.33              $0.083             $1.00
+    grok-2-vision-1212        $2.00              $0.50              $6.00
+    grok-vision-beta          $5.00              $1.25              $15.00
+    grok-beta                 $5.00              $1.25              $15.00
+
+    Args:
+        input_tokens: Number of input tokens
+        cached_tokens: Number of cached tokens
+        output_tokens: Number of output tokens
+        model: Model name to use for pricing calculation
+
+    Returns:
+        Total cost in USD or None if invalid input
+    """
+    pricing = {
+        "grok-3": {
+            "prompt": 0.000_003_5,
+            "cached": 0.000_000_875,
+            "completion": 0.000_010_5,
+        },
+        "grok-3-latest": {
+            "prompt": 0.000_003_5,
+            "cached": 0.000_000_875,
+            "completion": 0.000_010_5,
+        },
+        "grok-2": {
+            "prompt": 0.000_002,
+            "cached": 0.000_000_5,
+            "completion": 0.000_006,
+        },
+        "grok-latest": {
+            "prompt": 0.000_002,
+            "cached": 0.000_000_5,
+            "completion": 0.000_006,
+        },
+        "grok-2-1212": {
+            "prompt": 0.000_002,
+            "cached": 0.000_000_5,
+            "completion": 0.000_006,
+        },
+        "grok-2-mini": {
+            "prompt": 0.000_000_33,
+            "cached": 0.000_000_083,
+            "completion": 0.000_001,
+        },
+        "grok-2-vision-1212": {
+            "prompt": 0.000_002,
+            "cached": 0.000_000_5,
+            "completion": 0.000_006,
+        },
+        "grok-vision-beta": {
+            "prompt": 0.000_005,
+            "cached": 0.000_001_25,
+            "completion": 0.000_015,
+        },
+        "grok-beta": {
+            "prompt": 0.000_005,
+            "cached": 0.000_001_25,
+            "completion": 0.000_015,
+        },
+    }
+
+    if input_tokens is None or output_tokens is None:
+        return None
+
+    if cached_tokens is None:
+        cached_tokens = 0
+
+    try:
+        model_pricing = pricing[model]
+    except KeyError:
+        return None
+
+    prompt_price = model_pricing["prompt"]
+    cached_price = model_pricing["cached"]
+    completion_price = model_pricing["completion"]
+
+    prompt_cost = input_tokens * prompt_price
+    cached_cost = cached_tokens * cached_price
+    completion_cost = output_tokens * completion_price
+    total_cost = prompt_cost + cached_cost + completion_cost
+
+    return total_cost

mirascope/core/xai/_utils/_setup_call.py

@@ -0,0 +1,113 @@
+"""This module contains the setup_call function for OpenAI tools."""
+
+import os
+from collections.abc import Awaitable, Callable
+from typing import Any, overload
+
+from openai import AsyncOpenAI, OpenAI
+from openai.types.chat import (
+    ChatCompletion,
+    ChatCompletionChunk,
+    ChatCompletionMessageParam,
+)
+from pydantic import BaseModel
+
+from ...base import BaseTool
+from ...base._utils import AsyncCreateFn, CreateFn, fn_is_async
+from ...base.call_params import CommonCallParams
+from ...base.stream_config import StreamConfig
+from ...openai import (
+    AsyncOpenAIDynamicConfig,
+    OpenAICallParams,
+    OpenAIDynamicConfig,
+    OpenAITool,
+)
+from ...openai._call_kwargs import OpenAICallKwargs
+from ...openai._utils import setup_call as setup_call_openai
+
+
+@overload
+def setup_call(
+    *,
+    model: str,
+    client: AsyncOpenAI | None,
+    fn: Callable[..., Awaitable[AsyncOpenAIDynamicConfig]],
+    fn_args: dict[str, Any],
+    dynamic_config: AsyncOpenAIDynamicConfig,
+    tools: list[type[BaseTool] | Callable] | None,
+    json_mode: bool,
+    call_params: OpenAICallParams | CommonCallParams,
+    response_model: type[BaseModel] | None,
+    stream: bool | StreamConfig,
+) -> tuple[
+    AsyncCreateFn[ChatCompletion, ChatCompletionChunk],
+    str | None,
+    list[ChatCompletionMessageParam],
+    list[type[OpenAITool]] | None,
+    OpenAICallKwargs,
+]: ...
+
+
+@overload
+def setup_call(
+    *,
+    model: str,
+    client: OpenAI | None,
+    fn: Callable[..., OpenAIDynamicConfig],
+    fn_args: dict[str, Any],
+    dynamic_config: OpenAIDynamicConfig,
+    tools: list[type[BaseTool] | Callable] | None,
+    json_mode: bool,
+    call_params: OpenAICallParams | CommonCallParams,
+    response_model: type[BaseModel] | None,
+    stream: bool | StreamConfig,
+) -> tuple[
+    CreateFn[ChatCompletion, ChatCompletionChunk],
+    str | None,
+    list[ChatCompletionMessageParam],
+    list[type[OpenAITool]] | None,
+    OpenAICallKwargs,
+]: ...
+
+
+def setup_call(
+    *,
+    model: str,
+    client: OpenAI | AsyncOpenAI | None,
+    fn: Callable[..., OpenAIDynamicConfig]
+    | Callable[..., Awaitable[AsyncOpenAIDynamicConfig]],
+    fn_args: dict[str, Any],
+    dynamic_config: OpenAIDynamicConfig | AsyncOpenAIDynamicConfig,
+    tools: list[type[BaseTool] | Callable] | None,
+    json_mode: bool,
+    call_params: OpenAICallParams | CommonCallParams,
+    response_model: type[BaseModel] | None,
+    stream: bool | StreamConfig,
+) -> tuple[
+    CreateFn[ChatCompletion, ChatCompletionChunk]
+    | AsyncCreateFn[ChatCompletion, ChatCompletionChunk],
+    str | None,
+    list[ChatCompletionMessageParam],
+    list[type[OpenAITool]] | None,
+    OpenAICallKwargs,
+]:
+    if not client:
+        api_key = os.environ.get("XAI_API_KEY")
+        client = (
+            AsyncOpenAI(api_key=api_key, base_url="https://api.x.ai/v1")
+            if fn_is_async(fn)
+            else OpenAI(api_key=api_key, base_url="https://api.x.ai/v1")
+        )
+    create, prompt_template, messages, tool_types, call_kwargs = setup_call_openai(
+        model=model,  # pyright: ignore [reportCallIssue]
+        client=client,  # pyright: ignore [reportArgumentType]
+        fn=fn,  # pyright: ignore [reportArgumentType]
+        fn_args=fn_args,  # pyright: ignore [reportArgumentType]
+        dynamic_config=dynamic_config,
+        tools=tools,
+        json_mode=json_mode,
+        call_params=call_params,
+        response_model=response_model,
+        stream=stream,
+    )
+    return create, prompt_template, messages, tool_types, call_kwargs

mirascope/core/xai/call_params.py

@@ -0,0 +1,10 @@
+"""usage docs: learn/calls.md#provider-specific-parameters"""
+
+from ..openai import OpenAICallParams
+
+
+class XAICallParams(OpenAICallParams):
+    """A simple wrapper around `OpenAICallParams.`
+
+    Since xAI supports the OpenAI spec, we change nothing here.
+    """

mirascope/core/xai/call_response.py

@@ -0,0 +1,27 @@
+"""This module contains the `XAICallResponse` class.
+
+usage docs: learn/calls.md#handling-responses
+"""
+
+from pydantic import computed_field
+
+from ..openai import OpenAICallResponse
+from ._utils import calculate_cost
+
+
+class XAICallResponse(OpenAICallResponse):
+    """A simpler wrapper around `OpenAICallResponse`.
+
+    Everything is the same except the `cost` property, which has been updated to use
+    xAI's cost calculations so that cost tracking works for non-OpenAI models.
+    """
+
+    _provider = "xai"
+
+    @computed_field
+    @property
+    def cost(self) -> float | None:
+        """Returns the cost of the call."""
+        return calculate_cost(
+            self.input_tokens, self.cached_tokens, self.output_tokens, self.model
+        )

mirascope/core/xai/call_response_chunk.py

@@ -0,0 +1,14 @@
+"""This module contains the `XAICallResponseChunk` class.
+
+usage docs: learn/streams.md#handling-streamed-responses
+"""
+
+from ..openai import OpenAICallResponseChunk
+
+
+class XAICallResponseChunk(OpenAICallResponseChunk):
+    """A simpler wrapper around `OpenAICallResponse`.
+
+    Everything is the same except the `cost` property, which has been updated to use
+    xAI's cost calculations so that cost tracking works for non-OpenAI models.
+    """

mirascope/core/xai/dynamic_config.py

@@ -0,0 +1,8 @@
+"""This module defines the function return type for functions as LLM calls."""
+
+from typing import TypeAlias
+
+from ..openai import AsyncOpenAIDynamicConfig, OpenAIDynamicConfig
+
+AsyncXAIDynamicConfig: TypeAlias = AsyncOpenAIDynamicConfig
+XAIDynamicConfig: TypeAlias = OpenAIDynamicConfig

mirascope/core/xai/stream.py

@@ -0,0 +1,57 @@
+"""The `XAIStream` class for convenience around streaming xAI LLM calls.
+
+usage docs: learn/streams.md
+"""
+
+from collections.abc import AsyncGenerator, Generator
+
+from ..openai import OpenAIStream
+from .call_response import XAICallResponse
+from .call_response_chunk import XAICallResponseChunk
+from .tool import XAITool
+
+
+class XAIStream(OpenAIStream):
+    """A simple wrapper around `OpenAIStream`.
+
+    Everything is the same except updates to the `construct_call_response` method and
+    the `cost` property so that cost is properly calculated using xAI's cost
+    calculation method. This ensures cost calculation works for non-OpenAI models.
+    """
+
+    _provider = "xai"
+
+    def __iter__(
+        self,
+    ) -> Generator[tuple[XAICallResponseChunk, XAITool | None], None, None]:
+        yield from super().__iter__()  # pyright: ignore [reportReturnType]
+
+    def __aiter__(
+        self,
+    ) -> AsyncGenerator[tuple[XAICallResponseChunk, XAITool | None], None]:
+        return super().__aiter__()  # pyright: ignore [reportReturnType] # pragma: no cover
+
+    @property
+    def cost(self) -> float | None:
+        """Returns the cost of the call."""
+        response = self.construct_call_response()
+        return response.cost
+
+    def construct_call_response(self) -> XAICallResponse:
+        openai_call_response = super().construct_call_response()
+        response = XAICallResponse(
+            metadata=openai_call_response.metadata,
+            response=openai_call_response.response,
+            tool_types=openai_call_response.tool_types,
+            prompt_template=openai_call_response.prompt_template,
+            fn_args=openai_call_response.fn_args,
+            dynamic_config=openai_call_response.dynamic_config,
+            messages=openai_call_response.messages,
+            call_params=openai_call_response.call_params,
+            call_kwargs=openai_call_response.call_kwargs,
+            user_message_param=openai_call_response.user_message_param,
+            start_time=openai_call_response.start_time,
+            end_time=openai_call_response.end_time,
+        )
+        response._model = self.model
+        return response

mirascope/core/xai/tool.py

@@ -0,0 +1,13 @@
+"""The `XAITool` class for easy tool usage with xAI LLM calls.
+
+usage docs: learn/tools.md
+"""
+
+from ..openai import OpenAITool
+
+
+class XAITool(OpenAITool):
+    """A simple wrapper around `OpenAITool`.
+
+    Since xAI supports the OpenAI spec, we change nothing here.
+    """

mirascope/integrations/_middleware_factory.py

@@ -175,9 +175,9 @@ def middleware_factory(
                     def new_stream_aiter(
                         self: Any,  # noqa: ANN401
                     ) -> AsyncGenerator[tuple[Any, Any | None], Any]:  # noqa: ANN401
-                        async def generator() -> (
-                            AsyncGenerator[tuple[Any, Any | None], Any]
-                        ):
+                        async def generator() -> AsyncGenerator[
+                            tuple[Any, Any | None], Any
+                        ]:
                             try:
                                 async for chunk, tool in original_aiter():
                                     yield chunk, tool
@@ -226,9 +226,9 @@ def middleware_factory(
                     def new_aiter(
                         self: Any,  # noqa: ANN401
                     ) -> AsyncGenerator[tuple[Any, Any | None], Any]:  # noqa: ANN401
-                        async def generator() -> (
-                            AsyncGenerator[tuple[Any, Any | None], Any]
-                        ):
+                        async def generator() -> AsyncGenerator[
+                            tuple[Any, Any | None], Any
+                        ]:
                             try:
                                 async for chunk in original_aiter():
                                     yield chunk

mirascope/integrations/logfire/_utils.py

@@ -23,7 +23,7 @@ def custom_context_manager(
 ) -> Generator[logfire.LogfireSpan, Any, None]:
     metadata: Metadata = _utils.get_metadata(fn, None)
     tags = metadata.get("tags", [])
-    with logfire.with_settings(custom_scope_suffix="mirascope", tags=list(tags)).span(
+    with logfire.with_settings(custom_scope_suffix="mirascope", tags=list(tags)).span(  # pyright: ignore[reportGeneralTypeIssues]
         fn.__name__
     ) as logfire_span:
         yield logfire_span

mirascope/llm/__init__.py

@@ -1,6 +1,6 @@
-from ._protocols import Provider
+from ._protocols import LocalProvider, Provider
 from .call_response import CallResponse
 from .llm_call import call
 from .llm_override import override
 
-__all__ = ["call", "override", "CallResponse", "Provider"]
+__all__ = ["CallResponse", "LocalProvider", "Provider", "call", "override"]

mirascope/llm/_protocols.py

@@ -73,6 +73,12 @@ Provider: TypeAlias = Literal[
     "mistral",
     "openai",
     "vertex",
+    "xai",
+]
+
+LocalProvider: TypeAlias = Literal[
+    "ollama",
+    "vllm",
 ]
 
 
@@ -90,9 +96,9 @@ class _CallDecorator(
     ],
 ):
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -110,9 +116,9 @@ class _CallDecorator(
     ]: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -127,7 +133,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -140,9 +146,9 @@ class _CallDecorator(
     ) -> SyncLLMFunctionDecorator[_BaseDynamicConfigT, _BaseCallResponseT]: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -159,7 +165,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -174,7 +180,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -189,7 +195,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -207,7 +213,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig = True,
@@ -224,7 +230,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig = True,
@@ -239,7 +245,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig = True,
@@ -254,7 +260,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig = True,
@@ -272,7 +278,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig = True,
@@ -288,9 +294,9 @@ class _CallDecorator(
     ) -> NoReturn: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -307,7 +313,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -322,7 +328,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -337,7 +343,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -354,7 +360,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -369,7 +375,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[False] = False,
@@ -382,9 +388,9 @@ class _CallDecorator(
     ) -> SyncLLMFunctionDecorator[_BaseDynamicConfigT, _ParsedOutputT]: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig,
@@ -404,7 +410,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig,
@@ -421,7 +427,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig,
@@ -436,7 +442,7 @@ class _CallDecorator(
     @overload
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: Literal[True] | StreamConfig,
@@ -456,7 +462,7 @@ class _CallDecorator(
 
     def __call__(
         self,
-        provider: Provider,
+        provider: Provider | LocalProvider,
         model: str,
         *,
         stream: bool | StreamConfig = False,