pycityagent 2.0.0a66__cp310-cp310-macosx_11_0_arm64.whl → 2.0.0a67__cp310-cp310-macosx_11_0_arm64.whl

pycityagent/memory/faiss_query.py

@@ -14,12 +14,31 @@ from ..utils.decorators import lock_decorator
 
 
 class FaissQuery:
+    """
+    A class for handling similarity searches and document management using FAISS.
+
+    - **Description**:
+        - This class provides functionalities to manage embeddings and perform similarity searches over a set of documents.
+        - It allows adding, deleting, and querying documents based on their semantic content through embeddings.
+        - The class initializes with an optional embedding model and index type, setting up the environment for vector operations.
+        - If no embedding model is provided during initialization, certain methods will not be available until one is set.
+    """
+
     def __init__(
         self,
         embeddings: Optional[Embeddings] = None,
         index_type: Any = faiss.IndexFlatL2,
         dimension: Optional[int] = None,
     ) -> None:
+        """
+        Initialize the FaissQuery instance.
+
+        - **Parameters**:
+            - `embeddings` (Optional[Embeddings], optional): An embedding function that converts text into vectors. Defaults to None.
+            - `index_type` (Any, optional): The type of FAISS index to use for vector storage and retrieval. Defaults to faiss.IndexFlatL2.
+            - `dimension` (Optional[int], optional): The dimensionality of the vectors produced by the embedding function. Defaults to None.
+                If not specified, it's inferred from the embedding of a sample text ("hello world").
+        """
         self._embeddings = embeddings
         self._lock = asyncio.Lock()
         if embeddings is None:
@@ -40,6 +59,16 @@ class FaissQuery:
     def embeddings(
         self,
     ) -> Embeddings:
+        """
+        Access the current embedding model.
+
+        - **Description**:
+            - Getter for the embedding model used to convert text into vector representations.
+            - If no embedding model has been set, a RuntimeError is raised prompting the user to set one first.
+
+        - **Returns**:
+            - `Embeddings`: The current embedding model.
+        """
         if self._embeddings is None:
             raise RuntimeError(f"No embedding set, please `set_embeddings` first!")
         return self._embeddings
@@ -48,6 +77,16 @@ class FaissQuery:
     def vectors_store(
         self,
     ) -> FAISS:
+        """
+        Access the current vector store.
+
+        - **Description**:
+            - Getter for the FAISS vector store which holds the indexed documents and allows for similarity searches.
+            - If the vector store hasn't been initialized due to a missing embedding model, a RuntimeError is raised.
+
+        - **Returns**:
+            - `FAISS`: The current vector store instance.
+        """
         if self._vectors_store is None:
             raise RuntimeError(f"No embedding set, thus no vector stores initialized!")
         return self._vectors_store
@@ -59,6 +98,21 @@ class FaissQuery:
         documents: Union[str, Sequence[str]],
         extra_tags: Optional[dict] = None,
     ) -> list[str]:
+        """
+        Add documents to the vector store with metadata.
+
+        - **Description**:
+            - Asynchronously adds one or more documents to the vector store, associating them with an agent ID and optional extra tags.
+            - Each document is converted into a vector using the embedding model before being added to the index.
+
+        - **Args**:
+            - `agent_id` (int): Identifier of the agent to associate with the documents.
+            - `documents` (Union[str, Sequence[str]]): A single document string or a sequence of document strings to add.
+            - `extra_tags` (Optional[dict], optional): Additional metadata tags to associate with the documents. Defaults to None.
+
+        - **Returns**:
+            - `list[str]`: List of document IDs that were added to the vector store.
+        """
         if isinstance(documents, str):
             documents = [documents]
         _metadata = {"_id": agent_id}
@@ -76,6 +130,15 @@ class FaissQuery:
         self,
         to_delete_ids: list[str],
     ):
+        """
+        Delete documents from the vector store by IDs.
+
+        - **Description**:
+            - Asynchronously deletes documents from the vector store based on provided document IDs.
+
+        - **Args**:
+            - `to_delete_ids` (list[str]): List of document IDs to delete from the vector store.
+        """
         await self.vectors_store.adelete(
             ids=to_delete_ids,
         )
@@ -93,24 +156,26 @@ class FaissQuery:
         filter: Optional[dict] = None,
     ) -> Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]:
         """
-        Return content most similar to the given query.
-
-        Args:
-            query (str): The text to look up documents similar to.
-            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
-            k (int, optional): The number of top similar contents to return. Defaults to 4.
-            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
-            return_score_type (Union[Literal["none"], Literal["similarity_score"], Literal["L2-distance"]], optional):
-                Specifies whether and how to return similarity scores with the results:
-                    - "none": Do not return scores; only return the contents (default).
-                    - "similarity_score": Return a tuple of content and its similarity score.
-                    - "L2-distance": Return a tuple of content and its L2 distance from the query.
-            filter (dict, optional): The filter dict for metadata.
+        Perform a similarity search for documents related to the given query.
+
+        - **Description**:
+            - Conducts an asynchronous search for the top-k documents most similar to the query text.
+            - The search can be customized by specifying how many documents to fetch (`fetch_k`), how many to return (`k`),
+              and whether to include scores in the results (`return_score_type`).
+            - Filters can be applied to narrow down the search results based on metadata.
 
-        Returns:
-            Union[list[tuple[str,dict]], list[tuple[str, float,dict]]]:
-                Depending on the `return_score_type` parameter, returns either a list of strings representing the top-k similar contents,
-                or a list of tuples where each tuple contains a string and a floating-point score.
+        - **Args**:
+            - `query` (str): The text to look up documents similar to.
+            - `agent_id` (int): The identifier of the agent to filter specific documents.
+            - `k` (int, optional): The number of top similar contents to return. Defaults to 4.
+            - `fetch_k` (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            - `return_score_type` (Union[Literal["none"], Literal["similarity_score"], Literal["L2-distance"]], optional):
+                Specifies whether and how to return similarity scores with the results.
+            - `filter` (Optional[dict], optional): The filter dict for metadata.
+
+        - **Returns**:
+            - `Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]`: Depending on the `return_score_type` parameter,
+              returns either a list of tuples containing the content and its associated metadata, or also including a floating-point score.
         """
         _filter = {
             "_id": agent_id,
@@ -136,11 +201,13 @@ class FaissQuery:
                 )
                 return [(r.page_content, r.metadata) for r in _result]
             elif return_score_type == "similarity_score":
-                _result = await self.vectors_store.asimilarity_search_with_relevance_scores(
-                    query=query,
-                    k=k,
-                    filter=_filter,
-                    fetch_k=fetch_k,
+                _result = (
+                    await self.vectors_store.asimilarity_search_with_relevance_scores(
+                        query=query,
+                        k=k,
+                        filter=_filter,
+                        fetch_k=fetch_k,
+                    )
                 )
                 return [(r.page_content, s, r.metadata) for r, s in _result]
             else:
@@ -157,23 +224,28 @@ class FaissQuery:
         filter: Optional[dict] = None,
     ) -> Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]:
         """
-        Return content most similar to the given query.
-
-        Args:
-            embedding (list[float]): The vector to look up documents similar to.
-            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
-            k (int, optional): The number of top similar contents to return. Defaults to 4.
-            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
-            return_score_type (Union[Literal["none"], Literal["similarity_score"], Literal["L2-distance"]], optional):
+        Perform a similarity search for documents related to the given vector.
+
+        - **Description**:
+            - Conducts an asynchronous search for the top-k documents most similar to the provided embedding vector.
+            - The search can be customized by specifying how many documents to fetch (`fetch_k`), how many to return (`k`),
+              and whether to include L2 distances in the results (`return_score_type`).
+            - Filters can be applied to narrow down the search results based on metadata.
+
+        - **Args**:
+            - `embedding` (list[float]): The vector to look up documents similar to.
+            - `agent_id` (int): The identifier of the agent to filter specific documents.
+            - `k` (int, optional): The number of top similar contents to return. Defaults to 4.
+            - `fetch_k` (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            - `return_score_type` (Union[Literal["none"], Literal["L2-distance"]], optional):
                 Specifies whether and how to return similarity scores with the results:
                     - "none": Do not return scores; only return the contents (default).
                     - "L2-distance": Return a tuple of content and its L2 distance from the query.
-            filter (dict, optional): The filter dict for metadata.
+            - `filter` (Optional[dict], optional): The filter dict for metadata.
 
-        Returns:
-            Union[list[tuple[str,dict]], list[tuple[str, float,dict]]]:
-                Depending on the `return_score_type` parameter, returns either a list of strings representing the top-k similar contents,
-                or a list of tuples where each tuple contains a string and a floating-point score.
+        - **Returns**:
+            - `Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]`: Depending on the `return_score_type` parameter,
+              returns either a list of tuples containing the content and its associated metadata, or also including a floating-point score.
         """
         _filter = {
             "_id": agent_id,
@@ -211,21 +283,25 @@ class FaissQuery:
         filter: Optional[dict] = None,
     ) -> list[tuple[str, dict]]:
         """
-        Return contents selected using the maximal marginal relevance asynchronously.
-
-        Args:
-            query (str): The text to look up documents similar to.
-            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
-            k (int, optional): The number of top similar contents to return. Defaults to 4.
-            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
-            lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
-            return_score_type (Literal["none"].,optional):
+        Select contents using maximal marginal relevance asynchronously.
+
+        - **Description**:
+            - Asynchronously selects a set of documents that are relevant to the query while ensuring diversity among the results.
+            - The selection process balances between relevance to the query and diversity, controlled by `lambda_mult`.
+
+        - **Args**:
+            - `query` (str): The text to look up documents similar to.
+            - `agent_id` (int): The identifier of the agent to filter specific documents.
+            - `k` (int, optional): The number of top similar contents to return. Defaults to 4.
+            - `fetch_k` (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            - `lambda_mult` (float, optional): Number between 0 and 1 that determines the degree of diversity among the results. Defaults to 0.5.
+            - `return_score_type` (Literal["none"], optional):
                 Specifies whether and how to return similarity scores with the results:
                     - "none": Do not return scores; only return the contents (default).
-            filter (dict, optional): The filter dict for metadata.
+            - `filter` (Optional[dict], optional): The filter dict for metadata.
 
-        Returns:
-            list[tuple[str,dict]]: the result contents.
+        - **Returns**:
+            - `list[tuple[str, dict]]`: A list of tuples containing the content and its associated metadata.
         """
         _filter = {
             "_id": agent_id,
@@ -257,24 +333,27 @@ class FaissQuery:
         filter: Optional[dict] = None,
     ) -> Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]:
         """
-        Return contents selected using the maximal marginal relevance asynchronously.
-
-        Args:
-            embedding (list[float]): The vector to look up documents similar to.
-            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
-            k (int, optional): The number of top similar contents to return. Defaults to 4.
-            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
-            lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
-            return_score_type (Union[Literal["none"], Literal["similarity_score"]], optional):
+        Select contents using maximal marginal relevance asynchronously based on embedding.
+
+        - **Description**:
+            - Asynchronously selects a set of documents that are relevant to the provided embedding vector while ensuring diversity among the results.
+            - The selection process balances between relevance to the embedding and diversity, controlled by `lambda_mult`.
+
+        - **Args**:
+            - `embedding` (list[float]): The vector to look up documents similar to.
+            - `agent_id` (int): The identifier of the agent to filter specific documents.
+            - `k` (int, optional): The number of top similar contents to return. Defaults to 4.
+            - `fetch_k` (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            - `lambda_mult` (float, optional): Number between 0 and 1 that determines the degree of diversity among the results. Defaults to 0.5.
+            - `return_score_type` (Union[Literal["none"], Literal["similarity_score"]], optional):
                 Specifies whether and how to return similarity scores with the results:
                     - "none": Do not return scores; only return the contents (default).
                     - "similarity_score": Return a tuple of content and its similarity score.
-            filter (dict, optional): The filter dict for metadata.
+            - `filter` (Optional[dict], optional): The filter dict for metadata.
 
-        Returns:
-            Union[list[tuple[str,dict]], list[tuple[str, float,dict]]]:
-                Depending on the `return_score_type` parameter, returns either a list of strings representing the top-k similar contents,
-                or a list of tuples where each tuple contains a string and a floating-point score.
+        - **Returns**:
+            - `Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]`: Depending on the `return_score_type` parameter,
+              returns either a list of tuples containing the content and its associated metadata, or also including a floating-point score.
         """
 
         _filter = {