Skip to content

feat: refactored ContextGroundingVectorStore to improve type safety #241

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 29, 2025
Merged

feat: refactored ContextGroundingVectorStore to improve type safety #241

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
198 changes: 89 additions & 109 deletions src/uipath_langchain/vectorstores/context_grounding_vectorstore.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,66 +2,50 @@
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):
"""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.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
Expand All @@ -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)
Expand All @@ -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]]:
Expand All @@ -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
Expand All @@ -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]:
Expand All @@ -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]:
Expand All @@ -231,38 +161,88 @@ 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.

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,
metadatas: list[dict[str, Any]] | 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."
)

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."
Expand Down