diff --git a/libs/core/langchain_core/messages/content_blocks.py b/libs/core/langchain_core/messages/content_blocks.py index c0f6b2421d6..7bb2d9a1fe5 100644 --- a/libs/core/langchain_core/messages/content_blocks.py +++ b/libs/core/langchain_core/messages/content_blocks.py @@ -857,19 +857,34 @@ ContentBlock = Union[ def _extract_typedict_type_values(union_type: Any) -> set[str]: """Extract the values of the ``'type'`` field from a TypedDict union type.""" result: set[str] = set() + + def _extract_from_type(type_arg: Any) -> None: + # If it's a Union, recursively extract from its args + if get_origin(type_arg) is Union: + for sub_type in get_args(type_arg): + _extract_from_type(sub_type) + # If it's a TypedDict with a type annotation, extract the literal values + elif ( + hasattr(type_arg, "__annotations__") and "type" in type_arg.__annotations__ + ): + annotation = type_arg.__annotations__["type"] + # Handle NotRequired[Literal[...]] annotations + if get_origin(annotation) is not None and "NotRequired" in str( + get_origin(annotation) + ): + inner_annotation = get_args(annotation)[0] + if get_origin(inner_annotation) is Literal: + result.update(get_args(inner_annotation)) + elif get_origin(annotation) is Literal: + result.update(get_args(annotation)) + for value in get_args(union_type): - annotation = value.__annotations__["type"] - if get_origin(annotation) is Literal: - result.update(get_args(annotation)) - else: - msg = f"{value} 'type' is not a Literal" - raise ValueError(msg) + _extract_from_type(value) + return result -KNOWN_BLOCK_TYPES = { - bt for bt in get_args(ContentBlock) for bt in get_args(bt.__annotations__["type"]) -} +KNOWN_BLOCK_TYPES = _extract_typedict_type_values(ContentBlock) def is_data_content_block(block: dict) -> bool: diff --git a/libs/core/tests/unit_tests/test_messages.py b/libs/core/tests/unit_tests/test_messages.py index 77ddf96c974..13d86eefa25 100644 --- a/libs/core/tests/unit_tests/test_messages.py +++ b/libs/core/tests/unit_tests/test_messages.py @@ -1364,19 +1364,22 @@ def test_convert_to_openai_image_block() -> None: def test_known_block_types() -> None: assert { - "text", - "text-plain", - "tool_call", - "invalid_tool_call", - "reasoning", - "non_standard", - "image", "audio", - "file", - "video", + "citation", "code_interpreter_call", "code_interpreter_output", "code_interpreter_result", + "file", + "image", + "invalid_tool_call", + "non_standard", + "non_standard_annotation", + "reasoning", + "text", + "text-plain", + "tool_call", + "tool_call_chunk", + "video", "web_search_call", "web_search_result", } == KNOWN_BLOCK_TYPES