feat: add non persisted composite retrieval (#18908)

Author: EmanuelCampos

llama-index-integrations/indices/llama-index-indices-managed-llama-cloud/llama_index/indices/managed/llama_cloud/composite_retriever.py

@@ -45,6 +45,7 @@ def __init__(
         # composite retrieval params
         mode: Optional[CompositeRetrievalMode] = None,
         rerank_top_n: Optional[int] = None,
+        persisted: Optional[bool] = True,
         **kwargs: Any,
     ) -> None:
         """Initialize the Composite Retriever."""
@@ -58,13 +59,15 @@ def __init__(
             self._client, project_name, project_id, organization_id
         )
 
-        self.retriever = resolve_retriever(
-            self._client, self.project, name, retriever_id
-        )
         self.name = name
         self.project_name = self.project.name
+        self._persisted = persisted
+
+        self.retriever = resolve_retriever(
+            self._client, self.project, name, retriever_id, persisted
+        )
 
-        if self.retriever is None:
+        if self.retriever is None and persisted:
             if create_if_not_exists:
                 self.retriever = self._client.retrievers.upsert_retriever(
                     project_id=self.project.id,
@@ -91,9 +94,13 @@ def retriever_pipelines(self) -> List[RetrieverPipeline]:
     def update_retriever_pipelines(
         self, pipelines: List[RetrieverPipeline]
     ) -> Retriever:
-        self.retriever = self._client.retrievers.update_retriever(
-            self.retriever.id, pipelines=pipelines
-        )
+        if self._persisted:
+            self.retriever = self._client.retrievers.update_retriever(
+                self.retriever.id, pipelines=pipelines
+            )
+        else:
+            # Update in-memory retriever for non-persisted case using copy
+            self.retriever = self.retriever.copy(update={"pipelines": pipelines})
         return self.retriever
 
     def add_index(
@@ -138,9 +145,13 @@ def remove_index(self, name: str) -> bool:
     async def aupdate_retriever_pipelines(
         self, pipelines: List[RetrieverPipeline]
     ) -> Retriever:
-        self.retriever = await self._aclient.retrievers.update_retriever(
-            self.retriever.id, pipelines=pipelines
-        )
+        if self._persisted:
+            self.retriever = await self._aclient.retrievers.update_retriever(
+                self.retriever.id, pipelines=pipelines
+            )
+        else:
+            # Update in-memory retriever for non-persisted case using copy
+            self.retriever = self.retriever.copy(update={"pipelines": pipelines})
         return self.retriever
 
     async def async_add_index(
@@ -202,12 +213,21 @@ def _retrieve(
     ) -> List[NodeWithScore]:
         mode = mode if mode is not None else self._mode
         rerank_top_n = rerank_top_n if rerank_top_n is not None else self._rerank_top_n
-        result = self._client.retrievers.retrieve(
-            self.retriever.id,
-            mode=mode,
-            rerank_top_n=rerank_top_n,
-            query=query_bundle.query_str,
-        )
+        if self._persisted:
+            result = self._client.retrievers.retrieve(
+                self.retriever.id,
+                mode=mode,
+                rerank_top_n=rerank_top_n,
+                query=query_bundle.query_str,
+            )
+        else:
+            result = self._client.retrievers.direct_retrieve(
+                project_id=self.project.id,
+                mode=mode,
+                rerank_top_n=rerank_top_n,
+                query=query_bundle.query_str,
+                pipelines=self.retriever.pipelines,
+            )
         node_w_scores = [
             self._result_nodes_to_node_with_score(node) for node in result.nodes
         ]
@@ -226,12 +246,21 @@ async def _aretrieve(
     ) -> List[NodeWithScore]:
         mode = mode if mode is not None else self._mode
         rerank_top_n = rerank_top_n if rerank_top_n is not None else self._rerank_top_n
-        result = await self._aclient.retrievers.retrieve(
-            self.retriever.id,
-            mode=mode,
-            rerank_top_n=rerank_top_n,
-            query=query_bundle.query_str,
-        )
+        if self._persisted:
+            result = await self._aclient.retrievers.retrieve(
+                self.retriever.id,
+                mode=mode,
+                rerank_top_n=rerank_top_n,
+                query=query_bundle.query_str,
+            )
+        else:
+            result = await self._aclient.retrievers.direct_retrieve(
+                project_id=self.project.id,
+                mode=mode,
+                rerank_top_n=rerank_top_n,
+                query=query_bundle.query_str,
+                pipelines=self.retriever.pipelines,
+            )
         node_w_scores = [
             self._result_nodes_to_node_with_score(node) for node in result.nodes
         ]

llama-index-integrations/indices/llama-index-indices-managed-llama-cloud/pyproject.toml

@@ -28,7 +28,7 @@ dev = [
 
 [project]
 name = "llama-index-indices-managed-llama-cloud"
-version = "0.7.1"
+version = "0.7.2"
 description = "llama-index indices llama-cloud integration"
 authors = [{name = "Logan Markewich", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"