feat: add retriever and query engine implementations

Author: radu-mocanu

samples/simple-hitl-agent/uv.lock

@@ -0,0 +1,3 @@
+from .context_grounding_query_engine import ContextGroundingQueryEngine
+
+__all__ = ["ContextGroundingQueryEngine"]

src/uipath_llamaindex/query_engines/__init__.py

@@ -0,0 +1,42 @@
+from typing import Optional
+
+from llama_index.core.query_engine import CustomQueryEngine
+from llama_index.core.response_synthesizers import BaseSynthesizer
+from uipath import UiPath
+
+from uipath_llamaindex.retrievers import ContextGroundingRetriever
+
+
+class ContextGroundingQueryEngine(CustomQueryEngine):
+    """RAG Query Engine."""
+
+    def __init__(
+        self,
+        response_synthesizer: BaseSynthesizer,
+        index_name: str,
+        folder_path: Optional[str] = None,
+        folder_key: Optional[str] = None,
+        uipath: Optional[UiPath] = None,
+        number_of_results: Optional[int] = 10,
+        **kwargs,
+    ):
+        super().__init__()
+        self._retriever = ContextGroundingRetriever(
+            index_name=index_name,
+            folder_path=folder_path,
+            folder_key=folder_key,
+            number_of_results=number_of_results,
+            uipath=uipath,
+            **kwargs,
+        )
+        self._response_synthesizer = response_synthesizer
+
+    def custom_query(self, query_str: str):
+        nodes = self._retriever.retrieve(query_str)
+        response_obj = self._response_synthesizer.synthesize(query_str, nodes)
+        return response_obj
+
+    async def acustom_query(self, query_str: str):
+        nodes = await self._retriever.aretrieve(query_str)
+        response_obj = self._response_synthesizer.synthesize(query_str, nodes)
+        return response_obj

src/uipath_llamaindex/query_engines/context_grounding_query_engine.py

@@ -0,0 +1,3 @@
+from .context_grounding_retriever import ContextGroundingRetriever
+
+__all__ = ["ContextGroundingRetriever"]

src/uipath_llamaindex/retrievers/__init__.py

@@ -0,0 +1,63 @@
+from typing import List, Optional
+
+from llama_index.core.retrievers import (
+    BaseRetriever,
+)
+from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
+from uipath import UiPath
+from uipath.models import ContextGroundingQueryResponse
+
+
+class ContextGroundingRetriever(BaseRetriever):
+    def __init__(
+        self,
+        index_name: str,
+        folder_path: Optional[str] = None,
+        folder_key: Optional[str] = None,
+        uipath: Optional[UiPath] = None,
+        number_of_results: Optional[int] = 10,
+        **kwargs,
+    ):
+        super().__init__()
+        self._index_name = index_name
+        self._folder_path = folder_path
+        self._folder_key = folder_key
+        self._uipath = uipath or UiPath()
+        self._number_of_results = number_of_results
+        self._results: list[ContextGroundingQueryResponse] = []
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        self._results = self._uipath.context_grounding.search(
+            self._index_name,
+            query_bundle.query_str,
+            self._number_of_results,
+            folder_path=self._folder_path,
+            folder_key=self._folder_key,
+        )
+
+        return self._to_nodes_with_scores()
+
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        self._results = await self._uipath.context_grounding.search_async(
+            self._index_name,
+            query_bundle.query_str,
+            self._number_of_results,
+            folder_path=self._folder_path,
+            folder_key=self._folder_key,
+        )
+
+        return self._to_nodes_with_scores()
+
+    def _to_nodes_with_scores(self) -> List[NodeWithScore]:
+        nodes_with_scores = []
+        for chunk in self._results:
+            node = TextNode(
+                text=chunk.content,
+                metadata={
+                    "source_document_id": chunk.source_document_id,
+                    "source": chunk.source,
+                    "page_number": chunk.page_number,
+                },
+            )
+            nodes_with_scores.append(NodeWithScore(node=node, score=chunk.score))
+        return nodes_with_scores