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

mirascope/core/bedrock/_utils/_convert_message_params.py

@@ -1,111 +0,0 @@
-"""Utility for converting `BaseMessageParam` to `BedrockMessageParam`."""
-
-from typing import cast
-
-from ...base import BaseMessageParam
-from ...base._utils import get_image_type
-from ...base._utils._parse_content_template import _load_media
-from .._types import ConversationRoleType, InternalBedrockMessageParam
-
-
-def convert_message_params(
-    message_params: list[BaseMessageParam | InternalBedrockMessageParam],
-) -> list[InternalBedrockMessageParam]:
-    converted_message_params = []
-    for message_param in message_params:
-        if not isinstance(message_param, BaseMessageParam):
-            converted_message_params.append(message_param)
-        elif isinstance((content := message_param.content), str):
-            converted_message_params.append(
-                {
-                    "role": cast(ConversationRoleType, message_param.role),
-                    "content": [{"text": content}],
-                }
-            )
-        else:
-            converted_content = []
-            for part in content:
-                if part.type == "text":
-                    converted_content.append({"text": part.text})
-                elif part.type == "image":
-                    if part.media_type not in [
-                        "image/jpeg",
-                        "image/png",
-                        "image/gif",
-                        "image/webp",
-                    ]:
-                        raise ValueError(
-                            f"Unsupported image media type: {part.media_type}. Bedrock"
-                            " currently only supports JPEG, PNG, GIF, and WebP images."
-                        )
-                    converted_content.append(
-                        {
-                            "image": {
-                                "format": part.media_type.split("/")[-1],
-                                "source": {"bytes": part.image},
-                            }
-                        }
-                    )
-                elif part.type == "image_url":
-                    image = _load_media(part.url)
-                    converted_content.append(
-                        {
-                            "image": {
-                                "format": get_image_type(image),
-                                "source": {"bytes": image},
-                            },
-                        }
-                    )
-                elif part.type == "tool_result":
-                    if converted_content:
-                        converted_message_params.append(
-                            {"role": message_param.role, "content": converted_content}
-                        )
-                        converted_content = []
-
-                    converted_message_params.append(
-                        {
-                            "role": "user",
-                            "content": [
-                                {
-                                    "toolResult": {
-                                        "toolUseId": part.id,
-                                        "content": [{"text": part.content}]
-                                        if isinstance(part.content, str)
-                                        else part.content,
-                                    }
-                                }
-                            ],
-                        }
-                    )
-                elif part.type == "tool_call":
-                    if converted_content:
-                        converted_message_params.append(
-                            {"role": message_param.role, "content": converted_content}
-                        )
-                        converted_content = []
-
-                    converted_message_params.append(
-                        {
-                            "role": "assistant",
-                            "content": [
-                                {
-                                    "toolUse": {
-                                        "toolUseId": part.id,
-                                        "name": part.name,
-                                        "input": part.args,
-                                    }
-                                }
-                            ],
-                        }
-                    )
-                else:
-                    raise ValueError(
-                        "Bedrock currently only supports text and image parts. "
-                        f"Part provided: {part.type}"
-                    )
-            if converted_content:
-                converted_message_params.append(
-                    {"role": message_param.role, "content": converted_content}
-                )
-    return converted_message_params

mirascope/core/bedrock/_utils/_get_json_output.py

@@ -1,30 +0,0 @@
-"""Get the JSON output from a completion response."""
-
-import json
-
-from ..call_response import BedrockCallResponse
-from ..call_response_chunk import BedrockCallResponseChunk
-
-
-def get_json_output(
-    response: BedrockCallResponse | BedrockCallResponseChunk, json_mode: bool
-) -> str:
-    """Get the JSON output from a completion response."""
-    if isinstance(response, BedrockCallResponse):
-        if json_mode and response.content:
-            return response.content
-        elif message := response.message:
-            tool_calls = [t for c in message["content"] if (t := c.get("toolUse"))]
-            if tool_calls and (tool_call_input := tool_calls[0].get("input", {})):
-                return json.dumps(tool_call_input)
-        raise ValueError("No tool call or JSON object found in response.")
-    else:
-        if json_mode:
-            return response.content
-        elif (
-            (content_block_delta := response.chunk.get("contentBlockDelta"))
-            and (tool_use := content_block_delta["delta"].get("toolUse"))
-            and (tool_use_input := tool_use.get("input"))
-        ):
-            return tool_use_input
-    return ""

mirascope/core/bedrock/_utils/_handle_stream.py

@@ -1,104 +0,0 @@
-"""Handles the stream of completion chunks."""
-
-import json
-from collections.abc import AsyncGenerator, Generator
-
-from mypy_boto3_bedrock_runtime.type_defs import (
-    ToolUseBlockOutputTypeDef,
-)
-from typing_extensions import TypedDict
-
-from .._types import (
-    AsyncStreamOutputChunk,
-    StreamOutputChunk,
-    ToolUseBlockContentTypeDef,
-)
-from ..call_response_chunk import BedrockCallResponseChunk
-from ..tool import BedrockTool
-
-
-class ToolUseChunk(TypedDict):
-    tool_use_id: str
-    input_chunk: str
-    name: str
-    stop: bool
-
-
-def _handle_chunk(
-    chunk: StreamOutputChunk | AsyncStreamOutputChunk,
-    current_tool_use_chunk: ToolUseChunk | None,
-    tool_types: list[type[BedrockTool]] | None,
-) -> tuple[
-    BedrockCallResponseChunk | None,
-    BedrockTool | None,
-    ToolUseChunk | None,
-]:
-    """Handles a chunk of the stream."""
-    if not tool_types:
-        return BedrockCallResponseChunk(chunk=chunk), None, None
-    elif (content_block_start := chunk.get("contentBlockStart")) and (
-        tool_use := content_block_start["start"].get("toolUse")
-    ):
-        current_tool_use_chunk = ToolUseChunk(
-            tool_use_id=tool_use["toolUseId"],
-            input_chunk="",
-            name=tool_use["name"],
-            stop=False,
-        )
-    elif (
-        (content_block_delta := chunk.get("contentBlockDelta"))
-        and (tool_use := content_block_delta["delta"].get("toolUse"))
-        and current_tool_use_chunk
-        and not current_tool_use_chunk["stop"]
-    ):
-        current_tool_use_chunk["input_chunk"] += tool_use["input"]
-        return None, None, current_tool_use_chunk
-    elif "contentBlockStop" in chunk and current_tool_use_chunk:
-        current_tool_use_chunk["stop"] = True
-        return None, None, current_tool_use_chunk
-    elif current_tool_use_chunk and current_tool_use_chunk["stop"]:
-        for tool_type in tool_types:
-            if current_tool_use_chunk["name"] == tool_type._name():
-                current_tool_use = ToolUseBlockContentTypeDef(
-                    toolUse=ToolUseBlockOutputTypeDef(
-                        toolUseId=current_tool_use_chunk["tool_use_id"],
-                        input=json.loads(current_tool_use_chunk["input_chunk"]),
-                        name=current_tool_use_chunk["name"],
-                    )
-                )
-                return (
-                    BedrockCallResponseChunk(chunk=chunk),
-                    tool_type.from_tool_call(current_tool_use),
-                    None,
-                )
-    return BedrockCallResponseChunk(chunk=chunk), None, current_tool_use_chunk
-
-
-def handle_stream(
-    stream: Generator[StreamOutputChunk, None, None],
-    tool_types: list[type[BedrockTool]] | None,
-    partial_tools: bool = False,
-) -> Generator[tuple[BedrockCallResponseChunk, BedrockTool | None], None, None]:
-    """Iterator over the stream and constructs tools as they are streamed."""
-    current_tool_use_chunk = None
-    for chunk in stream:
-        call_response, tool, current_tool_use_chunk = _handle_chunk(
-            chunk, current_tool_use_chunk, tool_types
-        )
-        if call_response:
-            yield call_response, tool
-
-
-async def handle_stream_async(
-    stream: AsyncGenerator[AsyncStreamOutputChunk, None],
-    tool_types: list[type[BedrockTool]] | None,
-    partial_tools: bool = False,
-) -> AsyncGenerator[tuple[BedrockCallResponseChunk, BedrockTool | None], None]:
-    """Async iterator over the stream and constructs tools as they are streamed."""
-    current_tool_use_chunk = None
-    async for chunk in stream:
-        call_response, tool, current_tool_use_chunk = _handle_chunk(
-            chunk, current_tool_use_chunk, tool_types
-        )
-        if call_response:
-            yield call_response, tool

mirascope/core/bedrock/_utils/_message_param_converter.py

@@ -1,171 +0,0 @@
-from typing import cast
-
-from mirascope.core import BaseMessageParam
-from mirascope.core.base import (
-    DocumentPart,
-    ImagePart,
-    TextPart,
-    ToolCallPart,
-    ToolResultPart,
-)
-from mirascope.core.base._utils._base_message_param_converter import (
-    BaseMessageParamConverter,
-)
-
-from .._types import (
-    InternalBedrockMessageParam,
-)
-from . import convert_message_params
-
-IMAGE_FORMAT_MAP = {
-    "JPEG": "image/jpeg",
-    "PNG": "image/png",
-    "GIF": "image/gif",
-    "WEBP": "image/webp",
-}
-
-
-class BedrockMessageParamConverter(BaseMessageParamConverter):
-    """Converts between Bedrock `InternalBedrockMessageParam` and Mirascope `BaseMessageParam`."""
-
-    @staticmethod
-    def to_provider(
-        message_params: list[BaseMessageParam],
-    ) -> list[InternalBedrockMessageParam]:
-        """
-        Convert from Mirascope `BaseMessageParam` to Bedrock's `InternalBedrockMessageParam`.
-        """
-        return convert_message_params(
-            cast(list[BaseMessageParam | InternalBedrockMessageParam], message_params)
-        )
-
-    @staticmethod
-    def from_provider(
-        message_params: list[InternalBedrockMessageParam],
-    ) -> list[BaseMessageParam]:
-        """
-        Convert from Bedrock's `InternalBedrockMessageParam` to Mirascope `BaseMessageParam`.
-        """
-        converted = []
-        for message_param in message_params:
-            message_param["role"]
-            content_blocks = message_param["content"]
-
-            converted_content = []
-            for block in content_blocks:
-                if "text" in block:
-                    text = block["text"]
-                    if not isinstance(text, str):
-                        raise ValueError("Text content must be a string.")
-                    converted_content.append(TextPart(type="text", text=text))
-
-                elif "image" in block:
-                    image_block = block["image"]
-                    img_format = image_block["format"]  # e.g. "JPEG"
-                    source = image_block["source"]
-                    if "bytes" not in source or not isinstance(source["bytes"], bytes):
-                        raise ValueError(
-                            "Image block must have 'source.bytes' as bytes."
-                        )
-                    media_type = IMAGE_FORMAT_MAP.get(img_format.upper())
-                    if not media_type:
-                        raise ValueError(f"Unsupported image format: {img_format}")
-                    converted_content.append(
-                        ImagePart(
-                            type="image",
-                            media_type=media_type,
-                            image=source["bytes"],
-                            detail=None,
-                        )
-                    )
-
-                elif "document" in block:
-                    doc_block = block["document"]
-                    doc_format = doc_block["format"]  # e.g. "PDF"
-                    if doc_format.upper() != "PDF":
-                        raise ValueError(
-                            f"Unsupported document format: {doc_format}. Only PDF is supported."
-                        )
-                    source = doc_block["source"]
-                    if "bytes" not in source or not isinstance(source["bytes"], bytes):
-                        raise ValueError(
-                            "Document block must have 'source.bytes' as bytes."
-                        )
-                    converted_content.append(
-                        DocumentPart(
-                            type="document",
-                            media_type="application/pdf",
-                            document=source["bytes"],
-                        )
-                    )
-                elif "toolUse" in block:
-                    tool_use = block["toolUse"]
-                    if converted_content:
-                        converted.append(
-                            BaseMessageParam(
-                                role=message_param["role"], content=converted_content
-                            )
-                        )
-                        converted_content = []
-
-                    converted.append(
-                        BaseMessageParam(
-                            role="assistant",
-                            content=[
-                                ToolCallPart(
-                                    type="tool_call",
-                                    name=tool_use["name"],
-                                    id=tool_use["toolUseId"],
-                                    args=tool_use["input"],
-                                )
-                            ],
-                        )
-                    )
-                elif "toolResult" in block:
-                    tool_result = block["toolResult"]
-                    if converted_content:
-                        converted.append(
-                            BaseMessageParam(
-                                role=message_param["role"], content=converted_content
-                            )
-                        )
-                        converted_content = []
-                    converted.append(
-                        BaseMessageParam(
-                            role="user",
-                            content=[
-                                ToolResultPart(
-                                    type="tool_result",
-                                    id=tool_result["toolUseId"],
-                                    name=tool_result["name"],  # pyright: ignore [reportGeneralTypeIssues]
-                                    content=tool_result["content"]
-                                    if isinstance(tool_result["content"], str)
-                                    else tool_result["content"][0]["text"],  # pyright: ignore [reportTypedDictNotRequiredAccess]
-                                    is_error=tool_result.get("isError", False),
-                                )
-                            ],
-                        )
-                    )
-                else:
-                    raise ValueError(
-                        "Content block does not contain supported content."
-                    )
-            if len(converted_content) == 1 and isinstance(
-                converted_content[0], TextPart
-            ):
-                converted.append(
-                    BaseMessageParam(
-                        role=message_param["role"],
-                        content=converted_content[0].text,
-                    )
-                )
-            else:
-                if converted_content:
-                    converted.append(
-                        BaseMessageParam(
-                            role=message_param["role"],
-                            content=converted_content,
-                        )
-                    )
-
-        return converted

mirascope/core/bedrock/_utils/_setup_call.py

@@ -1,258 +0,0 @@
-"""This module contains the setup_call function for Bedrock tools."""
-
-from __future__ import annotations
-
-import os
-from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Generator
-from functools import wraps
-from typing import Any, ParamSpec, cast, overload
-
-import aiobotocore.client
-from aiobotocore.session import AioSession, get_session
-from boto3.session import Session
-from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
-from mypy_boto3_bedrock_runtime.type_defs import (
-    ConverseResponseTypeDef,
-    ConverseStreamResponseTypeDef,
-    ToolChoiceTypeDef,
-    ToolConfigurationTypeDef,
-)
-from pydantic import BaseModel
-from types_aiobotocore_bedrock_runtime import (
-    BedrockRuntimeClient as AsyncBedrockRuntimeClient,
-)
-from types_aiobotocore_bedrock_runtime.type_defs import (
-    ConverseRequestRequestTypeDef as AsyncConverseRequestRequestTypeDef,
-)
-from types_aiobotocore_bedrock_runtime.type_defs import (
-    ConverseResponseTypeDef as AsyncConverseResponseTypeDef,
-)
-from types_aiobotocore_bedrock_runtime.type_defs import (
-    ConverseStreamRequestRequestTypeDef as AsyncConverseStreamRequestRequestTypeDef,
-)
-from types_aiobotocore_bedrock_runtime.type_defs import (
-    ConverseStreamResponseTypeDef as AsyncConverseStreamResponseTypeDef,
-)
-from typing_extensions import Unpack
-
-from ... import BaseMessageParam
-from ...base import BaseTool, _utils
-from ...base._utils import (
-    AsyncCreateFn,
-    CreateFn,
-    fn_is_async,
-    get_async_create_fn,
-    get_create_fn,
-)
-from ...base.call_params import CommonCallParams
-from ...base.stream_config import StreamConfig
-from .._call_kwargs import BedrockCallKwargs
-from .._types import (
-    AsyncStreamOutputChunk,
-    InternalBedrockMessageParam,
-    StreamOutputChunk,
-)
-from ..call_params import BedrockCallParams
-from ..dynamic_config import AsyncBedrockDynamicConfig, BedrockDynamicConfig
-from ..tool import BedrockTool
-from ._convert_common_call_params import convert_common_call_params
-from ._convert_message_params import convert_message_params
-
-_P = ParamSpec("_P")
-
-
-def _extract_sync_stream_fn(
-    fn: Callable[_P, ConverseStreamResponseTypeDef], model: str
-) -> Callable[_P, Generator[StreamOutputChunk, None, None]]:
-    @wraps(fn)
-    def _inner(
-        *args: _P.args, **kwargs: _P.kwargs
-    ) -> Generator[StreamOutputChunk, None, None]:
-        response = fn(*args, **kwargs)
-        for chunk in response["stream"]:
-            yield StreamOutputChunk(
-                responseMetadata=response["ResponseMetadata"], model=model, **chunk
-            )
-
-    return _inner
-
-
-def _extract_async_stream_fn(
-    fn: Callable[_P, Coroutine[Any, Any, AsyncConverseStreamResponseTypeDef]],
-    model: str,
-) -> Callable[_P, AsyncGenerator[AsyncStreamOutputChunk, None]]:
-    @wraps(fn)
-    async def _inner(
-        *args: _P.args, **kwargs: _P.kwargs
-    ) -> AsyncGenerator[AsyncStreamOutputChunk, None]:
-        response = await fn(*args, **kwargs)
-        async for chunk in response["stream"]:
-            yield AsyncStreamOutputChunk(
-                responseMetadata=response["ResponseMetadata"], model=model, **chunk
-            )
-
-    return _inner
-
-
-class _AsyncBedrockRuntimeWrappedClient:
-    def __init__(self, session: AioSession, model: str) -> None:
-        self.session: AioSession = session
-        self.model: str = model
-
-    async def converse(
-        self, **kwargs: Unpack[AsyncConverseRequestRequestTypeDef]
-    ) -> AsyncConverseResponseTypeDef:
-        async with self.session.create_client("bedrock-runtime") as client:
-            return await client.converse(**kwargs)
-
-    async def converse_stream(
-        self, **kwargs: Unpack[AsyncConverseStreamRequestRequestTypeDef]
-    ) -> AsyncGenerator[AsyncStreamOutputChunk, None]:
-        async with self.session.create_client("bedrock-runtime") as client:
-            response = await client.converse_stream(**kwargs)
-            async for chunk in response["stream"]:
-                yield AsyncStreamOutputChunk(
-                    responseMetadata=response["ResponseMetadata"],
-                    model=self.model,
-                    **chunk,
-                )
-
-
-@overload
-def setup_call(
-    *,
-    model: str,
-    client: AsyncBedrockRuntimeClient | None,
-    fn: Callable[..., Awaitable[AsyncBedrockDynamicConfig]],
-    fn_args: dict[str, Any],
-    dynamic_config: AsyncBedrockDynamicConfig,
-    tools: list[type[BaseTool] | Callable] | None,
-    json_mode: bool,
-    call_params: BedrockCallParams | CommonCallParams,
-    response_model: type[BaseModel] | None,
-    stream: bool | StreamConfig,
-) -> tuple[
-    AsyncCreateFn[AsyncConverseResponseTypeDef, AsyncStreamOutputChunk],
-    str | None,
-    list[InternalBedrockMessageParam],
-    list[type[BedrockTool]] | None,
-    BedrockCallKwargs,
-]: ...
-
-
-@overload
-def setup_call(
-    *,
-    model: str,
-    client: BedrockRuntimeClient | None,
-    fn: Callable[..., BedrockDynamicConfig],
-    fn_args: dict[str, Any],
-    dynamic_config: BedrockDynamicConfig,
-    tools: list[type[BaseTool] | Callable] | None,
-    json_mode: bool,
-    call_params: BedrockCallParams | CommonCallParams,
-    response_model: type[BaseModel] | None,
-    stream: bool | StreamConfig,
-) -> tuple[
-    CreateFn[ConverseResponseTypeDef, StreamOutputChunk],
-    str | None,
-    list[InternalBedrockMessageParam],
-    list[type[BedrockTool]] | None,
-    BedrockCallKwargs,
-]: ...
-
-
-def setup_call(
-    *,
-    model: str,
-    client: BedrockRuntimeClient | AsyncBedrockRuntimeClient | None,
-    fn: Callable[..., BedrockDynamicConfig]
-    | Callable[..., Awaitable[AsyncBedrockDynamicConfig]],
-    fn_args: dict[str, Any],
-    dynamic_config: BedrockDynamicConfig | AsyncBedrockDynamicConfig,
-    tools: list[type[BaseTool] | Callable] | None,
-    json_mode: bool,
-    call_params: BedrockCallParams | CommonCallParams,
-    response_model: type[BaseModel] | None,
-    stream: bool | StreamConfig,
-) -> tuple[
-    AsyncCreateFn[AsyncConverseResponseTypeDef, AsyncStreamOutputChunk]
-    | CreateFn[ConverseResponseTypeDef, StreamOutputChunk],
-    str | None,
-    list[InternalBedrockMessageParam],
-    list[type[BedrockTool]] | None,
-    BedrockCallKwargs,
-]:
-    prompt_template, messages, tool_types, base_call_kwargs = _utils.setup_call(
-        fn,
-        fn_args,
-        dynamic_config,
-        tools,
-        BedrockTool,
-        call_params,
-        convert_common_call_params,
-    )
-    call_kwargs = cast(BedrockCallKwargs, base_call_kwargs)
-    messages = cast(list[InternalBedrockMessageParam | BaseMessageParam], messages)
-    messages = convert_message_params(messages)
-    if messages[0]["role"] == "system":
-        call_kwargs["system"] = [
-            {"text": text}
-            for c in messages.pop(0)["content"]
-            if (text := c.get("text"))
-        ]
-
-    if call_kwargs_tools := call_kwargs.pop("tools", None):
-        call_kwargs["toolConfig"] = cast(
-            ToolConfigurationTypeDef, {"tools": call_kwargs_tools}
-        )
-    if json_mode:
-        json_mode_content = _utils.json_mode_content(response_model)
-        if "text" in messages[-1]["content"][-1]:
-            messages[-1]["content"][-1]["text"] += json_mode_content
-        else:
-            messages[-1]["content"] = [
-                *messages[-1]["content"],
-                {"text": json_mode_content},
-            ]
-    elif response_model:
-        assert tool_types, "At least one tool must be provided for extraction."
-        if "toolConfig" in call_kwargs:
-            call_kwargs["toolConfig"]["toolChoice"] = cast(
-                ToolChoiceTypeDef, {"tool": {"name": tool_types[0]._name()}}
-            )
-
-    call_kwargs |= cast(BedrockCallKwargs, {"modelId": model, "messages": messages})
-
-    env_vars = {}
-    if access_key_id := os.getenv("AWS_ACCESS_KEY_ID"):
-        env_vars["aws_access_key_id"] = access_key_id
-    if secret_access_key := os.getenv("AWS_SECRET_ACCESS_KEY"):
-        env_vars["aws_secret_access_key"] = secret_access_key
-    if session_token := os.getenv("AWS_SESSION_TOKEN"):
-        env_vars["aws_session_token"] = session_token
-    if region_name := os.getenv("AWS_REGION_NAME"):
-        env_vars["region_name"] = region_name
-    if profile_name := os.getenv("AWS_PROFILE"):
-        env_vars["profile_name"] = profile_name
-    if client is None:
-        if fn_is_async(fn):
-            session = get_session(env_vars=env_vars)
-            _client = _AsyncBedrockRuntimeWrappedClient(session, model)
-        else:
-            session = Session(**env_vars)
-            _client = session.client("bedrock-runtime")
-    else:
-        _client = client
-    if isinstance(_client, aiobotocore.client.AioBaseClient):
-        create = get_async_create_fn(
-            _client.converse, _extract_async_stream_fn(_client.converse_stream, model)
-        )
-    elif isinstance(_client, _AsyncBedrockRuntimeWrappedClient):
-        create = get_async_create_fn(_client.converse, _client.converse_stream)
-
-    else:
-        create = get_create_fn(
-            _client.converse, _extract_sync_stream_fn(_client.converse_stream, model)
-        )
-    return create, prompt_template, messages, tool_types, call_kwargs