From 95e0385c332692ef198bdb9865b1fb1129f9561e Mon Sep 17 00:00:00 2001 From: Adrian Badea Date: Tue, 28 Oct 2025 12:33:57 +0200 Subject: [PATCH] feat: refactored ContextGroundingVectorStore to improve type safety --- pyproject.toml | 2 +- .../context_grounding_vectorstore.py | 198 ++++++++---------- 2 files changed, 90 insertions(+), 110 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 70a7bf90..197716e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.0.146" +version = "0.0.147" description = "UiPath Langchain" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.10" diff --git a/src/uipath_langchain/vectorstores/context_grounding_vectorstore.py b/src/uipath_langchain/vectorstores/context_grounding_vectorstore.py index 396434e4..2ab07686 100644 --- a/src/uipath_langchain/vectorstores/context_grounding_vectorstore.py +++ b/src/uipath_langchain/vectorstores/context_grounding_vectorstore.py @@ -2,23 +2,17 @@ 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 from langchain_core.documents import Document from langchain_core.embeddings import Embeddings from langchain_core.vectorstores import VectorStore +from typing_extensions import override from uipath import UiPath - -VST = TypeVar("VST", bound="ContextGroundingVectorStore") +from uipath.models.context_grounding import ContextGroundingQueryResponse class ContextGroundingVectorStore(VectorStore): @@ -26,42 +20,32 @@ class ContextGroundingVectorStore(VectorStore): 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.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 +54,24 @@ 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( - 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, + # 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, ) + ) - 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 @@ -201,6 +129,7 @@ async def asimilarity_search_with_relevance_scores( ) ] + @override def similarity_search( self, query: str, k: int = 4, **kwargs: Any ) -> list[Document]: @@ -216,6 +145,7 @@ def similarity_search( docs_and_scores = self.similarity_search_with_score(query, k, **kwargs) return [doc for doc, _ in docs_and_scores] + @override async def asimilarity_search( self, query: str, k: int = 4, **kwargs: Any ) -> list[Document]: @@ -231,14 +161,64 @@ 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[str, Any]] | 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. @@ -246,15 +226,14 @@ def from_texts( 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, + metadatas: list[dict[str, Any]] | None = None, **kwargs: Any, ) -> list[str]: """Not implemented for ContextGroundingVectorStore as this is a read-only wrapper.""" @@ -262,7 +241,8 @@ def add_texts( "ContextGroundingVectorStore is a read-only wrapper for UiPath Context Grounding." ) - 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."