core: Fix some missing generic types

Author: cbornet

libs/core/langchain_core/document_loaders/langsmith.py

@@ -125,7 +125,7 @@ def lazy_load(self) -> Iterator[Document]:
             yield Document(content_str, metadata=metadata)
 
 
-def _stringify(x: Union[str, dict]) -> str:
+def _stringify(x: Union[str, dict[str, Any]]) -> str:
     if isinstance(x, str):
         return x
     try:

libs/core/langchain_core/messages/utils.py

@@ -202,13 +202,17 @@ def message_chunk_to_message(chunk: BaseMessageChunk) -> BaseMessage:
 
 
 MessageLikeRepresentation = Union[
-    BaseMessage, list[str], tuple[str, str], str, dict[str, Any]
+    BaseMessage,
+    list[str],
+    tuple[str, Union[str, list[Union[str, dict[str, Any]]]]],
+    str,
+    dict[str, Any],
 ]
 
 
 def _create_message_from_message_type(
     message_type: str,
-    content: str,
+    content: Union[str, list[Union[str, dict[str, Any]]]],
     name: Optional[str] = None,
     tool_call_id: Optional[str] = None,
     tool_calls: Optional[list[dict[str, Any]]] = None,
@@ -218,13 +222,13 @@ def _create_message_from_message_type(
     """Create a message from a message type and content string.
 
     Args:
-        message_type: (str) the type of the message (e.g., "human", "ai", etc.).
-        content: (str) the content string.
-        name: (str) the name of the message. Default is None.
-        tool_call_id: (str) the tool call id. Default is None.
-        tool_calls: (list[dict[str, Any]]) the tool calls. Default is None.
-        id: (str) the id of the message. Default is None.
-        additional_kwargs: (dict[str, Any]) additional keyword arguments.
+        message_type: the type of the message (e.g., "human", "ai", etc.).
+        content: the content string.
+        name: the name of the message. Default is None.
+        tool_call_id: the tool call id. Default is None.
+        tool_calls: the tool calls. Default is None.
+        id: the id of the message. Default is None.
+        **additional_kwargs: additional keyword arguments.
 
     Returns:
         a message of the appropriate type.

libs/core/langchain_core/output_parsers/openai_functions.py

@@ -224,7 +224,7 @@ class Dog(BaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def validate_schema(cls, values: dict) -> Any:
+    def validate_schema(cls, values: dict[str, Any]) -> Any:
         """Validate the pydantic schema.
 
         Args:

libs/core/langchain_core/prompts/chat.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from pathlib import Path
 from typing import (
-    TYPE_CHECKING,
     Annotated,
     Any,
     Optional,
@@ -51,9 +51,6 @@
 from langchain_core.utils import get_colored_text
 from langchain_core.utils.interactive_env import is_interactive_env
 
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-
 
 class MessagesPlaceholder(BaseMessagePromptTemplate):
     """Prompt template that assumes variable is already list of messages.
@@ -771,7 +768,7 @@ def pretty_print(self) -> None:
     MessageLike,
     tuple[
         Union[str, type],
-        Union[str, list[dict], list[object]],
+        Union[str, Sequence[dict], Sequence[object]],
     ],
     str,
     dict[str, Any],

libs/core/langchain_core/tracers/evaluation.py

@@ -100,7 +100,7 @@ def __init__(
             )
         else:
             self.executor = None
-        self.futures: weakref.WeakSet[Future] = weakref.WeakSet()
+        self.futures: weakref.WeakSet[Future[None]] = weakref.WeakSet()
         self.skip_unfinished = skip_unfinished
         self.project_name = project_name
         self.logged_eval_results: dict[tuple[str, str], list[EvaluationResult]] = {}

libs/core/tests/unit_tests/example_selectors/test_base.py

@@ -10,7 +10,7 @@ def __init__(self) -> None:
     def add_example(self, example: dict[str, str]) -> None:
         self.example = example
 
-    def select_examples(self, input_variables: dict[str, str]) -> list[dict]:
+    def select_examples(self, input_variables: dict[str, str]) -> list[dict[str, str]]:
         return [input_variables]
 
 

libs/core/tests/unit_tests/fake/callbacks.py

@@ -276,7 +276,9 @@ def on_retriever_error(
         self.on_retriever_error_common()
 
     # Overriding since BaseModel has __deepcopy__ method as well
-    def __deepcopy__(self, memo: dict) -> "FakeCallbackHandler":  # type: ignore[override]
+    def __deepcopy__(
+        self, memo: Union[dict[int, Any], None] = None
+    ) -> "FakeCallbackHandler":
         return self
 
 
@@ -426,5 +428,7 @@ async def on_text(
         self.on_text_common()
 
     # Overriding since BaseModel has __deepcopy__ method as well
-    def __deepcopy__(self, memo: dict) -> "FakeAsyncCallbackHandler":  # type: ignore[override]
+    def __deepcopy__(
+        self, memo: Union[dict[int, Any], None] = None
+    ) -> "FakeAsyncCallbackHandler":
         return self

libs/core/tests/unit_tests/language_models/chat_models/test_base.py

@@ -41,22 +41,22 @@
 
 
 @pytest.fixture
-def messages() -> list:
+def messages() -> list[BaseMessage]:
     return [
         SystemMessage(content="You are a test user."),
         HumanMessage(content="Hello, I am a test user."),
     ]
 
 
 @pytest.fixture
-def messages_2() -> list:
+def messages_2() -> list[BaseMessage]:
     return [
         SystemMessage(content="You are a test user."),
         HumanMessage(content="Hello, I not a test user."),
     ]
 
 
-def test_batch_size(messages: list, messages_2: list) -> None:
+def test_batch_size(messages: list[BaseMessage], messages_2: list[BaseMessage]) -> None:
     # The base endpoint doesn't support native batching,
     # so we expect batch_size to always be 1
     llm = FakeListChatModel(responses=[str(i) for i in range(100)])
@@ -80,7 +80,9 @@ def test_batch_size(messages: list, messages_2: list) -> None:
         assert (cb.traced_runs[0].extra or {}).get("batch_size") == 1
 
 
-async def test_async_batch_size(messages: list, messages_2: list) -> None:
+async def test_async_batch_size(
+    messages: list[BaseMessage], messages_2: list[BaseMessage]
+) -> None:
     llm = FakeListChatModel(responses=[str(i) for i in range(100)])
     # The base endpoint doesn't support native batching,
     # so we expect batch_size to always be 1
@@ -262,7 +264,7 @@ def _llm_type(self) -> str:
 class FakeTracer(BaseTracer):
     def __init__(self) -> None:
         super().__init__()
-        self.traced_run_ids: list = []
+        self.traced_run_ids: list[uuid.UUID] = []
 
     def _persist_run(self, run: Run) -> None:
         """Persist a run."""
@@ -411,7 +413,7 @@ async def test_disable_streaming_no_streaming_model_async(
 class FakeChatModelStartTracer(FakeTracer):
     def __init__(self) -> None:
         super().__init__()
-        self.messages: list = []
+        self.messages: list[list[list[BaseMessage]]] = []
 
     def on_chat_model_start(self, *args: Any, **kwargs: Any) -> Run:
         _, messages = args

libs/core/tests/unit_tests/messages/test_utils.py

@@ -17,6 +17,7 @@
     ToolMessage,
 )
 from langchain_core.messages.utils import (
+    MessageLikeRepresentation,
     convert_to_messages,
     convert_to_openai_messages,
     count_tokens_approximately,
@@ -152,7 +153,7 @@ def test_merge_messages_tool_messages() -> None:
         {"include_names": ["blah", "blur"], "exclude_types": [SystemMessage]},
     ],
 )
-def test_filter_message(filters: dict) -> None:
+def test_filter_message(filters: dict[str, Any]) -> None:
     messages = [
         SystemMessage("foo", name="blah", id="1"),
         HumanMessage("bar", name="blur", id="2"),
@@ -672,7 +673,7 @@ def get_num_tokens_from_messages(
 
 
 def test_convert_to_messages() -> None:
-    message_like: list = [
+    message_like: list[MessageLikeRepresentation] = [
         # BaseMessage
         SystemMessage("1"),
         SystemMessage("1.1", additional_kwargs={"__openai_role__": "developer"}),
@@ -1178,7 +1179,7 @@ def test_convert_to_openai_messages_mixed_content_types() -> None:
 
 
 def test_convert_to_openai_messages_developer() -> None:
-    messages: list = [
+    messages: list[MessageLikeRepresentation] = [
         SystemMessage("a", additional_kwargs={"__openai_role__": "developer"}),
         {"role": "developer", "content": "a"},
     ]

libs/core/tests/unit_tests/output_parsers/test_openai_tools.py

@@ -21,7 +21,7 @@
     IS_PYDANTIC_V2,
 )
 
-STREAMED_MESSAGES: list = [
+STREAMED_MESSAGES: list[AIMessageChunk] = [
     AIMessageChunk(content=""),
     AIMessageChunk(
         content="",
@@ -335,7 +335,7 @@
         STREAMED_MESSAGES_WITH_TOOL_CALLS.append(message)
 
 
-EXPECTED_STREAMED_JSON = [
+EXPECTED_STREAMED_JSON: list[dict[str, Any]] = [
     {},
     {"names": ["suz"]},
     {"names": ["suzy"]},
@@ -396,7 +396,7 @@ def test_partial_json_output_parser() -> None:
         chain = input_iter | JsonOutputToolsParser()
 
         actual = list(chain.stream(None))
-        expected: list = [[]] + [
+        expected: list[list[dict[str, Any]]] = [[]] + [
             [{"type": "NameCollector", "args": chunk}]
             for chunk in EXPECTED_STREAMED_JSON
         ]
@@ -409,7 +409,7 @@ async def test_partial_json_output_parser_async() -> None:
         chain = input_iter | JsonOutputToolsParser()
 
         actual = [p async for p in chain.astream(None)]
-        expected: list = [[]] + [
+        expected: list[list[dict[str, Any]]] = [[]] + [
             [{"type": "NameCollector", "args": chunk}]
             for chunk in EXPECTED_STREAMED_JSON
         ]
@@ -422,7 +422,7 @@ def test_partial_json_output_parser_return_id() -> None:
         chain = input_iter | JsonOutputToolsParser(return_id=True)
 
         actual = list(chain.stream(None))
-        expected: list = [[]] + [
+        expected: list[list[dict[str, Any]]] = [[]] + [
             [
                 {
                     "type": "NameCollector",
@@ -441,7 +441,9 @@ def test_partial_json_output_key_parser() -> None:
         chain = input_iter | JsonOutputKeyToolsParser(key_name="NameCollector")
 
         actual = list(chain.stream(None))
-        expected: list = [[]] + [[chunk] for chunk in EXPECTED_STREAMED_JSON]
+        expected: list[list[dict[str, Any]]] = [[]] + [
+            [chunk] for chunk in EXPECTED_STREAMED_JSON
+        ]
         assert actual == expected
 
 
@@ -452,7 +454,9 @@ async def test_partial_json_output_parser_key_async() -> None:
         chain = input_iter | JsonOutputKeyToolsParser(key_name="NameCollector")
 
         actual = [p async for p in chain.astream(None)]
-        expected: list = [[]] + [[chunk] for chunk in EXPECTED_STREAMED_JSON]
+        expected: list[list[dict[str, Any]]] = [
+            [chunk] for chunk in EXPECTED_STREAMED_JSON
+        ]
         assert actual == expected
 
 

libs/core/tests/unit_tests/output_parsers/test_pydantic_parser.py

@@ -141,9 +141,7 @@ class TestModel(BaseModel):
 
 def test_pydantic_output_parser() -> None:
     """Test PydanticOutputParser."""
-    pydantic_parser: PydanticOutputParser = PydanticOutputParser(
-        pydantic_object=TestModel
-    )
+    pydantic_parser = PydanticOutputParser(pydantic_object=TestModel)
 
     result = pydantic_parser.parse(DEF_RESULT)
     assert result == DEF_EXPECTED_RESULT
@@ -152,9 +150,7 @@ def test_pydantic_output_parser() -> None:
 
 def test_pydantic_output_parser_fail() -> None:
     """Test PydanticOutputParser where completion result fails schema validation."""
-    pydantic_parser: PydanticOutputParser = PydanticOutputParser(
-        pydantic_object=TestModel
-    )
+    pydantic_parser = PydanticOutputParser(pydantic_object=TestModel)
 
     with pytest.raises(
         OutputParserException, match="Failed to parse TestModel from completion"

libs/core/tests/unit_tests/outputs/test_chat_generation.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Any, Union
 
 import pytest
 
@@ -19,14 +19,14 @@
         ],
     ],
 )
-def test_msg_with_text(content: Union[str, list]) -> None:
+def test_msg_with_text(content: Union[str, list[Union[str, dict[str, Any]]]]) -> None:
     expected = "foo"
     actual = ChatGeneration(message=AIMessage(content=content)).text
     assert actual == expected
 
 
 @pytest.mark.parametrize("content", [[], [{"tool_use": {}, "type": "tool_use"}]])
-def test_msg_no_text(content: Union[str, list]) -> None:
+def test_msg_no_text(content: Union[str, list[Union[str, dict[str, Any]]]]) -> None:
     expected = ""
     actual = ChatGeneration(message=AIMessage(content=content)).text
     assert actual == expected