Add type checking of ragstack-langchain (#619)

Author: cbornet

.github/workflows/ci-unit-tests.yml

@@ -95,6 +95,12 @@ jobs:
       - name: "Type check (knowledge-store)"
         run: tox -e type -c libs/knowledge-store && rm -rf libs/knowledge-store/.tox
 
+      - name: "Type check (langchain)"
+        run: tox -e type -c libs/langchain && rm -rf libs/langchain/.tox
+
+      - name: "Type check (llama-index)"
+        run: tox -e type -c libs/llamaindex && rm -rf libs/llamaindex/.tox
+
       - name: "Type check (ragulate)"
         run: tox -e type -c libs/ragulate && rm -rf libs/ragulate/.tox
 

libs/langchain/pyproject.toml

@@ -40,6 +40,26 @@ ragstack-ai-colbert = { path = "../colbert", develop = true }
 ragstack-ai-knowledge-store = { path = "../knowledge-store", develop = true }
 pytest-asyncio = "^0.23.6"
 
+[tool.poetry.group.dev.dependencies]
+mypy = "^1.11.0"
+
+[tool.mypy]
+disallow_any_generics = true
+disallow_incomplete_defs = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+follow_imports = "normal"
+ignore_missing_imports = true
+no_implicit_reexport = true
+show_error_codes = true
+show_error_context = true
+strict_equality = true
+strict_optional = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unused_ignores = true
+
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
 

libs/langchain/ragstack_langchain/colbert/colbert_retriever.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from langchain_core.callbacks.manager import (
     AsyncCallbackManagerForRetrieverRun,
@@ -39,11 +39,10 @@ class ColbertRetriever(BaseRetriever):
     def __init__(
         self,
         retriever: ColbertBaseRetriever,
-        k: Optional[int] = 5,
+        k: int = 5,
         query_maxlen: Optional[int] = None,
-        **kwargs: Any,
     ):
-        super().__init__(retriever=retriever, k=k, **kwargs)
+        super().__init__()
         self.retriever = retriever
         self.k = k
         self.query_maxlen = query_maxlen

libs/langchain/ragstack_langchain/colbert/colbert_vector_store.py

@@ -1,21 +1,25 @@
-from typing import Any, Iterable, List, Optional, Tuple, Type, TypeVar
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple
 
 from langchain_core.documents import Document
-from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
 from ragstack_colbert import Chunk
 from ragstack_colbert import ColbertVectorStore as RagstackColbertVectorStore
-from ragstack_colbert.base_database import BaseDatabase as ColbertBaseDatabase
-from ragstack_colbert.base_embedding_model import (
-    BaseEmbeddingModel as ColbertBaseEmbeddingModel,
-)
-from ragstack_colbert.base_retriever import BaseRetriever as ColbertBaseRetriever
-from ragstack_colbert.base_vector_store import BaseVectorStore as ColbertBaseVectorStore
-from typing_extensions import override
+from typing_extensions import Self, override
 
 from ragstack_langchain.colbert.embedding import TokensEmbeddings
 
-CVS = TypeVar("CVS", bound="ColbertVectorStore")
+if TYPE_CHECKING:
+    from langchain_core.embeddings import Embeddings
+    from ragstack_colbert.base_database import BaseDatabase as ColbertBaseDatabase
+    from ragstack_colbert.base_embedding_model import (
+        BaseEmbeddingModel as ColbertBaseEmbeddingModel,
+    )
+    from ragstack_colbert.base_retriever import BaseRetriever as ColbertBaseRetriever
+    from ragstack_colbert.base_vector_store import (
+        BaseVectorStore as ColbertBaseVectorStore,
+    )
 
 
 class ColbertVectorStore(VectorStore):
@@ -35,7 +39,7 @@ def _initialize(
         self,
         database: ColbertBaseDatabase,
         embedding_model: ColbertBaseEmbeddingModel,
-    ):
+    ) -> None:
         self._vector_store = RagstackColbertVectorStore(
             database=database, embedding_model=embedding_model
         )
@@ -45,7 +49,7 @@ def _initialize(
     def add_texts(
         self,
         texts: Iterable[str],
-        metadatas: Optional[List[dict]] = None,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
         doc_id: Optional[str] = None,
         **kwargs: Any,
     ) -> List[str]:
@@ -60,17 +64,18 @@ def add_texts(
         Returns:
             List of ids from adding the texts into the vectorstore.
         """
-        return self._vector_store.add_texts(
+        results = self._vector_store.add_texts(
             texts=list(texts), metadatas=metadatas, doc_id=doc_id
         )
+        return [results[0][0]] if results else []
 
     @override
     async def aadd_texts(
         self,
         texts: Iterable[str],
-        metadatas: Optional[List[dict]] = None,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
         doc_id: Optional[str] = None,
-        concurrent_inserts: Optional[int] = 100,
+        concurrent_inserts: int = 100,
         **kwargs: Any,
     ) -> List[str]:
         """Run more texts through the embeddings and add to the vectorstore.
@@ -86,51 +91,30 @@ async def aadd_texts(
         Returns:
             List of ids from adding the texts into the vectorstore.
         """
-        return await self._vector_store.aadd_texts(
+        results = await self._vector_store.aadd_texts(
             texts=list(texts),
             metadatas=metadatas,
             doc_id=doc_id,
             concurrent_inserts=concurrent_inserts,
         )
+        return [results[0][0]] if results else []
 
     @override
     def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
-        """Delete by vector ID or other criteria.
-
-        Args:
-            ids: List of ids to delete.
-            **kwargs: Other keyword arguments that subclasses might use.
-
-        Returns:
-            Optional[bool]: True if deletion is successful,
-            False otherwise, None if not implemented.
-        """
-        return None if ids is None else self._vector_store.delete(ids=ids)
+        return None if ids is None else self._vector_store.delete_chunks(doc_ids=ids)
 
     @override
     async def adelete(
         self,
         ids: Optional[List[str]] = None,
-        concurrent_deletes: Optional[int] = 100,
+        concurrent_deletes: int = 100,
         **kwargs: Any,
     ) -> Optional[bool]:
-        """Delete by vector ID or other criteria.
-
-        Args:
-            ids: List of ids to delete.
-            concurrent_deletes: How many concurrent deletes to make to the database.
-                Defaults to 100.
-            **kwargs: Other keyword arguments that subclasses might use.
-
-        Returns:
-            Optional[bool]: True if deletion is successful,
-            False otherwise, None if not implemented.
-        """
         return (
             None
             if ids is None
-            else await self._vector_store.adelete(
-                ids=ids, concurrent_deletes=concurrent_deletes
+            else await self._vector_store.adelete_chunks(
+                doc_ids=ids, concurrent_deletes=concurrent_deletes
             )
         )
 
@@ -215,7 +199,7 @@ def from_documents(
         *,
         database: Optional[ColbertBaseDatabase] = None,
         **kwargs: Any,
-    ) -> CVS:
+    ) -> Self:
         """Return VectorStore initialized from documents and embeddings."""
         texts = [d.page_content for d in documents]
         metadatas = [d.metadata for d in documents]
@@ -230,14 +214,14 @@ def from_documents(
     @classmethod
     @override
     async def afrom_documents(
-        cls: Type[CVS],
+        cls,
         documents: List[Document],
         embedding: Embeddings,
         *,
         database: Optional[ColbertBaseDatabase] = None,
-        concurrent_inserts: Optional[int] = 100,
+        concurrent_inserts: int = 100,
         **kwargs: Any,
-    ) -> CVS:
+    ) -> Self:
         """Return VectorStore initialized from documents and embeddings."""
         texts = [d.page_content for d in documents]
         metadatas = [d.metadata for d in documents]
@@ -253,14 +237,14 @@ async def afrom_documents(
     @classmethod
     @override
     def from_texts(
-        cls: Type[CVS],
+        cls,
         texts: List[str],
         embedding: Embeddings,
-        metadatas: Optional[List[dict]] = None,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
         *,
         database: Optional[ColbertBaseDatabase] = None,
         **kwargs: Any,
-    ) -> CVS:
+    ) -> Self:
         if not isinstance(embedding, TokensEmbeddings):
             raise TypeError("ColbertVectorStore requires a TokensEmbeddings embedding.")
         if database is None:
@@ -276,15 +260,15 @@ def from_texts(
     @classmethod
     @override
     async def afrom_texts(
-        cls: Type[CVS],
+        cls,
         texts: List[str],
         embedding: Embeddings,
-        metadatas: Optional[List[dict]] = None,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
         *,
         database: Optional[ColbertBaseDatabase] = None,
-        concurrent_inserts: Optional[int] = 100,
+        concurrent_inserts: int = 100,
         **kwargs: Any,
-    ) -> CVS:
+    ) -> Self:
         if not isinstance(embedding, TokensEmbeddings):
             raise TypeError("ColbertVectorStore requires a TokensEmbeddings embedding.")
         if database is None:

libs/langchain/ragstack_langchain/colbert/embedding.py

@@ -3,13 +3,13 @@
 from langchain_core.embeddings import Embeddings
 from ragstack_colbert import DEFAULT_COLBERT_MODEL, ColbertEmbeddingModel
 from ragstack_colbert.base_embedding_model import BaseEmbeddingModel
-from typing_extensions import override
+from typing_extensions import Self, override
 
 
 class TokensEmbeddings(Embeddings):
     """Adapter for token-based embedding models and the LangChain Embeddings."""
 
-    def __init__(self, embedding: BaseEmbeddingModel = None):
+    def __init__(self, embedding: Optional[BaseEmbeddingModel] = None):
         self.embedding = embedding or ColbertEmbeddingModel()
 
     @override
@@ -32,8 +32,9 @@ def get_embedding_model(self) -> BaseEmbeddingModel:
         """Get the embedding model."""
         return self.embedding
 
-    @staticmethod
+    @classmethod
     def colbert(
+        cls,
         checkpoint: str = DEFAULT_COLBERT_MODEL,
         doc_maxlen: int = 256,
         nbits: int = 2,
@@ -42,9 +43,9 @@ def colbert(
         query_maxlen: Optional[int] = None,
         verbose: int = 3,
         chunk_batch_size: int = 640,
-    ):
+    ) -> Self:
         """Create a new ColBERT embedding model."""
-        return TokensEmbeddings(
+        return cls(
             ColbertEmbeddingModel(
                 checkpoint,
                 doc_maxlen,

libs/langchain/tests/integration_tests/conftest.py

@@ -1,6 +1,7 @@
 from typing import Iterator
 
 import pytest
+from _pytest.fixtures import FixtureRequest
 from cassandra.cluster import Session
 from dotenv import load_dotenv
 from ragstack_tests_utils import AstraDBTestStore, LocalCassandraTestStore
@@ -21,13 +22,13 @@ def astra_db() -> AstraDBTestStore:
     return AstraDBTestStore()
 
 
-def get_session(request) -> Session:
+def get_session(request: FixtureRequest) -> Session:
     test_store = request.getfixturevalue(request.param)
     session = test_store.create_cassandra_session()
     session.default_timeout = 180
     return session
 
 
 @pytest.fixture()
-def session(request) -> Session:
+def session(request: FixtureRequest) -> Session:
     return get_session(request)