Add ruff rules for docstrings (#576)

Author: cbornet

libs/colbert/ragstack_colbert/__init__.py

@@ -1,4 +1,5 @@
-"""
+"""Ragstack Colbert: A ColBERT-based text retrieval system.
+
 This package provides a suite of tools for encoding and retrieving text using the
 ColBERT model, integrated with a Cassandra database for scalable storage and retrieval
 operations. It includes classes for token embeddings, managing the vector store, and

libs/colbert/ragstack_colbert/base_database.py

@@ -1,6 +1,7 @@
-"""
-This module defines abstract base classes for implementing storage mechanisms for text
-chunk embeddings, specifically designed to work with ColBERT or similar embedding
+"""Base Database module.
+
+This module defines abstract base classes for implementing storage mechanisms for
+text chunk embeddings, specifically designed to work with ColBERT or similar embedding
 models.
 """
 
@@ -11,7 +12,8 @@
 
 
 class BaseDatabase(ABC):
-    """
+    """Base Database abstract class for ColBERT.
+
     Abstract base class (ABC) for a storage system designed to hold vector
     representations of text chunks, typically generated by a ColBERT model or similar
     embedding model.
@@ -23,10 +25,9 @@ class BaseDatabase(ABC):
 
     @abstractmethod
     def add_chunks(self, chunks: List[Chunk]) -> List[Tuple[str, int]]:
-        """
-        Stores a list of embedded text chunks in the vector store
+        """Stores a list of embedded text chunks in the vector store.
 
-        Parameters:
+        Args:
             chunks (List[Chunk]): A list of `Chunk` instances to be stored.
 
         Returns:
@@ -35,10 +36,9 @@ def add_chunks(self, chunks: List[Chunk]) -> List[Tuple[str, int]]:
 
     @abstractmethod
     def delete_chunks(self, doc_ids: List[str]) -> bool:
-        """
-        Deletes chunks from the vector store based on their document id.
+        """Deletes chunks from the vector store based on their document id.
 
-        Parameters:
+        Args:
             doc_ids (List[str]): A list of document identifiers specifying the chunks
                 to be deleted.
 
@@ -50,10 +50,9 @@ def delete_chunks(self, doc_ids: List[str]) -> bool:
     async def aadd_chunks(
         self, chunks: List[Chunk], concurrent_inserts: Optional[int] = 100
     ) -> List[Tuple[str, int]]:
-        """
-        Stores a list of embedded text chunks in the vector store
+        """Stores a list of embedded text chunks in the vector store.
 
-        Parameters:
+        Args:
             chunks (List[Chunk]): A list of `Chunk` instances to be stored.
             concurrent_inserts (Optional[int]): How many concurrent inserts to make to
                 the database. Defaults to 100.
@@ -66,10 +65,9 @@ async def aadd_chunks(
     async def adelete_chunks(
         self, doc_ids: List[str], concurrent_deletes: Optional[int] = 100
     ) -> bool:
-        """
-        Deletes chunks from the vector store based on their document id.
+        """Deletes chunks from the vector store based on their document id.
 
-        Parameters:
+        Args:
             doc_ids (List[str]): A list of document identifiers specifying the chunks
                 to be deleted.
             concurrent_deletes (Optional[int]): How many concurrent deletes to make
@@ -81,8 +79,7 @@ async def adelete_chunks(
 
     @abstractmethod
     async def search_relevant_chunks(self, vector: Vector, n: int) -> List[Chunk]:
-        """
-        Retrieves 'n' ANN results for an embedded token vector.
+        """Retrieves 'n' ANN results for an embedded token vector.
 
         Returns:
             A list of Chunks with only `doc_id` and `chunk_id` set.
@@ -91,8 +88,7 @@ async def search_relevant_chunks(self, vector: Vector, n: int) -> List[Chunk]:
 
     @abstractmethod
     async def get_chunk_embedding(self, doc_id: str, chunk_id: int) -> Chunk:
-        """
-        Retrieve the embedding data for a chunk.
+        """Retrieve the embedding data for a chunk.
 
         Returns:
             A chunk with `doc_id`, `chunk_id`, and `embedding` set.
@@ -102,8 +98,7 @@ async def get_chunk_embedding(self, doc_id: str, chunk_id: int) -> Chunk:
     async def get_chunk_data(
         self, doc_id: str, chunk_id: int, include_embedding: Optional[bool]
     ) -> Chunk:
-        """
-        Retrieve the text and metadata for a chunk.
+        """Retrieve the text and metadata for a chunk.
 
         Returns:
             A chunk with `doc_id`, `chunk_id`, `text`, `metadata`, and optionally
@@ -112,6 +107,4 @@ async def get_chunk_data(
 
     @abstractmethod
     def close(self) -> None:
-        """
-        Cleans up any open resources.
-        """
+        """Cleans up any open resources."""

libs/colbert/ragstack_colbert/base_embedding_model.py

@@ -1,6 +1,7 @@
-"""
-This module defines an abstract base class (ABC) for generating token-based embeddings
-for text.
+"""Base embedding for ColBERT.
+
+This module defines an abstract base class (ABC) for generating token-based
+embeddings for text.
 """
 
 from abc import ABC, abstractmethod
@@ -10,8 +11,7 @@
 
 
 class BaseEmbeddingModel(ABC):
-    """
-    Abstract base class (ABC) for token-based embedding models.
+    """Abstract base class (ABC) for token-based embedding models.
 
     This class defines the interface for models that generate embeddings for text
     chunks and queries.
@@ -22,11 +22,9 @@ class BaseEmbeddingModel(ABC):
 
     @abstractmethod
     def embed_texts(self, texts: List[str]) -> List[Embedding]:
-        """
-        Embeds a list of texts into their corresponding vector embedding
-        representations.
+        """Embeds a list of texts into their vector embedding representations.
 
-        Parameters:
+        Args:
             texts (List[str]): A list of string texts.
 
         Returns:
@@ -40,13 +38,12 @@ def embed_query(
         full_length_search: Optional[bool] = False,
         query_maxlen: int = -1,
     ) -> Embedding:
-        """
-        Embeds a single query text into its vector representation.
+        """Embeds a single query text into its vector representation.
 
         If the query has fewer than query_maxlen tokens it will be padded with BERT
         special [mast] tokens.
 
-        Parameters:
+        Args:
             query (str): The query text to encode.
             full_length_search (Optional[bool]): Indicates whether to encode the
                 query for a full-length search. Defaults to False.

libs/colbert/ragstack_colbert/base_retriever.py

@@ -1,4 +1,5 @@
-"""
+"""Base retriever module.
+
 This module defines abstract base classes for implementing retrieval mechanisms for
 text chunk embeddings, specifically designed to work with ColBERT or similar embedding
 models.
@@ -11,9 +12,10 @@
 
 
 class BaseRetriever(ABC):
-    """
-    Abstract base class (ABC) for a retrieval system that operates on a ColBERT vector
-    store, facilitating the search and retrieval of text chunks based on query
+    """Base Retriever abstract class for ColBERT.
+
+    Abstract base class (ABC) for a retrieval system that operates on a ColBERT
+    vector store, facilitating the search and retrieval of text chunks based on query
     embeddings.
     """
 
@@ -26,11 +28,12 @@ def embedding_search(
         include_embedding: Optional[bool] = False,
         **kwargs: Any,
     ) -> List[Tuple[Chunk, float]]:
-        """
+        """Search for relevant text chunks based on a query embedding.
+
         Retrieves a list of text chunks relevant to a given query from the vector
         store, ranked by relevance or other metrics.
 
-        Parameters:
+        Args:
             query_embedding (Embedding): The query embedding to search for relevant
                 text chunks.
             k (Optional[int]): The number of top results to retrieve.
@@ -54,11 +57,12 @@ async def aembedding_search(
         include_embedding: Optional[bool] = False,
         **kwargs: Any,
     ) -> List[Tuple[Chunk, float]]:
-        """
+        """Search for relevant text chunks based on a query embedding.
+
         Retrieves a list of text chunks relevant to a given query from the vector
         store, ranked by relevance or other metrics.
 
-        Parameters:
+        Args:
             query_embedding (Embedding): The query embedding to search for relevant
                 text chunks.
             k (Optional[int]): The number of top results to retrieve.
@@ -83,11 +87,12 @@ def text_search(
         include_embedding: Optional[bool] = False,
         **kwargs: Any,
     ) -> List[Tuple[Chunk, float]]:
-        """
+        """Search for relevant text chunks based on a query text.
+
         Retrieves a list of text chunks relevant to a given query from the vector
         store, ranked by relevance or other metrics.
 
-        Parameters:
+        Args:
             query_text (str): The query text to search for relevant text chunks.
             k (Optional[int]): The number of top results to retrieve.
             query_maxlen (Optional[int]): The maximum length of the query to consider.
@@ -113,11 +118,12 @@ async def atext_search(
         include_embedding: Optional[bool] = False,
         **kwargs: Any,
     ) -> List[Tuple[Chunk, float]]:
-        """
+        """Search for relevant text chunks based on a query text.
+
         Retrieves a list of text chunks relevant to a given query from the vector
         store, ranked by relevance or other metrics.
 
-        Parameters:
+        Args:
             query_text (str): The query text to search for relevant text chunks.
             k (Optional[int]): The number of top results to retrieve.
             query_maxlen (Optional[int]): The maximum length of the query to consider.

libs/colbert/ragstack_colbert/base_vector_store.py

@@ -1,4 +1,5 @@
-"""
+"""Base Vector Store module for ColBERT.
+
 This module defines the abstract base class for a standard vector store
 specifically designed to work with ColBERT or similar dense embedding models,
 and can be used to create a LangChain or LlamaIndex ColBERT vector store.
@@ -23,7 +24,8 @@
 
 
 class BaseVectorStore(ABC):
-    """
+    """Base Vector Store abstract class for ColBERT.
+
     Abstract base class (ABC) for a storage system designed to hold vector
     representations of text chunks, typically generated by a ColBERT model or similar
     embedding model.
@@ -36,10 +38,9 @@ class BaseVectorStore(ABC):
     # handles LlamaIndex add
     @abstractmethod
     def add_chunks(self, chunks: List[Chunk]) -> List[Tuple[str, int]]:
-        """
-        Stores a list of embedded text chunks in the vector store
+        """Stores a list of embedded text chunks in the vector store.
 
-        Parameters:
+        Args:
             chunks (List[Chunk]): A list of `Chunk` instances to be stored.
 
         Returns:
@@ -54,11 +55,12 @@ def add_texts(
         metadatas: Optional[List[Metadata]],
         doc_id: Optional[str] = None,
     ) -> List[Tuple[str, int]]:
-        """
+        """Adds text chunks to the vector store.
+
         Embeds and stores a list of text chunks and optional metadata into the vector
         store.
 
-        Parameters:
+        Args:
             texts (List[str]): The list of text chunks to be embedded
             metadatas (Optional[List[Metadata]])): An optional list of Metadata to be
                 stored. If provided, these are set 1 to 1 with the texts list.
@@ -72,10 +74,9 @@ def add_texts(
     # handles LangChain and LlamaIndex delete
     @abstractmethod
     def delete_chunks(self, doc_ids: List[str]) -> bool:
-        """
-        Deletes chunks from the vector store based on their document id.
+        """Deletes chunks from the vector store based on their document id.
 
-        Parameters:
+        Args:
             doc_ids (List[str]): A list of document identifiers specifying the chunks
                 to be deleted.
 
@@ -88,10 +89,9 @@ def delete_chunks(self, doc_ids: List[str]) -> bool:
     async def aadd_chunks(
         self, chunks: List[Chunk], concurrent_inserts: Optional[int] = 100
     ) -> List[Tuple[str, int]]:
-        """
-        Stores a list of embedded text chunks in the vector store
+        """Stores a list of embedded text chunks in the vector store.
 
-        Parameters:
+        Args:
             chunks (List[Chunk]): A list of `Chunk` instances to be stored.
             concurrent_inserts (Optional[int]): How many concurrent inserts to make to
                 the database. Defaults to 100.
@@ -109,11 +109,12 @@ async def aadd_texts(
         doc_id: Optional[str] = None,
         concurrent_inserts: Optional[int] = 100,
     ) -> List[Tuple[str, int]]:
-        """
+        """Adds text chunks to the vector store.
+
         Embeds and stores a list of text chunks and optional metadata into the vector
         store.
 
-        Parameters:
+        Args:
             texts (List[str]): The list of text chunks to be embedded
             metadatas (Optional[List[Metadata]])): An optional list of Metadata to be
                 stored. If provided, these are set 1 to 1 with the texts list.
@@ -131,10 +132,9 @@ async def aadd_texts(
     async def adelete_chunks(
         self, doc_ids: List[str], concurrent_deletes: Optional[int] = 100
     ) -> bool:
-        """
-        Deletes chunks from the vector store based on their document id.
+        """Deletes chunks from the vector store based on their document id.
 
-        Parameters:
+        Args:
             doc_ids (List[str]): A list of document identifiers specifying the chunks
                 to be deleted.
             concurrent_deletes (Optional[int]): How many concurrent deletes to make to
@@ -147,6 +147,4 @@ async def adelete_chunks(
     # handles LangChain as_retriever
     @abstractmethod
     def as_retriever(self) -> BaseRetriever:
-        """
-        Gets a retriever using the vector store.
-        """
+        """Gets a retriever using the vector store."""