coco-framework 0.1.0__py3-none-any.whl

coco_core/messages/block_translators/bedrock_converse.py

@@ -0,0 +1,329 @@
+"""Derivations of standard content blocks from Amazon (Bedrock Converse) content."""
+
+import base64
+from collections.abc import Iterator
+from typing import Any, cast
+
+from coco_core.messages import AIMessage, AIMessageChunk
+from coco_core.messages import content as types
+
+
+def _bytes_to_b64_str(bytes_: bytes) -> str:
+    return base64.b64encode(bytes_).decode("utf-8")
+
+
+def _populate_extras(
+    standard_block: types.ContentBlock, block: dict[str, Any], known_fields: set[str]
+) -> types.ContentBlock:
+    """Mutate a block, populating extras."""
+    if standard_block.get("type") == "non_standard":
+        return standard_block
+
+    for key, value in block.items():
+        if key not in known_fields:
+            if "extras" not in standard_block:
+                # Below type-ignores are because mypy thinks a non-standard block can
+                # get here, although we exclude them above.
+                standard_block["extras"] = {}  # type: ignore[typeddict-unknown-key]
+            standard_block["extras"][key] = value  # type: ignore[typeddict-item]
+
+    return standard_block
+
+
+def _convert_to_v1_from_converse_input(
+    content: list[types.ContentBlock],
+) -> list[types.ContentBlock]:
+    """Convert Bedrock Converse format blocks to v1 format.
+
+    During the `content_blocks` parsing process, we wrap blocks not recognized as a v1
+    block as a `'non_standard'` block with the original block stored in the `value`
+    field. This function attempts to unpack those blocks and convert any blocks that
+    might be Converse format to v1 ContentBlocks.
+
+    If conversion fails, the block is left as a `'non_standard'` block.
+
+    Args:
+        content: List of content blocks to process.
+
+    Returns:
+        Updated list with Converse blocks converted to v1 format.
+    """
+
+    def _iter_blocks() -> Iterator[types.ContentBlock]:
+        blocks: list[dict[str, Any]] = [
+            cast("dict[str, Any]", block)
+            if block.get("type") != "non_standard"
+            else block["value"]  # type: ignore[typeddict-item]  # this is only non-standard blocks
+            for block in content
+        ]
+        for block in blocks:
+            num_keys = len(block)
+
+            if num_keys == 1 and (text := block.get("text")):
+                yield {"type": "text", "text": text}
+
+            elif (
+                num_keys == 1
+                and (document := block.get("document"))
+                and isinstance(document, dict)
+                and "format" in document
+            ):
+                if document.get("format") == "pdf":
+                    if "bytes" in document.get("source", {}):
+                        file_block: types.FileContentBlock = {
+                            "type": "file",
+                            "base64": _bytes_to_b64_str(document["source"]["bytes"]),
+                            "mime_type": "application/pdf",
+                        }
+                        _populate_extras(file_block, document, {"format", "source"})
+                        yield file_block
+
+                    else:
+                        yield {"type": "non_standard", "value": block}
+
+                elif document["format"] == "txt":
+                    if "text" in document.get("source", {}):
+                        plain_text_block: types.PlainTextContentBlock = {
+                            "type": "text-plain",
+                            "text": document["source"]["text"],
+                            "mime_type": "text/plain",
+                        }
+                        _populate_extras(
+                            plain_text_block, document, {"format", "source"}
+                        )
+                        yield plain_text_block
+                    else:
+                        yield {"type": "non_standard", "value": block}
+
+                else:
+                    yield {"type": "non_standard", "value": block}
+
+            elif (
+                num_keys == 1
+                and (image := block.get("image"))
+                and isinstance(image, dict)
+                and "format" in image
+            ):
+                if "bytes" in image.get("source", {}):
+                    image_block: types.ImageContentBlock = {
+                        "type": "image",
+                        "base64": _bytes_to_b64_str(image["source"]["bytes"]),
+                        "mime_type": f"image/{image['format']}",
+                    }
+                    _populate_extras(image_block, image, {"format", "source"})
+                    yield image_block
+
+                else:
+                    yield {"type": "non_standard", "value": block}
+
+            elif block.get("type") in types.KNOWN_BLOCK_TYPES:
+                yield cast("types.ContentBlock", block)
+
+            else:
+                yield {"type": "non_standard", "value": block}
+
+    return list(_iter_blocks())
+
+
+def _convert_citation_to_v1(citation: dict[str, Any]) -> types.Annotation:
+    standard_citation: types.Citation = {"type": "citation"}
+    if "title" in citation:
+        standard_citation["title"] = citation["title"]
+    if (
+        (source_content := citation.get("source_content"))
+        and isinstance(source_content, list)
+        and all(isinstance(item, dict) for item in source_content)
+    ):
+        standard_citation["cited_text"] = "".join(
+            item.get("text", "") for item in source_content
+        )
+
+    known_fields = {"type", "source_content", "title", "index", "extras"}
+
+    for key, value in citation.items():
+        if key not in known_fields:
+            if "extras" not in standard_citation:
+                standard_citation["extras"] = {}
+            standard_citation["extras"][key] = value
+
+    return standard_citation
+
+
+def _convert_to_v1_from_converse(message: AIMessage) -> list[types.ContentBlock]:
+    """Convert Bedrock Converse message content to v1 format."""
+    if (
+        message.content == ""
+        and not message.additional_kwargs
+        and not message.tool_calls
+    ):
+        # Converse outputs multiple chunks containing response metadata
+        return []
+
+    # Translated through a local list so that reading `content_blocks` leaves
+    # `message.content` untouched.
+    content: list[str | dict[str, Any]] = (
+        [{"type": "text", "text": message.content}]
+        if isinstance(message.content, str)
+        else message.content
+    )
+
+    def _iter_blocks() -> Iterator[types.ContentBlock]:
+        for block in content:
+            if not isinstance(block, dict):
+                continue
+            block_type = block.get("type")
+
+            if block_type == "text":
+                if citations := block.get("citations"):
+                    text_block: types.TextContentBlock = {
+                        "type": "text",
+                        "text": block.get("text", ""),
+                        "annotations": [_convert_citation_to_v1(a) for a in citations],
+                    }
+                else:
+                    text_block = {"type": "text", "text": block["text"]}
+                if "index" in block:
+                    text_block["index"] = block["index"]
+                yield text_block
+
+            elif block_type == "reasoning_content":
+                reasoning_block: types.ReasoningContentBlock = {"type": "reasoning"}
+                if reasoning_content := block.get("reasoning_content"):
+                    if reasoning := reasoning_content.get("text"):
+                        reasoning_block["reasoning"] = reasoning
+                    if signature := reasoning_content.get("signature"):
+                        if "extras" not in reasoning_block:
+                            reasoning_block["extras"] = {}
+                        reasoning_block["extras"]["signature"] = signature
+
+                if "index" in block:
+                    reasoning_block["index"] = block["index"]
+
+                known_fields = {"type", "reasoning_content", "index", "extras"}
+                for key in block:
+                    if key not in known_fields:
+                        if "extras" not in reasoning_block:
+                            reasoning_block["extras"] = {}
+                        reasoning_block["extras"][key] = block[key]
+                yield reasoning_block
+
+            elif block_type == "tool_use":
+                if (
+                    isinstance(message, AIMessageChunk)
+                    and len(message.tool_call_chunks) == 1
+                    and message.chunk_position != "last"
+                ):
+                    # Isolated chunk
+                    chunk = message.tool_call_chunks[0]
+                    tool_call_chunk = types.ToolCallChunk(
+                        name=chunk.get("name"),
+                        id=chunk.get("id"),
+                        args=chunk.get("args"),
+                        type="tool_call_chunk",
+                    )
+                    index = chunk.get("index")
+                    if index is not None:
+                        tool_call_chunk["index"] = index
+                    yield tool_call_chunk
+                else:
+                    tool_call_block: types.ToolCall | None = None
+                    # Non-streaming or gathered chunk
+                    if len(message.tool_calls) == 1:
+                        tool_call_block = {
+                            "type": "tool_call",
+                            "name": message.tool_calls[0]["name"],
+                            "args": message.tool_calls[0]["args"],
+                            "id": message.tool_calls[0].get("id"),
+                        }
+                    elif call_id := block.get("id"):
+                        for tc in message.tool_calls:
+                            if tc.get("id") == call_id:
+                                tool_call_block = {
+                                    "type": "tool_call",
+                                    "name": tc["name"],
+                                    "args": tc["args"],
+                                    "id": tc.get("id"),
+                                }
+                                break
+                    if not tool_call_block:
+                        tool_call_block = {
+                            "type": "tool_call",
+                            "name": block.get("name", ""),
+                            "args": block.get("input", {}),
+                            "id": block.get("id", ""),
+                        }
+                    if "index" in block:
+                        tool_call_block["index"] = block["index"]
+                    yield tool_call_block
+
+            elif (
+                block_type == "input_json_delta"
+                and isinstance(message, AIMessageChunk)
+                and len(message.tool_call_chunks) == 1
+            ):
+                chunk = message.tool_call_chunks[0]
+                tool_call_chunk = types.ToolCallChunk(
+                    name=chunk.get("name"),
+                    id=chunk.get("id"),
+                    args=chunk.get("args"),
+                    type="tool_call_chunk",
+                )
+                index = chunk.get("index")
+                if index is not None:
+                    tool_call_chunk["index"] = index
+                yield tool_call_chunk
+
+            else:
+                if "index" in block:
+                    # Copy before lifting `index` off the block.
+                    value = block.copy()
+                    index = value.pop("index")
+                    new_block: types.NonStandardContentBlock = {
+                        "type": "non_standard",
+                        "value": value,
+                        "index": index,
+                    }
+                else:
+                    new_block = {"type": "non_standard", "value": block}
+                yield new_block
+
+    return list(_iter_blocks())
+
+
+def translate_content(message: AIMessage) -> list[types.ContentBlock]:
+    """Derive standard content blocks from a message with Bedrock Converse content.
+
+    Args:
+        message: The message to translate.
+
+    Returns:
+        The derived content blocks.
+    """
+    return _convert_to_v1_from_converse(message)
+
+
+def translate_content_chunk(message: AIMessageChunk) -> list[types.ContentBlock]:
+    """Derive standard content blocks from a chunk with Bedrock Converse content.
+
+    Args:
+        message: The message chunk to translate.
+
+    Returns:
+        The derived content blocks.
+    """
+    return _convert_to_v1_from_converse(message)
+
+
+def _register_bedrock_converse_translator() -> None:
+    """Register the Bedrock Converse translator with the central registry.
+
+    Run automatically when the module is imported.
+    """
+    from coco_core.messages.block_translators import (  # noqa: PLC0415
+        register_translator,
+    )
+
+    register_translator("bedrock_converse", translate_content, translate_content_chunk)
+
+
+_register_bedrock_converse_translator()

coco_core/messages/block_translators/coco_v0.py

@@ -0,0 +1,303 @@
+"""Derivations of standard content blocks from Coco v0 multimodal content."""
+
+from typing import Any, cast
+
+from coco_core.messages import content as types
+
+
+def _convert_v0_multimodal_input_to_v1(
+    content: list[types.ContentBlock],
+) -> list[types.ContentBlock]:
+    """Convert v0 multimodal blocks to v1 format.
+
+    During the `content_blocks` parsing process, we wrap blocks not recognized as a v1
+    block as a `'non_standard'` block with the original block stored in the `value`
+    field. This function attempts to unpack those blocks and convert any v0 format
+    blocks to v1 format.
+
+    If conversion fails, the block is left as a `'non_standard'` block.
+
+    Args:
+        content: List of content blocks to process.
+
+    Returns:
+        v1 content blocks.
+    """
+    converted_blocks = []
+    unpacked_blocks: list[dict[str, Any]] = [
+        cast("dict[str, Any]", block)
+        if block.get("type") != "non_standard"
+        else block["value"]  # type: ignore[typeddict-item]  # this is only non-standard blocks
+        for block in content
+    ]
+    for block in unpacked_blocks:
+        if block.get("type") in {"image", "audio", "file"} and "source_type" in block:
+            converted_block = _convert_legacy_v0_content_block_to_v1(block)
+            converted_blocks.append(cast("types.ContentBlock", converted_block))
+        elif block.get("type") in types.KNOWN_BLOCK_TYPES:
+            # Guard in case this function is used outside of the .content_blocks flow
+            converted_blocks.append(cast("types.ContentBlock", block))
+        else:
+            converted_blocks.append({"type": "non_standard", "value": block})
+
+    return converted_blocks
+
+
+def _convert_legacy_v0_content_block_to_v1(
+    block: dict[str, Any],
+) -> types.ContentBlock | dict[str, Any]:
+    """Convert a Coco v0 content block to v1 format.
+
+    Preserves unknown keys as extras to avoid data loss.
+
+    Returns the original block unchanged if it's not in v0 format.
+    """
+
+    def _extract_v0_extras(
+        block_dict: dict[str, Any], known_keys: set[str]
+    ) -> dict[str, Any]:
+        """Extract unknown keys from v0 block to preserve as extras.
+
+        Args:
+            block_dict: The original v0 block dictionary.
+            known_keys: Set of keys known to be part of the v0 format for this block.
+
+        Returns:
+            A dictionary of extra keys not part of the known v0 format.
+        """
+        return {k: v for k, v in block_dict.items() if k not in known_keys}
+
+    # Check if this is actually a v0 format block
+    block_type = block.get("type")
+    if block_type not in {"image", "audio", "file"} or "source_type" not in block:
+        # Not a v0 format block, return unchanged
+        return block
+
+    if block.get("type") == "image":
+        source_type = block.get("source_type")
+        if source_type == "url":
+            # image-url
+            known_keys = {"mime_type", "type", "source_type", "url"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_image_block(
+                    url=block["url"],
+                    mime_type=block.get("mime_type"),
+                    id=block["id"],
+                    **extras,
+                )
+
+            # Don't construct with an ID if not present in original block
+            v1_image_url = types.ImageContentBlock(type="image", url=block["url"])
+            if block.get("mime_type"):
+                v1_image_url["mime_type"] = block["mime_type"]
+
+            v1_image_url["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_image_url["extras"][key] = value
+            if v1_image_url["extras"] == {}:
+                del v1_image_url["extras"]
+
+            return v1_image_url
+        if source_type == "base64":
+            # image-base64
+            known_keys = {"mime_type", "type", "source_type", "data"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_image_block(
+                    base64=block["data"],
+                    mime_type=block.get("mime_type"),
+                    id=block["id"],
+                    **extras,
+                )
+
+            v1_image_base64 = types.ImageContentBlock(
+                type="image", base64=block["data"]
+            )
+            if block.get("mime_type"):
+                v1_image_base64["mime_type"] = block["mime_type"]
+
+            v1_image_base64["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_image_base64["extras"][key] = value
+            if v1_image_base64["extras"] == {}:
+                del v1_image_base64["extras"]
+
+            return v1_image_base64
+        if source_type == "id":
+            # image-id
+            known_keys = {"type", "source_type", "id"}
+            extras = _extract_v0_extras(block, known_keys)
+            # For id `source_type`, `id` is the file reference, not block ID
+            v1_image_id = types.ImageContentBlock(type="image", file_id=block["id"])
+
+            v1_image_id["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_image_id["extras"][key] = value
+            if v1_image_id["extras"] == {}:
+                del v1_image_id["extras"]
+
+            return v1_image_id
+    elif block.get("type") == "audio":
+        source_type = block.get("source_type")
+        if source_type == "url":
+            # audio-url
+            known_keys = {"mime_type", "type", "source_type", "url"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_audio_block(
+                    url=block["url"],
+                    mime_type=block.get("mime_type"),
+                    id=block["id"],
+                    **extras,
+                )
+
+            # Don't construct with an ID if not present in original block
+            v1_audio_url: types.AudioContentBlock = types.AudioContentBlock(
+                type="audio", url=block["url"]
+            )
+            if block.get("mime_type"):
+                v1_audio_url["mime_type"] = block["mime_type"]
+
+            v1_audio_url["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_audio_url["extras"][key] = value
+            if v1_audio_url["extras"] == {}:
+                del v1_audio_url["extras"]
+
+            return v1_audio_url
+        if source_type == "base64":
+            # audio-base64
+            known_keys = {"mime_type", "type", "source_type", "data"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_audio_block(
+                    base64=block["data"],
+                    mime_type=block.get("mime_type"),
+                    id=block["id"],
+                    **extras,
+                )
+
+            v1_audio_base64: types.AudioContentBlock = types.AudioContentBlock(
+                type="audio", base64=block["data"]
+            )
+            if block.get("mime_type"):
+                v1_audio_base64["mime_type"] = block["mime_type"]
+
+            v1_audio_base64["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_audio_base64["extras"][key] = value
+            if v1_audio_base64["extras"] == {}:
+                del v1_audio_base64["extras"]
+
+            return v1_audio_base64
+        if source_type == "id":
+            # audio-id
+            known_keys = {"type", "source_type", "id"}
+            extras = _extract_v0_extras(block, known_keys)
+            v1_audio_id: types.AudioContentBlock = types.AudioContentBlock(
+                type="audio", file_id=block["id"]
+            )
+
+            v1_audio_id["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_audio_id["extras"][key] = value
+            if v1_audio_id["extras"] == {}:
+                del v1_audio_id["extras"]
+
+            return v1_audio_id
+    elif block.get("type") == "file":
+        source_type = block.get("source_type")
+        if source_type == "url":
+            # file-url
+            known_keys = {"mime_type", "type", "source_type", "url"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_file_block(
+                    url=block["url"],
+                    mime_type=block.get("mime_type"),
+                    id=block["id"],
+                    **extras,
+                )
+
+            v1_file_url: types.FileContentBlock = types.FileContentBlock(
+                type="file", url=block["url"]
+            )
+            if block.get("mime_type"):
+                v1_file_url["mime_type"] = block["mime_type"]
+
+            v1_file_url["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_file_url["extras"][key] = value
+            if v1_file_url["extras"] == {}:
+                del v1_file_url["extras"]
+
+            return v1_file_url
+        if source_type == "base64":
+            # file-base64
+            known_keys = {"mime_type", "type", "source_type", "data"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_file_block(
+                    base64=block["data"],
+                    mime_type=block.get("mime_type"),
+                    id=block["id"],
+                    **extras,
+                )
+
+            v1_file_base64: types.FileContentBlock = types.FileContentBlock(
+                type="file", base64=block["data"]
+            )
+            if block.get("mime_type"):
+                v1_file_base64["mime_type"] = block["mime_type"]
+
+            v1_file_base64["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_file_base64["extras"][key] = value
+            if v1_file_base64["extras"] == {}:
+                del v1_file_base64["extras"]
+
+            return v1_file_base64
+        if source_type == "id":
+            # file-id
+            known_keys = {"type", "source_type", "id"}
+            extras = _extract_v0_extras(block, known_keys)
+            return types.create_file_block(file_id=block["id"], **extras)
+        if source_type == "text":
+            # file-text
+            known_keys = {"mime_type", "type", "source_type", "url"}
+            extras = _extract_v0_extras(block, known_keys)
+            if "id" in block:
+                return types.create_plaintext_block(
+                    # In v0, URL points to the text file content
+                    # TODO: attribute this claim
+                    text=block["url"],
+                    id=block["id"],
+                    **extras,
+                )
+
+            v1_file_text: types.PlainTextContentBlock = types.PlainTextContentBlock(
+                type="text-plain", text=block["url"], mime_type="text/plain"
+            )
+            if block.get("mime_type"):
+                v1_file_text["mime_type"] = block["mime_type"]
+
+            v1_file_text["extras"] = {}
+            for key, value in extras.items():
+                if value is not None:
+                    v1_file_text["extras"][key] = value
+            if v1_file_text["extras"] == {}:
+                del v1_file_text["extras"]
+
+            return v1_file_text
+
+    # If we can't convert, return the block unchanged
+    return block