CrateDB: Format code. Satisfy linter and type checker. ruff + mypy

Author: amotl

libs/community/langchain_community/chat_message_histories/cratedb.py

@@ -2,8 +2,8 @@
 import typing as t
 
 import sqlalchemy as sa
-from sqlalchemy_cratedb.support import refresh_after_dml, refresh_table
 from langchain.schema import BaseMessage, _message_to_dict, messages_from_dict
+from sqlalchemy_cratedb.support import refresh_after_dml, refresh_table
 
 from langchain_community.chat_message_histories.sql import (
     BaseMessageConverter,

libs/community/langchain_community/vectorstores/cratedb/base.py

@@ -19,9 +19,9 @@
 from langchain.utils import get_from_dict_or_env
 from langchain.vectorstores.pgvector import PGVector
 from sqlalchemy.orm import sessionmaker
+from sqlalchemy_cratedb.support import refresh_after_dml, refresh_table
 
 from langchain_community.vectorstores.cratedb.model import ModelFactory
-from sqlalchemy_cratedb.support import refresh_after_dml, refresh_table
 
 
 class DistanceStrategy(str, enum.Enum):
@@ -87,7 +87,7 @@ def __post_init__(
         """
 
         self._engine = self._bind
-        self.Session = sessionmaker(self._engine)
+        self.Session = sessionmaker(bind=self._engine)  # type: ignore[call-overload]
 
         # Patch dialect to invoke `REFRESH TABLE` after each DML operation.
         refresh_after_dml(self._engine)
@@ -199,6 +199,7 @@ def drop_tables(self) -> None:
     def delete(
         self,
         ids: Optional[List[str]] = None,
+        collection_only: bool = False,
         **kwargs: Any,
     ) -> None:
         """
@@ -214,7 +215,7 @@ def delete(
                 patch, listening to `after_delete` events seems not be
                 able to catch it.
         """
-        super().delete(ids=ids, **kwargs)
+        super().delete(ids=ids, collection_only=collection_only, **kwargs)
 
         # CrateDB: Synchronize data because `on_flush` does not catch it.
         with self.Session() as session:

libs/community/langchain_community/vectorstores/cratedb/extended.py

@@ -8,7 +8,6 @@
 )
 
 import sqlalchemy
-
 from langchain.schema.embeddings import Embeddings
 
 from langchain_community.vectorstores.cratedb.base import (

libs/community/langchain_community/vectorstores/cratedb/model.py

@@ -2,8 +2,8 @@
 from typing import Any, List, Optional, Tuple
 
 import sqlalchemy
-from sqlalchemy_cratedb import ObjectType, FloatVector
 from sqlalchemy.orm import Session, declarative_base, relationship
+from sqlalchemy_cratedb import FloatVector, ObjectType
 
 
 def generate_uuid() -> str:
@@ -48,7 +48,7 @@ class CollectionStore(BaseModel):
             )
 
             @classmethod
-            def get_by_name(cls, session: Session, name: str) -> "CollectionStore":
+            def get_by_name(cls, session: Session, name: str) -> Optional["CollectionStore"]:
                 return session.query(cls).filter(cls.name == name).first()  # type: ignore[attr-defined]
 
             @classmethod
@@ -95,8 +95,8 @@ class EmbeddingStore(BaseModel):
             )
             collection = relationship("CollectionStore", back_populates="embeddings")
 
-            embedding = sqlalchemy.Column(FloatVector(self.dimensions))
-            document = sqlalchemy.Column(sqlalchemy.String, nullable=True)
+            embedding: sqlalchemy.Column = sqlalchemy.Column(FloatVector(self.dimensions))
+            document: sqlalchemy.Column = sqlalchemy.Column(sqlalchemy.String, nullable=True)
             cmetadata: sqlalchemy.Column = sqlalchemy.Column(ObjectType, nullable=True)
 
             # custom_id : any user defined id

libs/community/tests/integration_tests/document_loaders/test_sql_database.py

@@ -122,9 +122,7 @@ def pytest_generate_tests(metafunc: "Metafunc") -> None:
             # use `docker compose up postgres` to start the instance
             # it will have the appropriate credentials set up including
             # being exposed on the appropriate port.
-            urls.append(
-                "crate://crate@localhost/?schema=testdrive"
-            )
+            urls.append("crate://crate@localhost/?schema=testdrive")
             ids.append("cratedb")
 
         metafunc.parametrize("db_uri", urls, ids=ids)

libs/community/tests/integration_tests/vectorstores/test_cratedb.py

@@ -4,6 +4,7 @@
 cd tests/integration_tests/vectorstores/docker-compose
 docker-compose -f cratedb.yml up
 """
+
 import os
 import re
 from typing import Dict, Generator, List

libs/community/tests/unit_tests/test_sql_database.py

@@ -4,7 +4,6 @@
 import pytest
 import sqlalchemy as sa
 import sqlalchemy.orm
-from langchain_community.utilities.sql_database import SQLDatabase, truncate_word
 from packaging import version
 from sqlalchemy import (
     Column,
@@ -18,6 +17,8 @@
 )
 from sqlalchemy.engine import Engine, Result
 
+from langchain_community.utilities.sql_database import SQLDatabase, truncate_word
+
 is_sqlalchemy_v1 = version.parse(sa.__version__).major == 1
 
 metadata_obj = MetaData()

libs/langchain/langchain/memory/chat_message_histories/cratedb.py

@@ -13,7 +13,7 @@
 # handling optional imports.
 DEPRECATED_LOOKUP = {
     "CrateDBChatMessageHistory": "langchain_community.chat_message_histories",
-    "CrateDBMessageConverter": "langchain_community.chat_message_histories"
+    "CrateDBMessageConverter": "langchain_community.chat_message_histories",
 }
 
 _import_attribute = create_importer(__package__, deprecated_lookups=DEPRECATED_LOOKUP)