core: Add mypy strict-equality rule (#31286)

Author: cbornet

libs/core/langchain_core/language_models/chat_models.py

@@ -51,7 +51,6 @@
     AIMessage,
     AnyMessage,
     BaseMessage,
-    BaseMessageChunk,
     HumanMessage,
     convert_to_messages,
     convert_to_openai_image_block,
@@ -446,13 +445,10 @@ def stream(
         *,
         stop: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> Iterator[BaseMessageChunk]:
+    ) -> Iterator[BaseMessage]:
         if not self._should_stream(async_api=False, **{**kwargs, "stream": True}):
             # model doesn't implement streaming, so use default implementation
-            yield cast(
-                "BaseMessageChunk",
-                self.invoke(input, config=config, stop=stop, **kwargs),
-            )
+            yield self.invoke(input, config=config, stop=stop, **kwargs)
         else:
             config = ensure_config(config)
             messages = self._convert_input(input).to_messages()
@@ -537,13 +533,10 @@ async def astream(
         *,
         stop: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> AsyncIterator[BaseMessageChunk]:
+    ) -> AsyncIterator[BaseMessage]:
         if not self._should_stream(async_api=True, **{**kwargs, "stream": True}):
             # No async or sync stream is implemented, so fall back to ainvoke
-            yield cast(
-                "BaseMessageChunk",
-                await self.ainvoke(input, config=config, stop=stop, **kwargs),
-            )
+            yield await self.ainvoke(input, config=config, stop=stop, **kwargs)
             return
 
         config = ensure_config(config)
@@ -1454,7 +1447,7 @@ class AnswerWithJustification(BaseModel):
             PydanticToolsParser,
         )
 
-        if self.bind_tools is BaseChatModel.bind_tools:
+        if type(self).bind_tools is BaseChatModel.bind_tools:
             msg = "with_structured_output is not implemented for this model."
             raise NotImplementedError(msg)
 

libs/core/langchain_core/runnables/base.py

@@ -4331,8 +4331,9 @@ def __init__(
         self,
         func: Union[
             Union[
-                Callable[[Input], Output],
                 Callable[[Input], Iterator[Output]],
+                Callable[[Input], Runnable[Input, Output]],
+                Callable[[Input], Output],
                 Callable[[Input, RunnableConfig], Output],
                 Callable[[Input, CallbackManagerForChainRun], Output],
                 Callable[[Input, CallbackManagerForChainRun, RunnableConfig], Output],

libs/core/langchain_core/utils/pydantic.py

@@ -77,6 +77,7 @@ def get_pydantic_major_version() -> int:
     TypeBaseModel = type[BaseModel]
 elif IS_PYDANTIC_V2:
     from pydantic.v1.fields import FieldInfo as FieldInfoV1  # type: ignore[assignment]
+    from pydantic.v1.fields import ModelField
 
     # Union type needs to be last assignment to PydanticBaseModel to make mypy happy.
     PydanticBaseModel = Union[BaseModel, pydantic.BaseModel]  # type: ignore[assignment,misc]
@@ -373,20 +374,20 @@ def get_fields(model: type[BaseModelV2]) -> dict[str, FieldInfoV2]: ...
     def get_fields(model: BaseModelV2) -> dict[str, FieldInfoV2]: ...
 
     @overload
-    def get_fields(model: type[BaseModelV1]) -> dict[str, FieldInfoV1]: ...
+    def get_fields(model: type[BaseModelV1]) -> dict[str, ModelField]: ...
 
     @overload
-    def get_fields(model: BaseModelV1) -> dict[str, FieldInfoV1]: ...
+    def get_fields(model: BaseModelV1) -> dict[str, ModelField]: ...
 
     def get_fields(
         model: Union[type[Union[BaseModelV2, BaseModelV1]], BaseModelV2, BaseModelV1],
-    ) -> Union[dict[str, FieldInfoV2], dict[str, FieldInfoV1]]:
+    ) -> Union[dict[str, FieldInfoV2], dict[str, ModelField]]:
         """Get the field names of a Pydantic model."""
         if hasattr(model, "model_fields"):
             return model.model_fields
 
         if hasattr(model, "__fields__"):
-            return model.__fields__  # type: ignore[return-value]
+            return model.__fields__
         msg = f"Expected a Pydantic model. Got {type(model)}"
         raise TypeError(msg)
 

libs/core/pyproject.toml

@@ -74,7 +74,6 @@ report_deprecated_as_note = "True"
 # TODO: activate for 'strict' checking
 disallow_any_generics = "False"
 warn_return_any = "False"
-strict_equality = "False"
 
 
 [tool.ruff]

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

@@ -1,5 +1,7 @@
 """Tests for verifying that testing utility code works as expected."""
 
+import operator
+from functools import reduce
 from itertools import cycle
 from typing import Any, Optional, Union
 from uuid import UUID
@@ -115,12 +117,7 @@ async def test_generic_fake_chat_model_stream() -> None:
     ]
     assert len({chunk.id for chunk in chunks}) == 1
 
-    accumulate_chunks = None
-    for chunk in chunks:
-        if accumulate_chunks is None:
-            accumulate_chunks = chunk
-        else:
-            accumulate_chunks += chunk
+    accumulate_chunks = reduce(operator.add, chunks)
 
     assert accumulate_chunks == AIMessageChunk(
         content="",

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

@@ -1,7 +1,7 @@
 """Test PydanticOutputParser."""
 
 from enum import Enum
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 
 import pydantic
 import pytest
@@ -30,7 +30,7 @@ class ForecastV1(V1BaseModel):
 
 @pytest.mark.parametrize("pydantic_object", [ForecastV2, ForecastV1])
 def test_pydantic_parser_chaining(
-    pydantic_object: TBaseModel,
+    pydantic_object: Union[type[ForecastV2], type[ForecastV1]],
 ) -> None:
     prompt = PromptTemplate(
         template="""{{
@@ -43,11 +43,11 @@ def test_pydantic_parser_chaining(
 
     model = ParrotFakeChatModel()
 
-    parser = PydanticOutputParser(pydantic_object=pydantic_object)  # type: ignore[arg-type,var-annotated]
+    parser = PydanticOutputParser(pydantic_object=pydantic_object)  # type: ignore[type-var]
     chain = prompt | model | parser
 
     res = chain.invoke({})
-    assert type(res) is pydantic_object
+    assert isinstance(res, pydantic_object)
     assert res.f_or_c == "C"
     assert res.temperature == 20
     assert res.forecast == "Sunny"

libs/core/tests/unit_tests/prompts/test_prompt.py

@@ -441,7 +441,7 @@ def test_basic_sandboxing_with_jinja2() -> None:
     template = " {{''.__class__.__bases__[0] }} "  # malicious code
     prompt = PromptTemplate.from_template(template, template_format="jinja2")
     with pytest.raises(jinja2.exceptions.SecurityError):
-        assert prompt.format() == []
+        prompt.format()
 
 
 @pytest.mark.requires("jinja2")

libs/core/tests/unit_tests/prompts/test_structured.py

@@ -51,7 +51,7 @@ class OutputSchema(BaseModel):
 
     chain = prompt | model
 
-    assert chain.invoke({"hello": "there"}) == OutputSchema(name="yo", value=42)
+    assert chain.invoke({"hello": "there"}) == OutputSchema(name="yo", value=42)  # type: ignore[comparison-overlap]
 
 
 def test_structured_prompt_dict() -> None:
@@ -73,13 +73,13 @@ def test_structured_prompt_dict() -> None:
 
     chain = prompt | model
 
-    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 42}
+    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 42}  # type: ignore[comparison-overlap]
 
     assert loads(dumps(prompt)).model_dump() == prompt.model_dump()
 
     chain = loads(dumps(prompt)) | model
 
-    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 42}
+    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 42}  # type: ignore[comparison-overlap]
 
 
 def test_structured_prompt_kwargs() -> None:
@@ -99,10 +99,10 @@ def test_structured_prompt_kwargs() -> None:
     )
     model = FakeStructuredChatModel(responses=[])
     chain = prompt | model
-    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 7}
+    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 7}  # type: ignore[comparison-overlap]
     assert loads(dumps(prompt)).model_dump() == prompt.model_dump()
     chain = loads(dumps(prompt)) | model
-    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 7}
+    assert chain.invoke({"hello": "there"}) == {"name": 1, "value": 7}  # type: ignore[comparison-overlap]
 
     class OutputSchema(BaseModel):
         name: str
@@ -116,7 +116,7 @@ class OutputSchema(BaseModel):
 
     chain = prompt | model
 
-    assert chain.invoke({"hello": "there"}) == OutputSchema(name="yo", value=7)
+    assert chain.invoke({"hello": "there"}) == OutputSchema(name="yo", value=7)  # type: ignore[comparison-overlap]
 
 
 def test_structured_prompt_template_format() -> None:

libs/core/tests/unit_tests/runnables/test_runnable.py

@@ -671,7 +671,7 @@ def foo(x: int) -> None:
 
 def test_schema_with_itemgetter() -> None:
     """Test runnable with itemgetter."""
-    foo = RunnableLambda(itemgetter("hello"))
+    foo: Runnable = RunnableLambda(itemgetter("hello"))
     assert _schema(foo.input_schema) == {
         "properties": {"hello": {"title": "Hello"}},
         "required": ["hello"],
@@ -4001,7 +4001,7 @@ def test_runnable_lambda_stream() -> None:
     # sleep to better simulate a real stream
     llm = FakeStreamingListLLM(responses=[llm_res], sleep=0.01)
 
-    output = list(RunnableLambda(lambda _: llm).stream(""))
+    output = list(RunnableLambda[str, str](lambda _: llm).stream(""))
     assert output == list(llm_res)
 
 
@@ -4014,9 +4014,9 @@ def test_runnable_lambda_stream_with_callbacks() -> None:
     llm = FakeStreamingListLLM(responses=[llm_res], sleep=0.01)
     config: RunnableConfig = {"callbacks": [tracer]}
 
-    assert list(RunnableLambda(lambda _: llm).stream("", config=config)) == list(
-        llm_res
-    )
+    assert list(
+        RunnableLambda[str, str](lambda _: llm).stream("", config=config)
+    ) == list(llm_res)
 
     assert len(tracer.runs) == 1
     assert tracer.runs[0].error is None
@@ -4075,10 +4075,7 @@ async def afunc(*args: Any, **kwargs: Any) -> Any:
     assert output == list(llm_res)
 
     output = [
-        chunk
-        async for chunk in cast(
-            "AsyncIterator[str]", RunnableLambda(lambda _: llm).astream("")
-        )
+        chunk async for chunk in RunnableLambda[str, str](lambda _: llm).astream("")
     ]
     assert output == list(llm_res)
 
@@ -4093,7 +4090,10 @@ async def test_runnable_lambda_astream_with_callbacks() -> None:
     config: RunnableConfig = {"callbacks": [tracer]}
 
     assert [
-        _ async for _ in RunnableLambda(lambda _: llm).astream("", config=config)
+        _
+        async for _ in RunnableLambda[str, str](lambda _: llm).astream(
+            "", config=config
+        )
     ] == list(llm_res)
 
     assert len(tracer.runs) == 1
@@ -5300,7 +5300,7 @@ async def idchain_async(_input: dict, /) -> bool:
     def func(_input: dict, /) -> Runnable:
         return idchain
 
-    assert await RunnableLambda(func).ainvoke({})
+    assert await RunnableLambda[dict, bool](func).ainvoke({})
 
 
 def test_invoke_stream_passthrough_assign_trace() -> None:

libs/core/tests/unit_tests/test_tools.py

@@ -210,7 +210,7 @@ def tool_func_v1(*, arg1: int, arg2: bool, arg3: Optional[dict] = None) -> str:
         return f"{arg1} {arg2} {arg3}"
 
     assert isinstance(tool_func_v1, BaseTool)
-    assert tool_func_v1.args_schema == _MockSchemaV1
+    assert tool_func_v1.args_schema == cast("ArgsSchema", _MockSchemaV1)
 
 
 def test_decorated_function_schema_equivalent() -> None: