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

mirascope/core/base/stream.py

@@ -1,365 +0,0 @@
-"""This module contains the base classes for streaming responses from LLMs."""
-
-import datetime
-import inspect
-from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, Awaitable, Callable, Generator
-from functools import wraps
-from typing import (
-    Any,
-    ClassVar,
-    Generic,
-    ParamSpec,
-    TypeVar,
-    overload,
-)
-
-from ._utils import (
-    HandleStream,
-    HandleStreamAsync,
-    SetupCall,
-    get_fn_args,
-    get_metadata,
-    get_possible_user_message_param,
-)
-from .call_params import BaseCallParams
-from .call_response import BaseCallResponse
-from .call_response_chunk import BaseCallResponseChunk
-from .dynamic_config import BaseDynamicConfig
-from .metadata import Metadata
-from .tool import BaseTool
-
-_BaseCallResponseT = TypeVar("_BaseCallResponseT", bound=BaseCallResponse)
-_BaseCallResponseChunkT = TypeVar(
-    "_BaseCallResponseChunkT", bound=BaseCallResponseChunk
-)
-_UserMessageParamT = TypeVar("_UserMessageParamT")
-_AssistantMessageParamT = TypeVar("_AssistantMessageParamT")
-_ToolMessageParamT = TypeVar("_ToolMessageParamT")
-_MessageParamT = TypeVar("_MessageParamT")
-_BaseToolT = TypeVar("_BaseToolT", bound=BaseTool)
-_BaseCallParamsT = TypeVar("_BaseCallParamsT", bound=BaseCallParams)
-_BaseDynamicConfigT = TypeVar("_BaseDynamicConfigT", bound=BaseDynamicConfig)
-_FinishReason = TypeVar("_FinishReason")
-
-
-class BaseStream(
-    Generic[
-        _BaseCallResponseT,
-        _BaseCallResponseChunkT,
-        _UserMessageParamT,
-        _AssistantMessageParamT,
-        _ToolMessageParamT,
-        _MessageParamT,
-        _BaseToolT,
-        _BaseDynamicConfigT,
-        _BaseCallParamsT,
-        _FinishReason,
-    ],
-    ABC,
-):
-    """A base class for streaming responses from LLMs."""
-
-    stream: (
-        Generator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None, None]
-        | AsyncGenerator[
-            tuple[_BaseCallResponseChunkT, _BaseToolT | None],
-            None,
-        ]
-    )
-    content: str
-    metadata: Metadata
-    tool_types: list[type[_BaseToolT]] | None
-    call_response_type: type[_BaseCallResponseT]
-    model: str
-    prompt_template: str | None = None
-    fn_args: dict[str, Any] | None = None
-    dynamic_config: _BaseDynamicConfigT
-    messages: list[_MessageParamT]
-    call_params: _BaseCallParamsT
-    call_kwargs: dict[str, Any]
-    user_message_param: _UserMessageParamT | None = None
-    message_param: _AssistantMessageParamT
-    input_tokens: int | float | None = None
-    output_tokens: int | float | None = None
-    id: str | None = None
-    finish_reasons: list[_FinishReason] | None = None
-    start_time: float = 0
-    end_time: float = 0
-
-    _provider: ClassVar[str] = "NO PROVIDER"
-
-    def __init__(
-        self,
-        *,
-        stream: Generator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None, None]
-        | AsyncGenerator[
-            tuple[_BaseCallResponseChunkT, _BaseToolT | None],
-            None,
-        ],
-        metadata: Metadata,
-        tool_types: list[type[_BaseToolT]] | None,
-        call_response_type: type[_BaseCallResponseT],
-        model: str,
-        prompt_template: str,
-        fn_args: dict[str, Any],
-        dynamic_config: _BaseDynamicConfigT,
-        messages: list[_MessageParamT],
-        call_params: _BaseCallParamsT,
-        call_kwargs: dict[str, Any],
-    ) -> None:
-        """Initializes an instance of `BaseStream`."""
-        self.content = ""
-        self.stream = stream
-        self.metadata = metadata
-        self.tool_types = tool_types
-        self.call_response_type = call_response_type
-        self.model = model
-        self.prompt_template = prompt_template
-        self.fn_args = fn_args
-        self.dynamic_config = dynamic_config
-        self.messages = messages
-        self.call_params = call_params
-        self.call_kwargs = call_kwargs
-        self.user_message_param = get_possible_user_message_param(messages)  # type: ignore
-
-    def __iter__(
-        self,
-    ) -> Generator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None, None]:
-        """Iterator over the stream and stores useful information."""
-        assert isinstance(
-            self.stream, Generator
-        ), "Stream must be a generator for __iter__"
-        self.content, tool_calls = "", []
-        self.start_time = datetime.datetime.now().timestamp() * 1000
-        for chunk, tool in self.stream:
-            self._update_properties(chunk)
-            if tool:
-                tool_calls.append(tool.tool_call)  # type: ignore
-            yield chunk, tool
-        self.end_time = datetime.datetime.now().timestamp() * 1000
-        self.message_param = self._construct_message_param(
-            tool_calls or None, self.content
-        )
-
-    def __aiter__(
-        self,
-    ) -> AsyncGenerator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None]:
-        """Iterates over the stream and stores useful information."""
-        self.content = ""
-
-        async def generator() -> (
-            AsyncGenerator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None]
-        ):
-            assert isinstance(
-                self.stream, AsyncGenerator
-            ), "Stream must be an async generator for __aiter__"
-            tool_calls = []
-            async for chunk, tool in self.stream:
-                self._update_properties(chunk)
-                if tool:
-                    tool_calls.append(tool.tool_call)  # type: ignore
-                yield chunk, tool
-            self.message_param = self._construct_message_param(
-                tool_calls or None, self.content
-            )
-
-        return generator()
-
-    def _update_properties(self, chunk: _BaseCallResponseChunkT) -> None:
-        """Updates the properties of the stream."""
-        self.content += chunk.content
-        if chunk.input_tokens is not None:
-            self.input_tokens = (
-                chunk.input_tokens
-                if not self.input_tokens
-                else self.input_tokens + chunk.input_tokens
-            )
-        if chunk.output_tokens is not None:
-            self.output_tokens = (
-                chunk.output_tokens
-                if not self.output_tokens
-                else self.output_tokens + chunk.output_tokens
-            )
-        if chunk.model is not None:
-            self.model = chunk.model
-        if chunk.id is not None:
-            self.id = chunk.id
-        if chunk.finish_reasons is not None:
-            self.finish_reasons = chunk.finish_reasons
-
-    @property
-    @abstractmethod
-    def cost(self) -> float | None:
-        """Returns the cost of the stream."""
-        ...
-
-    @abstractmethod
-    def _construct_message_param(
-        self, tool_calls: list[Any] | None = None, content: str | None = None
-    ) -> _AssistantMessageParamT:
-        """Constructs the assistant message."""
-        ...
-
-    def tool_message_params(
-        self, tools_and_outputs: list[tuple[_BaseToolT, str]]
-    ) -> list[_ToolMessageParamT]:
-        """Returns the tool message parameters for tool call results.
-
-        Args:
-            tools_and_outputs: The list of tools and their outputs from which the tool
-                message parameters should be constructed.
-        """
-        return self.call_response_type.tool_message_params(tools_and_outputs)
-
-    @abstractmethod
-    def construct_call_response(self) -> _BaseCallResponseT:
-        """Constructs the call response."""
-        ...
-
-
-_BaseCallResponseChunkT = TypeVar(
-    "_BaseCallResponseChunkT", bound=BaseCallResponseChunk
-)
-_BaseDynamicConfigT = TypeVar("_BaseDynamicConfigT", bound=BaseDynamicConfig)
-_BaseClientT = TypeVar("_BaseClientT", bound=object)
-_BaseCallParamsT = TypeVar("_BaseCallParamsT", bound=BaseCallParams)
-_ResponseT = TypeVar("_ResponseT")
-_ResponseChunkT = TypeVar("_ResponseChunkT")
-_P = ParamSpec("_P")
-
-
-def stream_factory(  # noqa: ANN201
-    *,
-    TCallResponse: type[_BaseCallResponseT],
-    TStream: type[BaseStream],
-    setup_call: SetupCall[
-        _BaseClientT,
-        _BaseDynamicConfigT,
-        _BaseCallParamsT,
-        _ResponseT,
-        _ResponseChunkT,
-        _BaseToolT,
-    ],
-    handle_stream: HandleStream[_ResponseChunkT, _BaseCallResponseChunkT, _BaseToolT],
-    handle_stream_async: HandleStreamAsync[
-        _ResponseChunkT, _BaseCallResponseChunkT, _BaseToolT
-    ],
-):
-    @overload
-    def decorator(
-        fn: Callable[_P, _BaseDynamicConfigT],
-        model: str,
-        tools: list[type[BaseTool] | Callable] | None,
-        json_mode: bool,
-        client: _BaseClientT | None,
-        call_params: _BaseCallParamsT,
-    ) -> Callable[_P, TStream]: ...
-
-    @overload
-    def decorator(
-        fn: Callable[_P, Awaitable[_BaseDynamicConfigT]],
-        model: str,
-        tools: list[type[BaseTool] | Callable] | None,
-        json_mode: bool,
-        client: _BaseClientT | None,
-        call_params: _BaseCallParamsT,
-    ) -> Callable[_P, Awaitable[TStream]]: ...
-
-    def decorator(
-        fn: Callable[_P, _BaseDynamicConfigT | Awaitable[_BaseDynamicConfigT]],
-        model: str,
-        tools: list[type[BaseTool] | Callable] | None,
-        json_mode: bool,
-        client: _BaseClientT | None,
-        call_params: _BaseCallParamsT,
-    ) -> Callable[_P, TStream | Awaitable[TStream]]:
-        if inspect.iscoroutinefunction(fn):
-
-            @wraps(fn)
-            async def inner_async(*args: _P.args, **kwargs: _P.kwargs) -> TStream:
-                assert SetupCall.fn_is_async(fn)
-                fn_args = get_fn_args(fn, args, kwargs)
-                dynamic_config = await fn(*args, **kwargs)
-                create, prompt_template, messages, tool_types, call_kwargs = setup_call(
-                    model=model,
-                    client=client,
-                    fn=fn,
-                    fn_args=fn_args,
-                    dynamic_config=dynamic_config,
-                    tools=tools,
-                    json_mode=json_mode,
-                    call_params=call_params,
-                    extract=False,
-                )
-
-                async def generator() -> (
-                    AsyncGenerator[
-                        tuple[_BaseCallResponseChunkT, _BaseToolT | None], None
-                    ]
-                ):
-                    async for chunk, tool in handle_stream_async(
-                        await create(stream=True, **call_kwargs), tool_types
-                    ):
-                        yield chunk, tool
-
-                return TStream(
-                    stream=generator(),
-                    metadata=get_metadata(fn, dynamic_config),
-                    tool_types=tool_types,  # type: ignore
-                    call_response_type=TCallResponse,
-                    model=model,
-                    prompt_template=prompt_template,
-                    fn_args=fn_args,
-                    dynamic_config=dynamic_config,
-                    messages=messages,
-                    call_params=call_params,
-                    call_kwargs=call_kwargs,
-                )
-
-            return inner_async
-        else:
-
-            @wraps(fn)
-            def inner(*args: _P.args, **kwargs: _P.kwargs) -> TStream:
-                assert SetupCall.fn_is_sync(fn)
-                fn_args = get_fn_args(fn, args, kwargs)
-                dynamic_config = fn(*args, **kwargs)
-                create, prompt_template, messages, tool_types, call_kwargs = setup_call(
-                    model=model,
-                    client=client,
-                    fn=fn,
-                    fn_args=fn_args,
-                    dynamic_config=dynamic_config,
-                    tools=tools,
-                    json_mode=json_mode,
-                    call_params=call_params,
-                    extract=False,
-                )
-
-                def generator() -> (
-                    Generator[
-                        tuple[_BaseCallResponseChunkT, _BaseToolT | None], None, None
-                    ]
-                ):
-                    yield from handle_stream(
-                        create(stream=True, **call_kwargs), tool_types
-                    )
-
-                return TStream(
-                    stream=generator(),
-                    metadata=get_metadata(fn, dynamic_config),
-                    tool_types=tool_types,  # type: ignore
-                    call_response_type=TCallResponse,
-                    model=model,
-                    prompt_template=prompt_template,
-                    fn_args=fn_args,
-                    dynamic_config=dynamic_config,
-                    messages=messages,
-                    call_params=call_params,
-                    call_kwargs=call_kwargs,
-                )
-
-            return inner
-
-    return decorator

mirascope/core/base/structured_stream.py

@@ -1,251 +0,0 @@
-"""This module defines the base class for structured streams."""
-
-import inspect
-from collections.abc import (
-    AsyncGenerator,
-    AsyncIterable,
-    Awaitable,
-    Callable,
-    Generator,
-    Iterable,
-)
-from functools import wraps
-from typing import (
-    Generic,
-    ParamSpec,
-    TypeVar,
-    cast,
-    overload,
-)
-
-from pydantic import BaseModel
-
-from ._utils import (
-    BaseType,
-    GetJsonOutput,
-    SetupCall,
-    extract_tool_return,
-    setup_extract_tool,
-)
-from .call_params import BaseCallParams
-from .call_response import BaseCallResponse
-from .call_response_chunk import BaseCallResponseChunk
-from .dynamic_config import BaseDynamicConfig
-from .stream import BaseStream, stream_factory
-from .tool import BaseTool
-
-_BaseCallResponseT = TypeVar("_BaseCallResponseT", bound=BaseCallResponse)
-_BaseCallResponseChunkT = TypeVar(
-    "_BaseCallResponseChunkT", bound=BaseCallResponseChunk
-)
-_BaseToolT = TypeVar("_BaseToolT", bound=BaseTool)
-_BaseCallParamsT = TypeVar("_BaseCallParamsT", bound=BaseCallParams)
-_BaseDynamicConfigT = TypeVar("_BaseDynamicConfigT", bound=BaseDynamicConfig)
-_ResponseModelT = TypeVar("_ResponseModelT", bound=BaseModel | BaseType)
-
-
-class BaseStructuredStream(Generic[_ResponseModelT]):
-    """A base class for streaming structured outputs from LLMs."""
-
-    stream: BaseStream
-    response_model: type[_ResponseModelT]
-    constructed_response_model: _ResponseModelT
-
-    def __init__(
-        self,
-        *,
-        stream: BaseStream,
-        response_model: type[_ResponseModelT],
-    ) -> None:
-        """Initializes an instance of `BaseStructuredStream`."""
-        self.stream = stream
-        self.response_model = response_model
-
-    def __iter__(self) -> Generator[_ResponseModelT, None, None]:
-        """Iterates over the stream and extracts structured outputs."""
-        json_output = ""
-        for chunk, _ in self.stream:
-            json_output += chunk.content
-            if json_output and json_output[0] != "{":
-                try:
-                    json_start = json_output.index("{")
-                    json_output = json_output[json_start:]
-                except ValueError:
-                    json_output = ""
-            if chunk.model is not None:
-                self.stream.model = chunk.model
-            if json_output:
-                yield extract_tool_return(self.response_model, json_output, True)
-        if json_output:
-            json_output = json_output[: json_output.rfind("}") + 1]
-        self.constructed_response_model = extract_tool_return(
-            self.response_model, json_output, False
-        )
-        yield self.constructed_response_model
-
-    def __aiter__(self) -> AsyncGenerator[_ResponseModelT, None]:
-        """Iterates over the stream and extracts structured outputs."""
-
-        async def generator() -> AsyncGenerator[_ResponseModelT, None]:
-            json_output = ""
-            async for chunk, _ in self.stream:
-                json_output += chunk.content
-                if json_output and json_output[0] != "{":
-                    try:
-                        json_start = json_output.index("{")
-                        json_output = json_output[json_start:]
-                    except ValueError:
-                        json_output = ""
-                if chunk.model is not None:
-                    self.stream.model = chunk.model
-                if json_output:
-                    yield extract_tool_return(self.response_model, json_output, True)
-            if json_output:
-                json_output = json_output[: json_output.rfind("}") + 1]
-            self.constructed_response_model = extract_tool_return(
-                self.response_model, json_output, False
-            )
-            yield self.constructed_response_model
-
-        return generator()
-
-
-_BaseDynamicConfigT = TypeVar("_BaseDynamicConfigT", bound=BaseDynamicConfig)
-_BaseClientT = TypeVar("_BaseClientT", bound=object)
-_ResponseT = TypeVar("_ResponseT")
-_ResponseChunkT = TypeVar("_ResponseChunkT")
-_P = ParamSpec("_P")
-
-
-def structured_stream_factory(  # noqa: ANN201
-    *,
-    TCallResponse: type[_BaseCallResponseT],
-    TCallResponseChunk: type[_BaseCallResponseChunkT],
-    TStream: type[BaseStream],
-    TToolType: type[_BaseToolT],
-    setup_call: SetupCall[
-        _BaseClientT,
-        _BaseDynamicConfigT,
-        _BaseCallParamsT,
-        _ResponseT,
-        _ResponseChunkT,
-        _BaseToolT,
-    ],
-    get_json_output: GetJsonOutput[_BaseCallResponseChunkT],
-):
-    class CustomContentChunk(TCallResponseChunk):
-        json_output: str
-
-        @property
-        def content(self) -> str:
-            return self.json_output
-
-    @overload
-    def decorator(
-        fn: Callable[_P, _BaseDynamicConfigT],
-        model: str,
-        response_model: type[_ResponseModelT],
-        json_mode: bool,
-        client: _BaseClientT | None,
-        call_params: _BaseCallParamsT,
-    ) -> Callable[
-        _P,
-        Iterable[_ResponseModelT],
-    ]: ...
-
-    @overload
-    def decorator(
-        fn: Callable[_P, Awaitable[_BaseDynamicConfigT]],
-        model: str,
-        response_model: type[_ResponseModelT],
-        json_mode: bool,
-        client: _BaseClientT | None,
-        call_params: _BaseCallParamsT,
-    ) -> Callable[
-        _P,
-        Awaitable[AsyncIterable[_ResponseModelT]],
-    ]: ...
-
-    def decorator(
-        fn: Callable[_P, _BaseDynamicConfigT | Awaitable[_BaseDynamicConfigT]],
-        model: str,
-        response_model: type[_ResponseModelT],
-        json_mode: bool,
-        client: _BaseClientT | None,
-        call_params: _BaseCallParamsT,
-    ) -> Callable[
-        _P,
-        Iterable[_ResponseModelT] | Awaitable[AsyncIterable[_ResponseModelT]],
-    ]:
-        def handle_chunk(
-            chunk: _ResponseChunkT,
-        ) -> tuple[_BaseCallResponseChunkT, None]:
-            call_response_chunk = TCallResponseChunk(chunk=chunk)
-            json_output = get_json_output(call_response_chunk, json_mode)
-
-            call_response_chunk = cast(
-                _BaseCallResponseChunkT,
-                CustomContentChunk(chunk=chunk, json_output=json_output),  # pyright: ignore [reportAbstractUsage]
-            )
-            return call_response_chunk, None
-
-        def handle_stream(
-            stream: Generator[_ResponseChunkT, None, None],
-            tool_types: list[type[_BaseToolT]] | None,
-        ) -> Generator[tuple[_BaseCallResponseChunkT, None], None, None]:
-            for chunk in stream:
-                yield handle_chunk(chunk)
-
-        async def handle_stream_async(
-            stream: AsyncGenerator[_ResponseChunkT, None],
-            tool_types: list[type[_BaseToolT]] | None,
-        ) -> AsyncGenerator[tuple[_BaseCallResponseChunkT, None], None]:
-            async for chunk in stream:
-                yield handle_chunk(chunk)
-
-        stream_decorator = stream_factory(
-            TCallResponse=TCallResponse,
-            TStream=TStream,
-            setup_call=setup_call,
-            handle_stream=handle_stream,
-            handle_stream_async=handle_stream_async,
-        )
-
-        tool = setup_extract_tool(response_model, TToolType)
-        stream_decorator_kwargs = {
-            "model": model,
-            "tools": [tool],
-            "json_mode": json_mode,
-            "client": client,
-            "call_params": call_params,
-        }
-        if inspect.iscoroutinefunction(fn):
-
-            @wraps(fn)
-            async def inner_async(
-                *args: _P.args, **kwargs: _P.kwargs
-            ) -> AsyncIterable[_ResponseModelT]:
-                assert SetupCall.fn_is_async(fn)
-                return BaseStructuredStream[_ResponseModelT](
-                    stream=await stream_decorator(fn=fn, **stream_decorator_kwargs)(
-                        *args, **kwargs
-                    ),
-                    response_model=response_model,
-                )
-
-            return inner_async
-        else:
-
-            @wraps(fn)
-            def inner(*args: _P.args, **kwargs: _P.kwargs) -> Iterable[_ResponseModelT]:
-                assert SetupCall.fn_is_sync(fn)
-                return BaseStructuredStream[_ResponseModelT](
-                    stream=stream_decorator(fn=fn, **stream_decorator_kwargs)(
-                        *args, **kwargs
-                    ),
-                    response_model=response_model,
-                )
-
-            return inner
-
-    return decorator