Update context grounding vector store to be compatible with gym

Author: adrian-badea

src/uipath_langchain/vectorstores/context_grounding_vectorstore.py

@@ -2,66 +2,52 @@
 Vector store implementation that connects to UiPath Context Grounding as a backend.
 
 This is a read-only vector store that uses the UiPath Context Grounding API to retrieve documents.
-
-You need to set the following environment variables (also see .env.example):
-### - UIPATH_URL="https://alpha.uipath.com/{ORG_ID}/{TENANT_ID}"
-### - UIPATH_ACCESS_TOKEN={BEARER_TOKEN_WITH_CONTEXT_GROUNDING_PERMISSIONS}
-### - UIPATH_FOLDER_PATH="" - this can be left empty
-### - UIPATH_FOLDER_KEY="" - this can be left empty
 """
 
 from collections.abc import Iterable
-from typing import Any, Optional, TypeVar
+from typing import Any, Self, TypeVar, override
 
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore
 from uipath import UiPath
-
-VST = TypeVar("VST", bound="ContextGroundingVectorStore")
-
+from uipath.models.context_grounding import ContextGroundingQueryResponse
 
 class ContextGroundingVectorStore(VectorStore):
     """Vector store that uses UiPath Context Grounding (ECS) as a backend.
 
     This class provides a straightforward implementation that connects to the
     UiPath Context Grounding API for semantic searching.
-
-    Example:
-        .. code-block:: python
-
-            from uipath_agents_gym.tools.ecs_vectorstore import ContextGroundingVectorStore
-
-            # Initialize the vector store with an index name
-            vectorstore = ContextGroundingVectorStore(index_name="ECCN")
-
-            # Perform similarity search
-            docs_with_scores = vectorstore.similarity_search_with_score(
-                "How do I process an invoice?", k=5
-            )
     """
 
+
     def __init__(
         self,
         index_name: str,
-        folder_path: Optional[str] = None,
-        uipath_sdk: Optional[UiPath] = None,
+        uipath_sdk: UiPath | None = None,
+        folder_path: str | None = None,
     ):
         """Initialize the ContextGroundingVectorStore.
 
         Args:
-            index_name: Name of the context grounding index to use
-            uipath_sdk: Optional SDK instance to use. If not provided, a new instance will be created.
+            index_name: Name of the context grounding index to use (schema name)
+            uipath_sdk: Optional UiPath SDK instance.
+            folder_path: Optional folder path for folder-scoped operations
         """
+        self.index_name = index_name
+        self.folder_path = folder_path
+
         self.index_name = index_name
         self.folder_path = folder_path
         self.sdk = uipath_sdk or UiPath()
 
+    # VectorStore implementation methods
+
+    @override
     def similarity_search_with_score(
         self, query: str, k: int = 4, **kwargs: Any
     ) -> list[tuple[Document, float]]:
         """Return documents most similar to the query along with the distances.
-        The distance is 1 - score, where score is the relevance score returned by the Context Grounding API.
 
         Args:
             query: The query string
@@ -70,52 +56,22 @@ def similarity_search_with_score(
         Returns:
             list of tuples of (document, score)
         """
-        # Call the UiPath SDK to perform the search
-        results = self.sdk.context_grounding.search(
+        # Use the context grounding service to perform search
+        results: list[ContextGroundingQueryResponse] = self.sdk.context_grounding.search(
             name=self.index_name,
             query=query,
             number_of_results=k,
             folder_path=self.folder_path,
         )
 
-        # Convert the results to Documents with scores
-        docs_with_scores = []
-        for result in results:
-            # Create metadata from result fields
-            metadata = {
-                "source": result.source,
-                "id": result.id,
-                "reference": result.reference,
-                "page_number": result.page_number,
-                "source_document_id": result.source_document_id,
-                "caption": result.caption,
-            }
-
-            # Add any operation metadata if available
-            if result.metadata:
-                metadata["operation_id"] = result.metadata.operation_id
-                metadata["strategy"] = result.metadata.strategy
-
-            # Create a Document with the content and metadata
-            doc = Document(
-                page_content=result.content,
-                metadata=metadata,
-            )
-
-            score = 1.0 - float(result.score)
-
-            docs_with_scores.append((doc, score))
-
-        return docs_with_scores
+        return self._convert_results_to_documents(results)
 
+    @override
     def similarity_search_with_relevance_scores(
         self, query: str, k: int = 4, **kwargs: Any
     ) -> list[tuple[Document, float]]:
         """Return documents along with their relevance scores on a scale from 0 to 1.
 
-        This directly uses the scores provided by the Context Grounding API,
-        which are already normalized between 0 and 1.
-
         Args:
             query: The query string
             k: Number of documents to return (default=4)
@@ -128,6 +84,7 @@ def similarity_search_with_relevance_scores(
             for doc, score in self.similarity_search_with_score(query, k, **kwargs)
         ]
 
+    @override
     async def asimilarity_search_with_score(
         self, query: str, k: int = 4, **kwargs: Any
     ) -> list[tuple[Document, float]]:
@@ -140,52 +97,23 @@ async def asimilarity_search_with_score(
         Returns:
             list of tuples of (document, score)
         """
-        # Call the UiPath SDK to perform the search asynchronously
-        results = await self.sdk.context_grounding.search_async(
+        # Use the context grounding service to perform async search
+        results: list[
+            ContextGroundingQueryResponse
+        ] = await self.sdk.context_grounding.search_async(
             name=self.index_name,
             query=query,
             number_of_results=k,
             folder_path=self.folder_path,
         )
 
-        # Convert the results to Documents with scores
-        docs_with_scores = []
-        for result in results:
-            # Create metadata from result fields
-            metadata = {
-                "source": result.source,
-                "id": result.id,
-                "reference": result.reference,
-                "page_number": result.page_number,
-                "source_document_id": result.source_document_id,
-                "caption": result.caption,
-            }
-
-            # Add any operation metadata if available
-            if result.metadata:
-                metadata["operation_id"] = result.metadata.operation_id
-                metadata["strategy"] = result.metadata.strategy
-
-            # Create a Document with the content and metadata
-            doc = Document(
-                page_content=result.content,
-                metadata=metadata,
-            )
-
-            # Get the distance score as 1 - ecs_score
-            score = 1.0 - float(result.score)
-
-            docs_with_scores.append((doc, score))
-
-        return docs_with_scores
+        return self._convert_results_to_documents(results)
 
+    @override
     async def asimilarity_search_with_relevance_scores(
         self, query: str, k: int = 4, **kwargs: Any
     ) -> list[tuple[Document, float]]:
-        """Asynchronously return documents along with their relevance scores on a scale from 0 to 1.
-
-        This directly uses the scores provided by the Context Grounding API,
-        which are already normalized between 0 and 1.
+        """Asynchronously return documents along with their relevance scores.
 
         Args:
             query: The query string
@@ -196,14 +124,11 @@ async def asimilarity_search_with_relevance_scores(
         """
         return [
             (doc, 1.0 - score)
-            for doc, score in await self.asimilarity_search_with_score(
-                query, k, **kwargs
-            )
+            for doc, score in await self.asimilarity_search_with_score(query, k, **kwargs)
         ]
 
-    def similarity_search(
-        self, query: str, k: int = 4, **kwargs: Any
-    ) -> list[Document]:
+    @override
+    def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> list[Document]:
         """Return documents most similar to the query.
 
         Args:
@@ -216,9 +141,8 @@ def similarity_search(
         docs_and_scores = self.similarity_search_with_score(query, k, **kwargs)
         return [doc for doc, _ in docs_and_scores]
 
-    async def asimilarity_search(
-        self, query: str, k: int = 4, **kwargs: Any
-    ) -> list[Document]:
+    @override
+    async def asimilarity_search(self, query: str, k: int = 4, **kwargs: Any) -> list[Document]:
         """Asynchronously return documents most similar to the query.
 
         Args:
@@ -231,38 +155,85 @@ async def asimilarity_search(
         docs_and_scores = await self.asimilarity_search_with_score(query, k, **kwargs)
         return [doc for doc, _ in docs_and_scores]
 
+    def _convert_results_to_documents(
+        self, results: list[ContextGroundingQueryResponse]
+    ) -> list[tuple[Document, float]]:
+        """Convert API results to Document objects with scores.
+
+        Args:
+            results: List of ContextGroundingQueryResponse objects
+
+        Returns:
+            List of tuples containing (Document, score)
+        """
+        docs_with_scores = []
+
+        for result in results:
+            # Create metadata from result fields
+            metadata = {}
+
+            # Add string fields with proper defaults
+            if result.source:
+                metadata["source"] = str(result.source)
+            if result.reference:
+                metadata["reference"] = str(result.reference)
+            if result.page_number:
+                metadata["page_number"] = str(result.page_number)
+            if result.source_document_id:
+                metadata["source_document_id"] = str(result.source_document_id)
+            if result.caption:
+                metadata["caption"] = str(result.caption)
+
+            # Add any operation metadata if available
+            if result.metadata:
+                if result.metadata.operation_id:
+                    metadata["operation_id"] = str(result.metadata.operation_id)
+                if result.metadata.strategy:
+                    metadata["strategy"] = str(result.metadata.strategy)
+
+            # Create a Document with the content and metadata
+            doc = Document(
+                page_content=result.content or "",
+                metadata=metadata,
+            )
+
+            # Convert score to distance (1 - score)
+            score = 1.0 - float(result.score or 0.0)
+
+            docs_with_scores.append((doc, score))
+
+        return docs_with_scores
+
     @classmethod
+    @override
     def from_texts(
-        cls: type[VST],
+        cls,
         texts: list[str],
         embedding: Embeddings,
-        metadatas: Optional[list[dict[str, Any]]] = None,
+        metadatas: list[dict] | None = None,
         **kwargs: Any,
-    ) -> VST:
+    ) -> Self:
         """This method is required by the VectorStore abstract class, but is not supported
         by ContextGroundingVectorStore which is read-only.
 
         Raises:
             NotImplementedError: This method is not supported by ContextGroundingVectorStore
         """
         raise NotImplementedError(
-            "ContextGroundingVectorStore is a read-only wrapper for UiPath Context Grounding. "
-            "Creating a vector store from texts is not supported."
+            "ContextGroundingVectorStore is a read-only wrapper for UiPath Context Grounding."
         )
 
-    # Other required methods with minimal implementation to satisfy the interface
+    @override
     def add_texts(
-        self,
-        texts: Iterable[str],
-        metadatas: Optional[list[dict[str, Any]]] = None,
-        **kwargs: Any,
+        self, texts: Iterable[str], metadatas: list[dict] | None = None, **kwargs: Any
     ) -> list[str]:
         """Not implemented for ContextGroundingVectorStore as this is a read-only wrapper."""
         raise NotImplementedError(
             "ContextGroundingVectorStore is a read-only wrapper for UiPath Context Grounding."
-        )
+        )  # type: ignore[unreachable]
 
-    def delete(self, ids: Optional[list[str]] = None, **kwargs: Any) -> Optional[bool]:
+    @override
+    def delete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None:
         """Not implemented for ContextGroundingVectorStore as this is a read-only wrapper."""
         raise NotImplementedError(
             "ContextGroundingVectorStore is a read-only wrapper for UiPath Context Grounding."