fix tests

This commit is contained in:
Chester Curme
2025-07-22 16:45:19 -04:00
parent b24f90dabe
commit b1a02f971b
4 changed files with 58 additions and 40 deletions

View File

@@ -47,11 +47,11 @@ if TYPE_CHECKING:
NonStandardContentBlock, NonStandardContentBlock,
PlainTextContentBlock, PlainTextContentBlock,
ReasoningContentBlock, ReasoningContentBlock,
SearchCall,
SearchResult,
TextContentBlock, TextContentBlock,
ToolCallContentBlock, ToolCallContentBlock,
VideoContentBlock, VideoContentBlock,
WebSearchCall,
WebSearchResult,
convert_to_openai_data_block, convert_to_openai_data_block,
convert_to_openai_image_block, convert_to_openai_image_block,
is_data_content_block, is_data_content_block,
@@ -110,8 +110,6 @@ __all__ = (
"PlainTextContentBlock", "PlainTextContentBlock",
"ReasoningContentBlock", "ReasoningContentBlock",
"RemoveMessage", "RemoveMessage",
"SearchCall",
"SearchResult",
"SystemMessage", "SystemMessage",
"SystemMessageChunk", "SystemMessageChunk",
"TextContentBlock", "TextContentBlock",
@@ -121,6 +119,8 @@ __all__ = (
"ToolMessage", "ToolMessage",
"ToolMessageChunk", "ToolMessageChunk",
"VideoContentBlock", "VideoContentBlock",
"WebSearchCall",
"WebSearchResult",
"_message_from_dict", "_message_from_dict",
"convert_to_messages", "convert_to_messages",
"convert_to_openai_data_block", "convert_to_openai_data_block",
@@ -168,8 +168,8 @@ _dynamic_imports = {
"RemoveMessage": "modifier", "RemoveMessage": "modifier",
"SystemMessage": "system", "SystemMessage": "system",
"SystemMessageChunk": "system", "SystemMessageChunk": "system",
"SearchCall": "content_blocks", "WebSearchCall": "content_blocks",
"SearchResult": "content_blocks", "WebSearchResult": "content_blocks",
"ImageContentBlock": "content_blocks", "ImageContentBlock": "content_blocks",
"InvalidToolCall": "tool", "InvalidToolCall": "tool",
"TextContentBlock": "content_blocks", "TextContentBlock": "content_blocks",

View File

@@ -80,7 +80,6 @@ The module defines several types of content blocks, including:
import warnings import warnings
from typing import Any, Literal, Union from typing import Any, Literal, Union
from pydantic import TypeAdapter, ValidationError
from typing_extensions import NotRequired, TypedDict, get_args, get_origin from typing_extensions import NotRequired, TypedDict, get_args, get_origin
# --- Text and annotations --- # --- Text and annotations ---
@@ -212,10 +211,10 @@ class ToolCallContentBlock(TypedDict):
# Web search # Web search
class SearchCall(TypedDict): class WebSearchCall(TypedDict):
"""Content block for a built-in web search tool call.""" """Content block for a built-in web search tool call."""
type: Literal["search_call"] type: Literal["web_search_call"]
"""Type of the content block.""" """Type of the content block."""
id: NotRequired[str] id: NotRequired[str]
@@ -229,10 +228,10 @@ class SearchCall(TypedDict):
"""The search query used in the web search tool call.""" """The search query used in the web search tool call."""
class SearchResult(TypedDict): class WebSearchResult(TypedDict):
"""Content block for the result of a built-in search tool call.""" """Content block for the result of a built-in web search tool call."""
type: Literal["search_result"] type: Literal["web_search_result"]
"""Type of the content block.""" """Type of the content block."""
id: NotRequired[str] id: NotRequired[str]
@@ -603,8 +602,8 @@ ToolContentBlock = Union[
CodeInterpreterCall, CodeInterpreterCall,
CodeInterpreterOutput, CodeInterpreterOutput,
CodeInterpreterResult, CodeInterpreterResult,
SearchCall, WebSearchCall,
SearchResult, WebSearchResult,
] ]
ContentBlock = Union[ ContentBlock = Union[
@@ -634,9 +633,6 @@ KNOWN_BLOCK_TYPES = {
bt for bt in get_args(ContentBlock) for bt in get_args(bt.__annotations__["type"]) bt for bt in get_args(ContentBlock) for bt in get_args(bt.__annotations__["type"])
} }
# Adapter for DataContentBlock
_DataAdapter: TypeAdapter[DataContentBlock] = TypeAdapter(DataContentBlock)
def is_data_content_block(block: dict) -> bool: def is_data_content_block(block: dict) -> bool:
"""Check if the content block is a standard data content block. """Check if the content block is a standard data content block.
@@ -647,12 +643,22 @@ def is_data_content_block(block: dict) -> bool:
Returns: Returns:
True if the content block is a data content block, False otherwise. True if the content block is a data content block, False otherwise.
""" """
try: return block.get("type") in (
_DataAdapter.validate_python(block) "audio",
except ValidationError: "image",
return False "video",
else: "file",
return True "text-plain",
) and any(
key in block
for key in (
"url",
"base64",
"file_id",
"text",
"source_type", # backwards compatibility
)
)
# TODO: don't use `source_type` anymore # TODO: don't use `source_type` anymore

View File

@@ -5,21 +5,29 @@ EXPECTED_ALL = [
"_message_from_dict", "_message_from_dict",
"AIMessage", "AIMessage",
"AIMessageChunk", "AIMessageChunk",
"Annotation",
"AnyMessage", "AnyMessage",
"Base64ContentBlock", "AudioContentBlock",
"BaseMessage", "BaseMessage",
"BaseMessageChunk", "BaseMessageChunk",
"ContentBlock", "ContentBlock",
"ChatMessage", "ChatMessage",
"ChatMessageChunk", "ChatMessageChunk",
"DocumentCitation", "Citation",
"CodeInterpreterCall",
"CodeInterpreterOutput",
"CodeInterpreterResult",
"DataContentBlock",
"FileContentBlock",
"FunctionMessage", "FunctionMessage",
"FunctionMessageChunk", "FunctionMessageChunk",
"HumanMessage", "HumanMessage",
"HumanMessageChunk", "HumanMessageChunk",
"ImageContentBlock",
"InvalidToolCall", "InvalidToolCall",
"NonStandardAnnotation", "NonStandardAnnotation",
"NonStandardContentBlock", "NonStandardContentBlock",
"PlainTextContentBlock",
"SystemMessage", "SystemMessage",
"SystemMessageChunk", "SystemMessageChunk",
"TextContentBlock", "TextContentBlock",
@@ -28,7 +36,9 @@ EXPECTED_ALL = [
"ToolCallContentBlock", "ToolCallContentBlock",
"ToolMessage", "ToolMessage",
"ToolMessageChunk", "ToolMessageChunk",
"UrlCitation", "VideoContentBlock",
"WebSearchCall",
"WebSearchResult",
"ReasoningContentBlock", "ReasoningContentBlock",
"RemoveMessage", "RemoveMessage",
"convert_to_messages", "convert_to_messages",

View File

@@ -1112,23 +1112,20 @@ def test_is_data_content_block() -> None:
assert is_data_content_block( assert is_data_content_block(
{ {
"type": "image", "type": "image",
"source_type": "url",
"url": "https://...", "url": "https://...",
} }
) )
assert is_data_content_block( assert is_data_content_block(
{ {
"type": "image", "type": "image",
"source_type": "base64", "base64": "<base64 data>",
"data": "<base64 data>",
"mime_type": "image/jpeg", "mime_type": "image/jpeg",
} }
) )
assert is_data_content_block( assert is_data_content_block(
{ {
"type": "image", "type": "image",
"source_type": "base64", "base64": "<base64 data>",
"data": "<base64 data>",
"mime_type": "image/jpeg", "mime_type": "image/jpeg",
"cache_control": {"type": "ephemeral"}, "cache_control": {"type": "ephemeral"},
} }
@@ -1136,13 +1133,17 @@ def test_is_data_content_block() -> None:
assert is_data_content_block( assert is_data_content_block(
{ {
"type": "image", "type": "image",
"source_type": "base64", "base64": "<base64 data>",
"data": "<base64 data>",
"mime_type": "image/jpeg", "mime_type": "image/jpeg",
"metadata": {"cache_control": {"type": "ephemeral"}}, "metadata": {"cache_control": {"type": "ephemeral"}},
} }
) )
assert is_data_content_block(
{
"type": "image",
"source_type": "base64", # backward compatibility
}
)
assert not is_data_content_block( assert not is_data_content_block(
{ {
"type": "text", "type": "text",
@@ -1155,12 +1156,6 @@ def test_is_data_content_block() -> None:
"image_url": {"url": "https://..."}, "image_url": {"url": "https://..."},
} }
) )
assert not is_data_content_block(
{
"type": "image",
"source_type": "base64",
}
)
assert not is_data_content_block( assert not is_data_content_block(
{ {
"type": "image", "type": "image",
@@ -1203,10 +1198,17 @@ def test_convert_to_openai_image_block() -> None:
def test_known_block_types() -> None: def test_known_block_types() -> None:
assert { assert {
"text", "text",
"text-plain",
"tool_call", "tool_call",
"reasoning", "reasoning",
"non_standard", "non_standard",
"image", "image",
"audio", "audio",
"file", "file",
"video",
"code_interpreter_call",
"code_interpreter_output",
"code_interpreter_result",
"web_search_call",
"web_search_result",
} == KNOWN_BLOCK_TYPES } == KNOWN_BLOCK_TYPES