fix test

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:

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