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

mirascope/ops/_internal/utils.py

@@ -0,0 +1,75 @@
+"""Internal utility functions for the Mirascope."""
+
+import inspect
+from collections.abc import Callable
+from typing import Any, TypeAlias
+
+import orjson
+
+from .protocols import P
+
+ORJSON_OPTS = (
+    orjson.OPT_NON_STR_KEYS
+    | orjson.OPT_NAIVE_UTC
+    | orjson.OPT_SERIALIZE_NUMPY
+    | orjson.OPT_SERIALIZE_DATACLASS
+    | orjson.OPT_SERIALIZE_UUID
+)
+
+PrimitiveType: TypeAlias = str | int | float | bool
+
+
+def json_dumps(obj: Any) -> str:  # noqa: ANN401
+    """Serialize Python objects to JSON using orjson."""
+    # json should be utf-8 encoded, json key should be str
+    return orjson.dumps(obj, option=ORJSON_OPTS).decode("utf-8")
+
+
+def get_qualified_name(fn: Callable[..., Any]) -> str:
+    """Return the simplified qualified name of a function.
+
+    If the function is defined locally, return the name after '<locals>.'; otherwise,
+    return the last non-empty part after splitting by '.'.
+    """
+    qualified_name = fn.__qualname__
+    if "<locals>." in qualified_name:
+        return qualified_name.split("<locals>.")[-1]
+    else:
+        parts = [part for part in qualified_name.split(".") if part]
+        return parts[-1] if parts else qualified_name
+
+
+def extract_arguments(
+    fn: Callable[P, Any],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> tuple[dict[str, str], dict[str, Any]]:
+    """Returns a tuple of (arg_types, arg_values) dictionaries from function call."""
+    signature = inspect.signature(fn)
+    bound_arguments = signature.bind(*args, **kwargs)
+    bound_arguments.apply_defaults()
+
+    arg_types: dict[str, str] = {}
+    arg_values: dict[str, Any] = {}
+
+    for param_name, param_value in bound_arguments.arguments.items():
+        parameter = signature.parameters[param_name]
+        if parameter.annotation != inspect.Parameter.empty:
+            if hasattr(parameter.annotation, "__name__"):
+                type_str = parameter.annotation.__name__
+            else:
+                type_str = str(parameter.annotation)
+            arg_types[param_name] = type_str
+        else:
+            arg_types[param_name] = type(param_value).__name__
+
+        if isinstance(param_value, PrimitiveType) or param_value is None:
+            arg_values[param_name] = param_value
+        else:
+            try:
+                json_dumps(param_value)
+                arg_values[param_name] = param_value
+            except (TypeError, ValueError):
+                arg_values[param_name] = repr(param_value)
+
+    return arg_types, arg_values

mirascope/ops/_internal/versioned_calls.py

@@ -0,0 +1,512 @@
+"""Versioned call wrappers for @ops.version decorated @llm.call functions."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Concatenate, Generic
+from typing_extensions import TypeIs
+
+from ...llm.calls import AsyncCall, AsyncContextCall, Call, ContextCall
+from ...llm.context import Context, DepsT
+from ...llm.formatting import FormattableT
+from ...llm.responses import (
+    AsyncContextResponse,
+    AsyncContextStreamResponse,
+    AsyncResponse,
+    AsyncStreamResponse,
+    ContextResponse,
+    ContextStreamResponse,
+    Response,
+    StreamResponse,
+)
+from ...llm.types import P
+from ..exceptions import ClosureComputationError
+from .closure import Closure
+from .protocols import (
+    AsyncFunction,
+    R,
+    SyncFunction,
+)
+from .versioned_functions import (
+    AsyncVersionedFunction,
+    AsyncVersionedResult,
+    VersionedFunction,
+    VersionedResult,
+    VersionInfo,
+    _BaseVersionedFunction,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _compute_closure_from_call(
+    call_obj: (
+        ContextCall[P, DepsT, FormattableT]
+        | AsyncContextCall[P, DepsT, FormattableT]
+        | Call[P, FormattableT]
+        | AsyncCall[P, FormattableT]
+    ),
+) -> Closure | None:
+    """Compute closure from a Call object's prompt function.
+
+    Extracts the original function from the Call's prompt closure and computes
+    a Closure object for version tracking. Returns None if closure computation
+    fails.
+    """
+    fn = call_obj.prompt.fn
+    closure = getattr(fn, "__closure__", None)
+    original_fn = closure[0].cell_contents if closure and len(closure) > 0 else fn
+    try:
+        return Closure.from_fn(original_fn)
+    except ClosureComputationError as e:
+        logger.warning(
+            "Failed to build closure for %s; continuing without version registration: %s",
+            e.qualified_name,
+            e,
+        )
+        return None
+
+
+def is_version_call_type(
+    fn: (
+        SyncFunction[P, R]
+        | AsyncFunction[P, R]
+        | ContextCall[P, DepsT, FormattableT]
+        | AsyncContextCall[P, DepsT, FormattableT]
+        | Call[P, FormattableT]
+        | AsyncCall[P, FormattableT]
+    ),
+) -> TypeIs[
+    ContextCall[P, DepsT, FormattableT]
+    | AsyncContextCall[P, DepsT, FormattableT]
+    | Call[P, FormattableT]
+    | AsyncCall[P, FormattableT]
+]:
+    """Check if fn is any of the Call types."""
+    return isinstance(fn, Call | AsyncCall | ContextCall | AsyncContextCall)
+
+
+def wrap_version_call(
+    fn: (
+        ContextCall[P, DepsT, FormattableT]
+        | AsyncContextCall[P, DepsT, FormattableT]
+        | Call[P, FormattableT]
+        | AsyncCall[P, FormattableT]
+    ),
+    tags: tuple[str, ...],
+    name: str | None,
+    metadata: dict[str, str],
+) -> (
+    VersionedContextCall[P, DepsT, FormattableT]
+    | VersionedAsyncContextCall[P, DepsT, FormattableT]
+    | VersionedCall[P, FormattableT]
+    | VersionedAsyncCall[P, FormattableT]
+):
+    """Wrap a Call object with the appropriate VersionedCall type."""
+    if isinstance(fn, AsyncContextCall):
+        return VersionedAsyncContextCall(
+            _call=fn, tags=tags, name=name, metadata=metadata
+        )
+    elif isinstance(fn, ContextCall):
+        return VersionedContextCall(_call=fn, tags=tags, name=name, metadata=metadata)
+    elif isinstance(fn, AsyncCall):
+        return VersionedAsyncCall(_call=fn, tags=tags, name=name, metadata=metadata)
+    else:
+        return VersionedCall(_call=fn, tags=tags, name=name, metadata=metadata)
+
+
+@dataclass(kw_only=True)
+class _BaseVersionedCall:
+    """Base class for versioned call wrappers.
+
+    Provides common fields and methods for all versioned call types,
+    reducing code duplication across VersionedCall, VersionedAsyncCall,
+    VersionedContextCall, and VersionedAsyncContextCall.
+    """
+
+    tags: tuple[str, ...]
+    """Tags to be associated with versioned calls."""
+
+    name: str | None = None
+    """Optional custom name for the versioned function."""
+
+    metadata: dict[str, str] = field(default_factory=dict)
+    """Arbitrary key-value pairs for additional metadata."""
+
+    closure: Closure | None = field(init=False, default=None)
+    """Cached closure for the original function."""
+
+    def _compute_closure(
+        self,
+        wrapped_call: (
+            ContextCall[P, DepsT, FormattableT]
+            | AsyncContextCall[P, DepsT, FormattableT]
+            | Call[P, FormattableT]
+            | AsyncCall[P, FormattableT]
+        ),
+        versioned_call: _BaseVersionedFunction[..., object],
+        versioned_stream: _BaseVersionedFunction[..., object],
+    ) -> None:
+        """Initialize closure from wrapped call and assign to versioned wrappers."""
+        self.closure = _compute_closure_from_call(wrapped_call)
+        if self.closure is not None:
+            versioned_call.closure = self.closure
+            versioned_stream.closure = self.closure
+
+    @property
+    def version_info(self) -> VersionInfo | None:
+        """Returns static version metadata for this versioned call."""
+        if self.closure is None:
+            return None
+        return VersionInfo(
+            uuid=None,
+            hash=self.closure.hash,
+            signature_hash=self.closure.signature_hash,
+            name=self.name or self.closure.name,
+            description=self.closure.docstring,
+            version=_BaseVersionedFunction._compute_version(self.closure.hash),
+            tags=self.tags,
+            metadata=self.metadata,
+        )
+
+
+@dataclass(kw_only=True)
+class VersionedCall(_BaseVersionedCall, Generic[P, FormattableT]):
+    """Versioned wrapper for synchronous Call objects.
+
+    When @ops.version decorates an @llm.call, it returns a VersionedCall that wraps
+    the call and stream methods with versioning capabilities.
+
+    Example:
+        ```python
+        @ops.version
+        @llm.call("gpt-4o-mini")
+        def recommend_book(genre: str):
+            return f"Recommend a {genre} book"
+
+        # Returns Response directly (but execution is versioned)
+        response = recommend_book("fantasy")
+        print(response.content)
+
+        # Same as __call__
+        response = recommend_book.call("fantasy")
+
+        # Streaming returns StreamResponse (versioned)
+        stream = recommend_book.stream("fantasy")
+        for chunk in stream:
+            print(chunk)
+
+        # Use wrapped to get VersionedResult[Response] with version info
+        result = recommend_book.wrapped("fantasy")
+        print(result.result.content)
+        print(result.function_uuid)
+
+        # Use wrapped_stream to get VersionedResult[StreamResponse]
+        result = recommend_book.wrapped_stream("fantasy")
+
+        # Access version info
+        info = recommend_book.version_info
+        print(info.version, info.hash)
+        ```
+    """
+
+    _call: Call[P, FormattableT]
+    """The wrapped Call object."""
+
+    call: VersionedFunction[P, Response | Response[FormattableT]] = field(init=False)
+    """VersionedFunction wrapping the call method."""
+
+    stream: VersionedFunction[P, StreamResponse | StreamResponse[FormattableT]] = field(
+        init=False
+    )
+    """VersionedFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize VersionedFunction wrappers for call and stream methods."""
+        self.call = VersionedFunction(
+            fn=self._call.call,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self.stream = VersionedFunction(
+            fn=self._call.stream,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self._compute_closure(self._call, self.call, self.stream)
+
+    def __call__(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> Response | Response[FormattableT]:
+        """Call the versioned function and return Response directly."""
+        return self.call(*args, **kwargs)
+
+    def wrapped(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[Response | Response[FormattableT]]:
+        """Call the versioned function and return a wrapped Response."""
+        return self.call.wrapped(*args, **kwargs)
+
+    def wrapped_stream(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[StreamResponse | StreamResponse[FormattableT]]:
+        """Stream the versioned function and return a wrapped StreamResponse."""
+        return self.stream.wrapped(*args, **kwargs)
+
+
+@dataclass(kw_only=True)
+class VersionedAsyncCall(_BaseVersionedCall, Generic[P, FormattableT]):
+    """Versioned wrapper for asynchronous AsyncCall objects.
+
+    Example:
+        ```python
+        @ops.version
+        @llm.call("gpt-4o-mini")
+        async def recommend_book(genre: str):
+            return f"Recommend a {genre} book"
+
+        # Returns AsyncResponse directly (but execution is versioned)
+        response = await recommend_book("fantasy")
+        print(response.content)
+
+        # Use wrapped to get AsyncVersionedResult[AsyncResponse] with version info
+        result = await recommend_book.wrapped("fantasy")
+        print(result.result.content)
+        print(result.function_uuid)
+        ```
+    """
+
+    _call: AsyncCall[P, FormattableT]
+    """The wrapped AsyncCall object."""
+
+    call: AsyncVersionedFunction[P, AsyncResponse | AsyncResponse[FormattableT]] = (
+        field(init=False)
+    )
+    """AsyncVersionedFunction wrapping the call method."""
+
+    stream: AsyncVersionedFunction[
+        P, AsyncStreamResponse | AsyncStreamResponse[FormattableT]
+    ] = field(init=False)
+    """AsyncVersionedFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize AsyncVersionedFunction wrappers for call and stream methods."""
+        self.call = AsyncVersionedFunction(
+            fn=self._call.call,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self.stream = AsyncVersionedFunction(
+            fn=self._call.stream,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self._compute_closure(self._call, self.call, self.stream)
+
+    async def __call__(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncResponse | AsyncResponse[FormattableT]:
+        """Call the versioned function and return AsyncResponse directly."""
+        return await self.call(*args, **kwargs)
+
+    async def wrapped(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[AsyncResponse | AsyncResponse[FormattableT]]:
+        """Call the versioned function and return a wrapped Response."""
+        return await self.call.wrapped(*args, **kwargs)
+
+    async def wrapped_stream(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[AsyncStreamResponse | AsyncStreamResponse[FormattableT]]:
+        """Stream the versioned function and return a wrapped StreamResponse."""
+        return await self.stream.wrapped(*args, **kwargs)
+
+
+@dataclass(kw_only=True)
+class VersionedContextCall(_BaseVersionedCall, Generic[P, DepsT, FormattableT]):
+    """Versioned wrapper for synchronous ContextCall objects.
+
+    Example:
+        ```python
+        @ops.version
+        @llm.call("gpt-4o-mini")
+        def recommend_book(ctx: llm.Context[str], genre: str):
+            return f"{ctx.deps} Recommend a {genre} book"
+
+        ctx = llm.Context(deps="As a librarian,")
+
+        # Returns ContextResponse directly (but execution is versioned)
+        response = recommend_book(ctx, "fantasy")
+        print(response.content)
+
+        # Use wrapped to get VersionedResult[ContextResponse] with version info
+        result = recommend_book.wrapped(ctx, "fantasy")
+        print(result.result.content)
+        ```
+    """
+
+    _call: ContextCall[P, DepsT, FormattableT]
+    """The wrapped ContextCall object."""
+
+    _call_versioned: VersionedFunction[
+        Concatenate[Context[DepsT], P],
+        ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """VersionedFunction wrapping the call method."""
+
+    _stream_versioned: VersionedFunction[
+        Concatenate[Context[DepsT], P],
+        ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """VersionedFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize VersionedFunction wrappers for call and stream methods."""
+        self._call_versioned = VersionedFunction(
+            fn=self._call.call,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self._stream_versioned = VersionedFunction(
+            fn=self._call.stream,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self._compute_closure(self._call, self._call_versioned, self._stream_versioned)
+
+    def call(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
+        """Call the versioned function and return ContextResponse directly."""
+        return self._call_versioned(ctx, *args, **kwargs)
+
+    def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
+        """Call the versioned function and return ContextResponse directly."""
+        return self.call(ctx, *args, **kwargs)
+
+    def stream(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> (
+        ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]
+    ):
+        """Stream the versioned function and return a ContextStreamResponse."""
+        return self._stream_versioned(ctx, *args, **kwargs)
+
+    def wrapped(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[
+        ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]
+    ]:
+        """Call the versioned function and return a wrapped Response."""
+        return self._call_versioned.wrapped(ctx, *args, **kwargs)
+
+    def wrapped_stream(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> VersionedResult[
+        ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]
+    ]:
+        """Stream the versioned function and return a wrapped StreamResponse."""
+        return self._stream_versioned.wrapped(ctx, *args, **kwargs)
+
+
+@dataclass(kw_only=True)
+class VersionedAsyncContextCall(_BaseVersionedCall, Generic[P, DepsT, FormattableT]):
+    """Versioned wrapper for asynchronous AsyncContextCall objects.
+
+    Example:
+        ```python
+        @ops.version
+        @llm.call("gpt-4o-mini")
+        async def recommend_book(ctx: llm.Context[str], genre: str):
+            return f"{ctx.deps} Recommend a {genre} book"
+
+        ctx = llm.Context(deps="As a librarian,")
+
+        # Returns AsyncContextResponse directly (but execution is versioned)
+        response = await recommend_book(ctx, "fantasy")
+        print(response.content)
+
+        # Use wrapped to get AsyncVersionedResult[AsyncContextResponse] with version info
+        result = await recommend_book.wrapped(ctx, "fantasy")
+        print(result.result.content)
+        ```
+    """
+
+    _call: AsyncContextCall[P, DepsT, FormattableT]
+    """The wrapped AsyncContextCall object."""
+
+    _call_versioned: AsyncVersionedFunction[
+        Concatenate[Context[DepsT], P],
+        AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """AsyncVersionedFunction wrapping the call method."""
+
+    _stream_versioned: AsyncVersionedFunction[
+        Concatenate[Context[DepsT], P],
+        AsyncContextStreamResponse[DepsT, None]
+        | AsyncContextStreamResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """AsyncVersionedFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize AsyncVersionedFunction wrappers for call and stream methods."""
+        self._call_versioned = AsyncVersionedFunction(
+            fn=self._call.call,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self._stream_versioned = AsyncVersionedFunction(
+            fn=self._call.stream,
+            tags=self.tags,
+            name=self.name,
+            metadata=self.metadata,
+        )
+        self._compute_closure(self._call, self._call_versioned, self._stream_versioned)
+
+    async def call(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
+        """Call the versioned function and return AsyncContextResponse directly."""
+        return await self._call_versioned(ctx, *args, **kwargs)
+
+    async def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
+        """Call the versioned function and return AsyncContextResponse directly."""
+        return await self.call(ctx, *args, **kwargs)
+
+    async def stream(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> (
+        AsyncContextStreamResponse[DepsT, None]
+        | AsyncContextStreamResponse[DepsT, FormattableT]
+    ):
+        """Stream the versioned function and return an AsyncContextStreamResponse."""
+        return await self._stream_versioned(ctx, *args, **kwargs)
+
+    async def wrapped(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[
+        AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]
+    ]:
+        """Call the versioned function and return a wrapped Response."""
+        return await self._call_versioned.wrapped(ctx, *args, **kwargs)
+
+    async def wrapped_stream(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncVersionedResult[
+        AsyncContextStreamResponse[DepsT, None]
+        | AsyncContextStreamResponse[DepsT, FormattableT]
+    ]:
+        """Stream the versioned function and return a wrapped StreamResponse."""
+        return await self._stream_versioned.wrapped(ctx, *args, **kwargs)