feat(langchain): add ruff rules DTZ (#32021)

See https://docs.astral.sh/ruff/rules/#flake8-datetimez-dtz

libs/langchain/langchain/chains/query_constructor/parser.py

@@ -157,7 +157,7 @@ class QueryTransformer(Transformer):
     def date(self, item: Any) -> ISO8601Date:
         item = str(item).strip("\"'")
         try:
-            datetime.datetime.strptime(item, "%Y-%m-%d")
+            datetime.datetime.strptime(item, "%Y-%m-%d")  # noqa: DTZ007
         except ValueError:
             warnings.warn(
                 "Dates are expected to be provided in ISO 8601 date format "
@@ -173,7 +173,7 @@ class QueryTransformer(Transformer):
             datetime.datetime.strptime(item, "%Y-%m-%dT%H:%M:%S%z")
         except ValueError:
             try:
-                datetime.datetime.strptime(item, "%Y-%m-%dT%H:%M:%S")
+                datetime.datetime.strptime(item, "%Y-%m-%dT%H:%M:%S")  # noqa: DTZ007
             except ValueError as e:
                 msg = "Datetime values are expected to be in ISO 8601 format."
                 raise ValueError(msg) from e

libs/langchain/langchain/output_parsers/datetime.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 from langchain_core.exceptions import OutputParserException
 from langchain_core.output_parsers import BaseOutputParser
@@ -26,7 +26,7 @@ class DatetimeOutputParser(BaseOutputParser[datetime]):
             )
         else:
             try:
-                now = datetime.now()
+                now = datetime.now(tz=timezone.utc)
                 examples = comma_list(
                     [
                         now.strftime(self.format),
@@ -48,7 +48,7 @@ class DatetimeOutputParser(BaseOutputParser[datetime]):
     def parse(self, response: str) -> datetime:
         """Parse a string into a datetime object."""
         try:
-            return datetime.strptime(response.strip(), self.format)
+            return datetime.strptime(response.strip(), self.format)  # noqa: DTZ007
         except ValueError as e:
             msg = f"Could not parse datetime string: {response}"
             raise OutputParserException(msg) from e

libs/langchain/pyproject.toml

@@ -152,6 +152,7 @@ select = [
     "COM",    # flake8-commas
     "D",      # pydocstyle
     "DOC",    # pydoclint
+    "DTZ",    # flake8-datetimez
     "E",      # pycodestyle error
     "EM",     # flake8-errmsg
     "F",      # pyflakes
@@ -217,6 +218,11 @@ pyupgrade.keep-runtime-typing = true
 "langchain/chains/constitutional_ai/principles.py" = [
     "E501", # Line too long
 ]
+"**/retrievers/*time_weighted_retriever.py" = [
+    "DTZ001", # Use of non timezone-aware datetime
+    "DTZ005", # Use of non timezone-aware datetime
+    "DTZ006", # Use of non timezone-aware datetime
+]
 
 [tool.coverage.run]
 omit = ["tests/*"]

libs/langchain/tests/unit_tests/indexes/test_indexing.py

@@ -1,5 +1,5 @@
 from collections.abc import AsyncIterator, Iterable, Iterator, Sequence
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import (
     Any,
     Optional,
@@ -166,6 +166,10 @@ def upserting_vector_store() -> InMemoryVectorStore:
     return InMemoryVectorStore(permit_upserts=True)
 
 
+_JANUARY_FIRST = datetime(2021, 1, 1, tzinfo=timezone.utc).timestamp()
+_JANUARY_SECOND = datetime(2021, 1, 2, tzinfo=timezone.utc).timestamp()
+
+
 def test_indexing_same_content(
     record_manager: SQLRecordManager,
     vector_store: InMemoryVectorStore,
@@ -256,7 +260,7 @@ def test_index_simple_delete_full(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 1).timestamp(),
+        return_value=_JANUARY_FIRST,
     ):
         assert index(loader, record_manager, vector_store, cleanup="full") == {
             "num_added": 2,
@@ -268,7 +272,7 @@ def test_index_simple_delete_full(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 1).timestamp(),
+        return_value=_JANUARY_FIRST,
     ):
         assert index(loader, record_manager, vector_store, cleanup="full") == {
             "num_added": 0,
@@ -291,7 +295,7 @@ def test_index_simple_delete_full(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(loader, record_manager, vector_store, cleanup="full") == {
             "num_added": 1,
@@ -311,7 +315,7 @@ def test_index_simple_delete_full(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(loader, record_manager, vector_store, cleanup="full") == {
             "num_added": 0,
@@ -341,7 +345,7 @@ async def test_aindex_simple_delete_full(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 1).timestamp(),
+        return_value=_JANUARY_FIRST,
     ):
         assert await aindex(loader, arecord_manager, vector_store, cleanup="full") == {
             "num_added": 2,
@@ -353,7 +357,7 @@ async def test_aindex_simple_delete_full(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 1).timestamp(),
+        return_value=_JANUARY_FIRST,
     ):
         assert await aindex(loader, arecord_manager, vector_store, cleanup="full") == {
             "num_added": 0,
@@ -376,7 +380,7 @@ async def test_aindex_simple_delete_full(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(loader, arecord_manager, vector_store, cleanup="full") == {
             "num_added": 1,
@@ -396,7 +400,7 @@ async def test_aindex_simple_delete_full(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(loader, arecord_manager, vector_store, cleanup="full") == {
             "num_added": 0,
@@ -507,7 +511,7 @@ def test_no_delete(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -526,7 +530,7 @@ def test_no_delete(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -558,7 +562,7 @@ def test_no_delete(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -596,7 +600,7 @@ async def test_ano_delete(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(
             loader,
@@ -615,7 +619,7 @@ async def test_ano_delete(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(
             loader,
@@ -647,7 +651,7 @@ async def test_ano_delete(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(
             loader,
@@ -684,7 +688,7 @@ def test_incremental_delete(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -710,7 +714,7 @@ def test_incremental_delete(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -747,7 +751,7 @@ def test_incremental_delete(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 3).timestamp(),
+        return_value=datetime(2021, 1, 3, tzinfo=timezone.utc).timestamp(),
     ):
         assert index(
             loader,
@@ -803,7 +807,7 @@ def test_incremental_indexing_with_batch_size(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -870,7 +874,7 @@ def test_incremental_delete_with_batch_size(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -897,7 +901,7 @@ def test_incremental_delete_with_batch_size(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert index(
             loader,
@@ -917,7 +921,7 @@ def test_incremental_delete_with_batch_size(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2022, 1, 3).timestamp(),
+        return_value=datetime(2022, 1, 3, tzinfo=timezone.utc).timestamp(),
     ):
         # Docs with same content
         docs = [
@@ -948,7 +952,7 @@ def test_incremental_delete_with_batch_size(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2023, 1, 3).timestamp(),
+        return_value=datetime(2023, 1, 3, tzinfo=timezone.utc).timestamp(),
     ):
         # Docs with same content
         docs = [
@@ -979,7 +983,7 @@ def test_incremental_delete_with_batch_size(
     with patch.object(
         record_manager,
         "get_time",
-        return_value=datetime(2024, 1, 3).timestamp(),
+        return_value=datetime(2024, 1, 3, tzinfo=timezone.utc).timestamp(),
     ):
         # Docs with same content
         docs = [
@@ -1028,7 +1032,7 @@ async def test_aincremental_delete(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(
             loader.lazy_load(),
@@ -1054,7 +1058,7 @@ async def test_aincremental_delete(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 2).timestamp(),
+        return_value=_JANUARY_SECOND,
     ):
         assert await aindex(
             loader.lazy_load(),
@@ -1091,7 +1095,7 @@ async def test_aincremental_delete(
     with patch.object(
         arecord_manager,
         "aget_time",
-        return_value=datetime(2021, 1, 3).timestamp(),
+        return_value=datetime(2021, 1, 3, tzinfo=timezone.utc).timestamp(),
     ):
         assert await aindex(
             loader.lazy_load(),

libs/langchain/tests/unit_tests/output_parsers/test_datetime_parser.py

@@ -10,14 +10,13 @@ def test_datetime_output_parser_parse() -> None:
     parser = DatetimeOutputParser()
 
     # Test valid input
-    date = datetime.now()
+    date = datetime.now()  # noqa: DTZ005
     datestr = date.strftime(parser.format)
     result = parser.parse(datestr)
     assert result == date
 
     # Test valid input
     parser.format = "%Y-%m-%dT%H:%M:%S"
-    date = datetime.now()
     datestr = date.strftime(parser.format)
     result = parser.parse(datestr)
     assert (
@@ -31,7 +30,6 @@ def test_datetime_output_parser_parse() -> None:
 
     # Test valid input
     parser.format = "%H:%M:%S"
-    date = datetime.now()
     datestr = date.strftime(parser.format)
     result = parser.parse(datestr)
     assert (

libs/langchain/tests/unit_tests/output_parsers/test_fix.py

@@ -1,4 +1,5 @@
 from datetime import datetime as dt
+from datetime import timezone
 from typing import Any, Callable, Optional, TypeVar
 
 import pytest
@@ -151,17 +152,17 @@ def test_output_fixing_parser_output_type(
     [
         (
             "2024/07/08",
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             NAIVE_FIX_PROMPT | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
         (
             # Case: retry_chain.InputType does not have 'instructions' key
             "2024/07/08",
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             PromptTemplate.from_template("{completion}\n{error}")
             | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
     ],
 )
@@ -188,17 +189,17 @@ def test_output_fixing_parser_parse_with_retry_chain(
     [
         (
             "2024/07/08",
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             NAIVE_FIX_PROMPT | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
         (
             # Case: retry_chain.InputType does not have 'instructions' key
             "2024/07/08",
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             PromptTemplate.from_template("{completion}\n{error}")
             | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
     ],
 )

libs/langchain/tests/unit_tests/output_parsers/test_retry.py

@@ -1,4 +1,5 @@
 from datetime import datetime as dt
+from datetime import timezone
 from typing import Any, Callable, Optional, TypeVar
 
 import pytest
@@ -209,10 +210,10 @@ def test_retry_with_error_output_parser_parse_is_not_implemented() -> None:
         (
             "2024/07/08",
             StringPromptValue(text="dummy"),
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             NAIVE_RETRY_PROMPT
             | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
     ],
 )
@@ -237,10 +238,10 @@ def test_retry_output_parser_parse_with_prompt_with_retry_chain(
         (
             "2024/07/08",
             StringPromptValue(text="dummy"),
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             NAIVE_RETRY_PROMPT
             | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
     ],
 )
@@ -266,10 +267,10 @@ async def test_retry_output_parser_aparse_with_prompt_with_retry_chain(
         (
             "2024/07/08",
             StringPromptValue(text="dummy"),
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             NAIVE_RETRY_WITH_ERROR_PROMPT
             | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
     ],
 )
@@ -295,10 +296,10 @@ def test_retry_with_error_output_parser_parse_with_prompt_with_retry_chain(
         (
             "2024/07/08",
             StringPromptValue(text="dummy"),
-            DatetimeOutputParser(),
+            DatetimeOutputParser(format="%Y-%m-%dT%H:%M:%S.%f%z"),
             NAIVE_RETRY_WITH_ERROR_PROMPT
             | RunnableLambda(lambda _: "2024-07-08T00:00:00.000000Z"),
-            dt(2024, 7, 8),
+            dt(2024, 7, 8, tzinfo=timezone.utc),
         ),
     ],
 )

libs/langchain/tests/unit_tests/smith/evaluation/test_runner_utils.py

@@ -2,7 +2,7 @@
 
 import uuid
 from collections.abc import Iterator
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Optional, Union
 from unittest import mock
 
@@ -27,7 +27,7 @@ from langchain.smith.evaluation.runner_utils import (
 from tests.unit_tests.llms.fake_chat_model import FakeChatModel
 from tests.unit_tests.llms.fake_llm import FakeLLM
 
-_CREATED_AT = datetime(2015, 1, 1, 0, 0, 0)
+_CREATED_AT = datetime(2015, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
 _TENANT_ID = "7a3d2b56-cd5b-44e5-846f-7eb6e8144ce4"
 _EXAMPLE_MESSAGE = {
     "data": {"content": "Foo", "example": False, "additional_kwargs": {}},