mirascope 1.25.7__py3-none-any.whl → 2.0.0a0__py3-none-any.whl

mirascope/llm/responses/streams.py

@@ -0,0 +1,363 @@
+"""Stream classes for streaming assistant content parts."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterator
+from typing import Generic, Literal, TypeAlias, TypeVar
+
+from ..content import (
+    AssistantContentPart,
+    Text,
+    TextChunk,
+    Thought,
+    ThoughtChunk,
+    ToolCall,
+    ToolCallChunk,
+)
+
+ContentT = TypeVar("ContentT", bound=AssistantContentPart)
+"""Type variable for content types (Text, Thought, ToolCall)."""
+
+DeltaT = TypeVar("DeltaT")
+"""Type variable for the deltas that the stream provides on iteration."""
+
+
+class BaseStream(ABC, Generic[ContentT, DeltaT]):
+    """Base class for synchronous streaming content."""
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[DeltaT]:
+        """Iterate over content deltas as they arrive."""
+        ...
+
+    @abstractmethod
+    def collect(self) -> ContentT:
+        """Collect all chunks and return the final content."""
+        ...
+
+
+class BaseAsyncStream(ABC, Generic[ContentT, DeltaT]):
+    """Base class for asynchronous streaming content."""
+
+    @abstractmethod
+    def __aiter__(self) -> AsyncIterator[DeltaT]:
+        """Asynchronously iterate over content deltas as they arrive."""
+        ...
+
+    @abstractmethod
+    async def collect(self) -> ContentT:
+        """Asynchronously collect all chunks and return the final content."""
+        ...
+
+
+class TextStream(BaseStream[Text, str]):
+    """Synchronous text stream implementation."""
+
+    type: Literal["text_stream"] = "text_stream"
+
+    content_type: Literal["text"] = "text"
+    """The type of content stored in this stream."""
+
+    partial_text: str
+    """The accumulated text content as chunks are received."""
+
+    def __init__(
+        self,
+        chunk_iterator: Iterator[TextChunk],
+    ) -> None:
+        """Initialize TextStream with a start chunk and chunk iterator.
+
+        Args:
+            chunk_iterator: Iterator providing subsequent chunks.
+        """
+        self.partial_text = ""
+        self._chunk_iterator = chunk_iterator
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over text deltas as they are received.
+
+        Yields:
+            str: Each delta string containing text.
+        """
+        for chunk in self._chunk_iterator:
+            delta = chunk.delta
+            self.partial_text += delta
+            yield delta
+
+    def collect(self) -> Text:
+        """Collect all chunks and return the final Text content.
+
+        Returns:
+            Text: The complete text content after consuming all chunks.
+        """
+        for _ in self:
+            pass
+        return Text(text=self.partial_text)
+
+
+class AsyncTextStream(BaseAsyncStream[Text, str]):
+    """Asynchronous text stream implementation."""
+
+    type: Literal["async_text_stream"] = "async_text_stream"
+
+    content_type: Literal["text"] = "text"
+    """The type of content stored in this stream."""
+
+    partial_text: str
+    """The accumulated text content as chunks are received."""
+
+    def __init__(
+        self,
+        chunk_iterator: AsyncIterator[TextChunk],
+    ) -> None:
+        """Initialize AsyncTextStream with a chunk iterator.
+
+        Args:
+            chunk_iterator: AsyncIterator providing subsequent chunks.
+        """
+        self.partial_text = ""
+        self._chunk_iterator = chunk_iterator
+
+    async def __aiter__(self) -> AsyncIterator[str]:
+        """Asynchronously iterate over text deltas as they are received.
+
+        Yields:
+            str: Each delta string containing text.
+        """
+        async for chunk in self._chunk_iterator:
+            delta = chunk.delta
+            self.partial_text += delta
+            yield delta
+
+    async def collect(self) -> Text:
+        """Asynchronously collect all chunks and return the final Text content.
+
+        Returns:
+            Text: The complete text content after consuming all chunks.
+        """
+        async for _ in self:
+            pass
+        return Text(text=self.partial_text)
+
+
+class ThoughtStream(BaseStream[Thought, str]):
+    """Synchronous thought stream implementation."""
+
+    type: Literal["thought_stream"] = "thought_stream"
+
+    content_type: Literal["thought"] = "thought"
+    """The type of content stored in this stream."""
+
+    partial_thought: str
+    """The accumulated thought content as chunks are received."""
+
+    def __init__(
+        self,
+        chunk_iterator: Iterator[ThoughtChunk],
+    ) -> None:
+        """Initialize ThoughtStream with a chunk iterator.
+
+        Args:
+            chunk_iterator: Iterator providing subsequent chunks.
+        """
+        self.partial_thought = ""
+        self._chunk_iterator = chunk_iterator
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over thought deltas as they are received.
+
+        Yields:
+            str: Each delta string containing thought text.
+        """
+        for chunk in self._chunk_iterator:
+            delta = chunk.delta
+            self.partial_thought += delta
+            yield delta
+
+    def collect(self) -> Thought:
+        """Collect all chunks and return the final Thought content.
+
+        Returns:
+            Thought: The complete thought content after consuming all chunks.
+        """
+        for _ in self:
+            pass
+        return Thought(thought=self.partial_thought)
+
+
+class AsyncThoughtStream(BaseAsyncStream[Thought, str]):
+    """Asynchronous thought stream implementation."""
+
+    type: Literal["async_thought_stream"] = "async_thought_stream"
+
+    content_type: Literal["thought"] = "thought"
+    """The type of content stored in this stream."""
+
+    partial_thought: str
+    """The accumulated thought content as chunks are received."""
+
+    def __init__(
+        self,
+        chunk_iterator: AsyncIterator[ThoughtChunk],
+    ) -> None:
+        """Initialize AsyncThoughtStream with a chunk iterator.
+
+        Args:
+            chunk_iterator: AsyncIterator providing subsequent chunks.
+        """
+        self.partial_thought = ""
+        self._chunk_iterator = chunk_iterator
+
+    async def __aiter__(self) -> AsyncIterator[str]:
+        """Asynchronously iterate over thought deltas as they are received.
+
+        Yields:
+            str: Each delta string containing thought text.
+        """
+        async for chunk in self._chunk_iterator:
+            delta = chunk.delta
+            self.partial_thought += delta
+            yield delta
+
+    async def collect(self) -> Thought:
+        """Asynchronously collect all chunks and return the final Thought content.
+
+        Returns:
+            Thought: The complete thought content after consuming all chunks.
+        """
+        async for _ in self:
+            pass
+        return Thought(thought=self.partial_thought)
+
+
+class ToolCallStream(BaseStream[ToolCall, str]):
+    """Synchronous tool call stream implementation."""
+
+    type: Literal["tool_call_stream"] = "tool_call_stream"
+
+    content_type: Literal["tool_call"] = "tool_call"
+    """The type of content stored in this stream."""
+
+    tool_id: str
+    """A unique identifier for this tool call."""
+
+    tool_name: str
+    """The name of the tool being called."""
+
+    partial_args: str
+    """The accumulated tool arguments as chunks are received."""
+
+    def __init__(
+        self,
+        tool_id: str,
+        tool_name: str,
+        chunk_iterator: Iterator[ToolCallChunk],
+    ) -> None:
+        """Initialize ToolCallStream with tool metadata and chunk iterator.
+
+        Args:
+            tool_id: A unique identifier for this tool call.
+            tool_name: The name of the tool being called.
+            chunk_iterator: Iterator providing subsequent chunks.
+        """
+        self.tool_id = tool_id
+        self.tool_name = tool_name
+        self.partial_args = ""
+        self._chunk_iterator = chunk_iterator
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over tool call argument deltas as they are received.
+
+        Yields:
+            str: Each delta string containing JSON argument fragments.
+        """
+        for chunk in self._chunk_iterator:
+            delta = chunk.delta
+            self.partial_args += delta
+            yield delta
+
+    def collect(self) -> ToolCall:
+        """Collect all chunks and return the final ToolCall content.
+
+        Returns:
+            ToolCall: The complete tool call after consuming all chunks.
+        """
+        for _ in self:
+            pass
+        return ToolCall(id=self.tool_id, name=self.tool_name, args=self.partial_args)
+
+
+class AsyncToolCallStream(BaseAsyncStream[ToolCall, str]):
+    """Asynchronous tool call stream implementation."""
+
+    type: Literal["async_tool_call_stream"] = "async_tool_call_stream"
+
+    content_type: Literal["tool_call"] = "tool_call"
+    """The type of content stored in this stream."""
+
+    tool_id: str
+    """A unique identifier for this tool call."""
+
+    tool_name: str
+    """The name of the tool being called."""
+
+    partial_args: str
+    """The accumulated tool arguments as chunks are received."""
+
+    def __init__(
+        self,
+        tool_id: str,
+        tool_name: str,
+        chunk_iterator: AsyncIterator[ToolCallChunk],
+    ) -> None:
+        """Initialize AsyncToolCallStream with tool metadata and chunk iterator.
+
+        Args:
+            tool_id: A unique identifier for this tool call.
+            tool_name: The name of the tool being called.
+            chunk_iterator: AsyncIterator providing subsequent chunks.
+        """
+        self.tool_id = tool_id
+        self.tool_name = tool_name
+        self.partial_args = ""
+        self._chunk_iterator = chunk_iterator
+
+    async def __aiter__(self) -> AsyncIterator[str]:
+        """Asynchronously iterate over tool call argument deltas as they are received.
+
+        Yields:
+            str: Each delta string containing JSON argument fragments.
+        """
+        async for chunk in self._chunk_iterator:
+            delta = chunk.delta
+            self.partial_args += delta
+            yield delta
+
+    async def collect(self) -> ToolCall:
+        """Asynchronously collect all chunks and return the final ToolCall content.
+
+        Returns:
+            ToolCall: The complete tool call after consuming all chunks.
+        """
+        async for _ in self:
+            pass
+        return ToolCall(id=self.tool_id, name=self.tool_name, args=self.partial_args)
+
+
+Stream: TypeAlias = TextStream | ToolCallStream | ThoughtStream
+"""A synchronous assistant content stream."""
+
+AsyncStream: TypeAlias = AsyncTextStream | AsyncToolCallStream | AsyncThoughtStream
+"""An asynchronous assistant content stream."""
+
+__all__ = [
+    "AsyncStream",
+    "AsyncTextStream",
+    "AsyncThoughtStream",
+    "AsyncToolCallStream",
+    "BaseAsyncStream",
+    "BaseStream",
+    "Stream",
+    "TextStream",
+    "ThoughtStream",
+    "ToolCallStream",
+]

mirascope/llm/tools/__init__.py

@@ -0,0 +1,40 @@
+"""The Tools module for LLMs."""
+
+from . import protocols
+from .decorator import ToolDecorator, tool
+from .tool_schema import (
+    FORMAT_TOOL_NAME,
+    ToolParameterSchema,
+    ToolSchema,
+    ToolSchemaT,
+)
+from .toolkit import (
+    AsyncContextToolkit,
+    AsyncToolkit,
+    BaseToolkit,
+    ContextToolkit,
+    Toolkit,
+    ToolkitT,
+)
+from .tools import AsyncContextTool, AsyncTool, ContextTool, Tool, ToolT
+
+__all__ = [
+    "FORMAT_TOOL_NAME",
+    "AsyncContextTool",
+    "AsyncContextToolkit",
+    "AsyncTool",
+    "AsyncToolkit",
+    "BaseToolkit",
+    "ContextTool",
+    "ContextToolkit",
+    "Tool",
+    "ToolDecorator",
+    "ToolParameterSchema",
+    "ToolSchema",
+    "ToolSchemaT",
+    "ToolT",
+    "Toolkit",
+    "ToolkitT",
+    "protocols",
+    "tool",
+]

mirascope/llm/tools/_utils.py

@@ -0,0 +1,25 @@
+"""Utilities for tool type checking and validation."""
+
+import inspect
+from typing_extensions import TypeIs
+
+from ..context import DepsT, _utils as _context_utils
+from ..types import JsonableCovariantT, P
+from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn
+
+
+def is_context_tool_fn(
+    fn: ToolFn | AsyncToolFn | ContextToolFn | AsyncContextToolFn,
+) -> TypeIs[
+    ContextToolFn[DepsT, P, JsonableCovariantT]
+    | AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+]:
+    """Type guard to check if a function is a context tool function."""
+    return _context_utils.first_param_is_context(fn)
+
+
+def is_async_tool_fn(
+    fn: ToolFn | AsyncToolFn | ContextToolFn | AsyncContextToolFn,
+) -> TypeIs[AsyncToolFn | AsyncContextToolFn]:
+    """Type guard to check if a function is an async tool function."""
+    return inspect.iscoroutinefunction(fn)

mirascope/llm/tools/decorator.py

@@ -0,0 +1,175 @@
+"""The `llm.tool` decorator for turning functions into tools."""
+
+from dataclasses import dataclass
+from typing import overload
+
+from ..context import DepsT
+from ..types import JsonableCovariantT, P
+from . import _utils as _tool_utils
+from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn
+from .tools import AsyncContextTool, AsyncTool, ContextTool, Tool
+
+
+@dataclass(kw_only=True)
+class ToolDecorator:
+    """Protocol for the tool decorator."""
+
+    strict: bool
+    """Whether to use strict tool calling, if supported by the provider."""
+
+    @overload
+    def __call__(  # pyright:ignore[reportOverlappingOverload]
+        self, fn: ContextToolFn[DepsT, P, JsonableCovariantT]
+    ) -> ContextTool[DepsT, P, JsonableCovariantT]:
+        """Call the decorator with a context function."""
+        ...
+
+    @overload
+    def __call__(  # pyright:ignore[reportOverlappingOverload]
+        self, fn: AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+    ) -> AsyncContextTool[DepsT, P, JsonableCovariantT]:
+        """Call the decorator with an async context function."""
+        ...
+
+    @overload
+    def __call__(
+        self, fn: ToolFn[P, JsonableCovariantT]
+    ) -> Tool[P, JsonableCovariantT]:
+        """Call the decorator with a sync function."""
+        ...
+
+    @overload
+    def __call__(
+        self, fn: AsyncToolFn[P, JsonableCovariantT]
+    ) -> AsyncTool[P, JsonableCovariantT]:
+        """Call the decorator with an async function."""
+        ...
+
+    def __call__(
+        self,
+        fn: ContextToolFn[DepsT, P, JsonableCovariantT]
+        | AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+        | ToolFn[P, JsonableCovariantT]
+        | AsyncToolFn[P, JsonableCovariantT],
+    ) -> (
+        ContextTool[DepsT, P, JsonableCovariantT]
+        | AsyncContextTool[DepsT, P, JsonableCovariantT]
+        | Tool[P, JsonableCovariantT]
+        | AsyncTool[P, JsonableCovariantT]
+    ):
+        """Call the decorator with a function."""
+        is_context = _tool_utils.is_context_tool_fn(fn)
+        is_async = _tool_utils.is_async_tool_fn(fn)
+
+        if is_context and is_async:
+            return AsyncContextTool[DepsT, P, JsonableCovariantT](
+                fn, strict=self.strict
+            )
+        elif is_context:
+            return ContextTool[DepsT, P, JsonableCovariantT](fn, strict=self.strict)
+        elif is_async:
+            return AsyncTool[P, JsonableCovariantT](fn, strict=self.strict)
+        else:
+            return Tool[P, JsonableCovariantT](fn, strict=self.strict)
+
+
+@overload
+def tool(  # pyright:ignore[reportOverlappingOverload]
+    __fn: AsyncContextToolFn[DepsT, P, JsonableCovariantT],
+) -> AsyncContextTool[DepsT, P, JsonableCovariantT]:
+    """Overload for async context tool functions."""
+    ...
+
+
+@overload
+def tool(  # pyright:ignore[reportOverlappingOverload]
+    __fn: ContextToolFn[DepsT, P, JsonableCovariantT],
+) -> ContextTool[DepsT, P, JsonableCovariantT]:
+    """Overload for context tool functions."""
+    ...
+
+
+@overload
+def tool(__fn: AsyncToolFn[P, JsonableCovariantT]) -> AsyncTool[P, JsonableCovariantT]:
+    """Overload for regular async tool functions."""
+    ...
+
+
+@overload
+def tool(__fn: ToolFn[P, JsonableCovariantT]) -> Tool[P, JsonableCovariantT]:
+    """Overload for regular sync tool functions."""
+    ...
+
+
+@overload
+def tool(*, strict: bool = False) -> ToolDecorator:
+    """Overload for setting non-default arguments."""
+    ...
+
+
+def tool(
+    __fn: ContextToolFn[DepsT, P, JsonableCovariantT]
+    | AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+    | ToolFn[P, JsonableCovariantT]
+    | AsyncToolFn[P, JsonableCovariantT]
+    | None = None,
+    *,
+    strict: bool = False,
+) -> (
+    ContextTool[DepsT, P, JsonableCovariantT]
+    | AsyncContextTool[DepsT, P, JsonableCovariantT]
+    | Tool[P, JsonableCovariantT]
+    | AsyncTool[P, JsonableCovariantT]
+    | ToolDecorator
+):
+    '''Decorator that turns a function into a tool definition.
+
+    This decorator creates a `Tool` or `ContextTool` that can be used with `llm.call`.
+    The function's name, docstring, and type hints are used to generate the
+    tool's metadata.
+
+    If the first parameter is named 'ctx' or typed as `llm.Context[T]`, it creates
+    a ContextTool. Otherwise, it creates a regular Tool.
+
+    Args:
+        strict: Whether the tool should use strict mode when supported by the model.
+
+    Returns:
+        A decorator function that converts the function into a Tool or ContextTool.
+
+    Examples:
+
+        Regular tool:
+        ```python
+        from mirascope import llm
+
+        @llm.tool
+        def available_books() -> list[str]:
+            """Returns the list of available books."""
+            return ["The Name of the Wind"]
+        ```
+
+        Context tool:
+        ```python
+        from dataclasses import dataclass
+
+        from mirascope import llm
+
+
+        @dataclass
+        class Library:
+            books: list[str]
+
+
+        library = Library(books=["Mistborn", "Gödel, Escher, Bach", "Dune"])
+
+        @llm.tool
+        def available_books(ctx: llm.Context[Library]) -> list[str]:
+            """Returns the list of available books."""
+            return ctx.deps.books
+        ```
+    '''
+    decorator = ToolDecorator(strict=strict)
+    if __fn is None:
+        return decorator
+    return decorator(__fn)

mirascope/llm/tools/protocols.py

@@ -0,0 +1,96 @@
+from typing import Any, Protocol
+
+from ..context import Context, DepsT
+from ..types import JsonableCovariantT, P
+
+
+class ToolFn(Protocol[P, JsonableCovariantT]):
+    """Protocol for the tool function."""
+
+    __name__: str
+    """The name of the tool."""
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> JsonableCovariantT:
+        """Call the function with the given arguments."""
+        raise NotImplementedError()
+
+
+class AsyncToolFn(Protocol[P, JsonableCovariantT]):
+    """Protocol for the async tool function."""
+
+    __name__: str
+    """The name of the tool."""
+
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> JsonableCovariantT:
+        """Call the function with the given arguments."""
+        raise NotImplementedError()
+
+
+class ContextToolFn(Protocol[DepsT, P, JsonableCovariantT]):
+    """Protocol for the context tool function."""
+
+    __name__: str
+    """The name of the tool."""
+
+    def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> JsonableCovariantT:
+        """Call the function with the given arguments."""
+        raise NotImplementedError()
+
+
+class AsyncContextToolFn(Protocol[DepsT, P, JsonableCovariantT]):
+    """Protocol for the async context tool function."""
+
+    __name__: str
+    """The name of the tool."""
+
+    async def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> JsonableCovariantT:
+        """Call the function with the given arguments."""
+        raise NotImplementedError()
+
+
+class _KwargsCallable(Protocol[JsonableCovariantT]):
+    """Protocol for functions that can be called with `Any`-typed kwargs.
+
+    Used internally to type-cast tool functions for compatibility with
+    json.loads() output when executing tool calls.
+    """
+
+    def __call__(self, **kwargs: dict[str, Any]) -> JsonableCovariantT: ...
+
+
+class _AsyncKwargsCallable(Protocol[JsonableCovariantT]):
+    """Protocol for async functions that can be called with `Any`-typed kwargs.
+
+    Used internally to type-cast async tool functions for compatibility with
+    json.loads() output when executing tool calls.
+    """
+
+    async def __call__(self, **kwargs: dict[str, Any]) -> JsonableCovariantT: ...
+
+
+class _ContextKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
+    """Protocol for context functions that can be called with `Any`-typed kwargs.
+
+    Used internally to type-cast context tool functions for compatibility with
+    json.loads() output when executing tool calls.
+    """
+
+    def __call__(
+        self, ctx: Context[DepsT], **kwargs: dict[str, Any]
+    ) -> JsonableCovariantT: ...
+
+
+class _AsyncJsonKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
+    """Protocol for async context functions that can be called with `Any`-typed kwargs.
+
+    Used internally to type-cast async context tool functions for compatibility with
+    json.loads() output when executing tool calls.
+    """
+
+    async def __call__(
+        self, ctx: Context[DepsT], **kwargs: dict[str, Any]
+    ) -> JsonableCovariantT: ...