llama-cloud 0.1.41__py3-none-any.whl → 1.0.0b4__py3-none-any.whl

llama_cloud/lib/index/composite_retriever.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+from typing import Any, List, Optional
+from typing_extensions import override
+
+import httpx
+from llama_index.core.schema import TextNode, QueryBundle, NodeWithScore
+from llama_index.core.constants import DEFAULT_PROJECT_NAME
+from llama_index.core.base.base_retriever import BaseRetriever
+
+from llama_cloud import LlamaCloud, AsyncLlamaCloud, omit
+from llama_cloud.types import (
+    Retriever,
+    ReRankConfigParam,
+    RetrieverPipeline,
+    PresetRetrievalParams,
+    CompositeRetrievalMode,
+    RetrieverPipelineParam,
+)
+from llama_cloud.types.composite_retrieval_result import Node
+
+from .base import LlamaCloudIndex
+from .api_utils import (
+    resolve_project,
+    resolve_retriever,
+    page_screenshot_nodes_to_node_with_score,
+    apage_screenshot_nodes_to_node_with_score,
+)
+
+
+class LlamaCloudCompositeRetriever(BaseRetriever):
+    def __init__(
+        self,
+        # retriever identifier
+        name: Optional[str] = None,
+        retriever_id: Optional[str] = None,
+        # project identifier
+        project_name: Optional[str] = DEFAULT_PROJECT_NAME,
+        project_id: Optional[str] = None,
+        organization_id: Optional[str] = None,
+        # creation options
+        create_if_not_exists: bool = False,
+        # connection params
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: int = 60,
+        httpx_client: Optional[httpx.Client] = None,
+        async_httpx_client: Optional[httpx.AsyncClient] = None,
+        # composite retrieval params
+        mode: Optional[CompositeRetrievalMode] = None,
+        rerank_top_n: Optional[int] = None,
+        rerank_config: Optional[ReRankConfigParam] = None,
+        persisted: Optional[bool] = True,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the Composite Retriever."""
+        # initialize clients
+        self._client = LlamaCloud(api_key=api_key, base_url=base_url, http_client=httpx_client, timeout=timeout)
+        self._aclient = AsyncLlamaCloud(
+            api_key=api_key, base_url=base_url, http_client=async_httpx_client, timeout=timeout
+        )
+
+        self.project = resolve_project(self._client, project_name, project_id, organization_id)
+
+        self.name = name
+        self.project_name = self.project.name
+        self._persisted = persisted
+
+        retriever = resolve_retriever(self._client, self.project, name, retriever_id, persisted)
+
+        if retriever is None and persisted and self.name is not None:
+            if create_if_not_exists:
+                retriever = self._client.retrievers.upsert(
+                    project_id=self.project.id,
+                    name=self.name,
+                    pipelines=[],
+                )
+            else:
+                raise ValueError(f"Retriever with name '{self.name}' does not exist in project.")
+
+        if retriever is None:
+            raise ValueError("Failed to resolve retriever")
+        self.retriever = retriever
+
+        # composite retrieval params
+        self._mode = mode if mode is not None else omit
+        self._rerank_top_n = rerank_top_n if rerank_top_n is not None else omit
+        self._rerank_config = rerank_config if rerank_config is not None else omit
+
+        super().__init__(  # type: ignore
+            callback_manager=kwargs.get("callback_manager"),
+            verbose=kwargs.get("verbose", False),
+        )
+
+    @property
+    def retriever_pipelines(self) -> List[RetrieverPipeline]:
+        return self.retriever.pipelines or []
+
+    def update_retriever_pipelines(self, pipelines: List[RetrieverPipeline]) -> Retriever:
+        if self._persisted:
+            self.retriever = self._client.retrievers.update(
+                retriever_id=self.retriever.id,
+                pipelines=[RetrieverPipelineParam(**pipeline.model_dump()) for pipeline in pipelines],  # type: ignore [typeddict-item]
+            )
+        else:
+            # Update in-memory retriever for non-persisted case using copy
+            self.retriever = self.retriever.model_copy(update={"pipelines": pipelines})
+
+        return self.retriever
+
+    def add_index(
+        self,
+        index: LlamaCloudIndex,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        preset_retrieval_parameters: Optional[PresetRetrievalParams] = None,
+    ) -> Retriever:
+        name = name or index.name
+        preset_retrieval_parameters = preset_retrieval_parameters or index.pipeline.preset_retrieval_parameters
+        retriever_pipeline = RetrieverPipeline(
+            pipeline_id=index.id,
+            name=name,
+            description=description,
+            preset_retrieval_parameters=preset_retrieval_parameters,
+        )
+        current_retriever_pipelines_by_name = {pipeline.name: pipeline for pipeline in (self.retriever_pipelines or [])}
+        current_retriever_pipelines_by_name[retriever_pipeline.name] = retriever_pipeline
+        return self.update_retriever_pipelines(list(current_retriever_pipelines_by_name.values()))
+
+    def remove_index(self, name: str) -> bool:
+        current_retriever_pipeline_names = self.retriever.pipelines or []  # type: ignore [union-attr]
+        new_retriever_pipelines = [pipeline for pipeline in current_retriever_pipeline_names if pipeline.name != name]
+        if len(new_retriever_pipelines) == len(current_retriever_pipeline_names):
+            return False
+        self.update_retriever_pipelines(new_retriever_pipelines)
+        return True
+
+    async def aupdate_retriever_pipelines(self, pipelines: List[RetrieverPipeline]) -> Retriever:
+        if self._persisted:
+            self.retriever = await self._aclient.retrievers.update(
+                retriever_id=self.retriever.id,
+                pipelines=pipelines,  # type: ignore [arg-type]
+            )
+        else:
+            # Update in-memory retriever for non-persisted case using copy
+            self.retriever = self.retriever.copy(update={"pipelines": pipelines})  # type: ignore [union-attr]
+        return self.retriever
+
+    async def async_add_index(
+        self,
+        index: LlamaCloudIndex,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        preset_retrieval_parameters: Optional[PresetRetrievalParams] = None,
+    ) -> Retriever:
+        name = name or index.name
+        preset_retrieval_parameters = preset_retrieval_parameters or index.pipeline.preset_retrieval_parameters
+        retriever_pipeline = RetrieverPipeline(
+            pipeline_id=index.id,
+            name=name,
+            description=description,
+            preset_retrieval_parameters=preset_retrieval_parameters,
+        )
+        current_retriever_pipelines_by_name = {pipeline.name: pipeline for pipeline in (self.retriever_pipelines or [])}
+        current_retriever_pipelines_by_name[retriever_pipeline.name] = retriever_pipeline
+        return await self.aupdate_retriever_pipelines(list(current_retriever_pipelines_by_name.values()))
+
+    async def aremove_index(self, name: str) -> bool:
+        current_retriever_pipeline_names = self.retriever.pipelines or []  # type: ignore [union-attr]
+        new_retriever_pipelines = [pipeline for pipeline in current_retriever_pipeline_names if pipeline.name != name]
+        if len(new_retriever_pipelines) == len(current_retriever_pipeline_names):
+            return False
+        await self.aupdate_retriever_pipelines(new_retriever_pipelines)
+        return True
+
+    def _result_nodes_to_node_with_score(self, composite_retrieval_node: Node) -> NodeWithScore:
+        return NodeWithScore(
+            node=TextNode(
+                id=composite_retrieval_node.node.id,
+                text=composite_retrieval_node.node.text,
+                metadata=composite_retrieval_node.node.metadata,
+            ),
+            score=composite_retrieval_node.score,
+        )
+
+    @override
+    def _retrieve(
+        self,
+        query_bundle: QueryBundle,
+        mode: Optional[CompositeRetrievalMode] = None,
+        rerank_top_n: Optional[int] = None,
+        rerank_config: Optional[ReRankConfigParam] = None,
+    ) -> List[NodeWithScore]:
+        mode = mode if mode is not None else self._mode  # type: ignore
+
+        rerank_top_n = rerank_top_n if rerank_top_n is not None else self._rerank_top_n  # type: ignore
+        rerank_config = (
+            rerank_config if rerank_config is not None else self._rerank_config  # type: ignore
+        )
+
+        # Inject rerank_top_n into rerank_config if specified
+        if rerank_top_n is not None:
+            if rerank_config is None:
+                rerank_config = ReRankConfigParam(top_n=rerank_top_n)
+            else:
+                # Update existing rerank_config with top_n
+                rerank_config = ReRankConfigParam(top_n=rerank_top_n, type=rerank_config.get("type", "system_default"))
+
+        if self._persisted:
+            result = self._client.retrievers.retriever.search(
+                retriever_id=self.retriever.id,
+                mode=mode,  # type: ignore
+                rerank_config=rerank_config,  # type: ignore
+                query=query_bundle.query_str,
+            )
+        else:
+            result = self._client.retrievers.search(
+                project_id=self.project.id,
+                mode=mode,  # type: ignore
+                rerank_config=rerank_config,  # type: ignore
+                query=query_bundle.query_str,
+                pipelines=self.retriever.pipelines,  # type: ignore
+            )
+        node_w_scores = [self._result_nodes_to_node_with_score(node) for node in (result.nodes or [])]
+        image_nodes_w_scores = page_screenshot_nodes_to_node_with_score(
+            self._client, result.image_nodes, self.retriever.project_id
+        )
+        return sorted(node_w_scores + image_nodes_w_scores, key=lambda x: x.score or 1.0, reverse=True)
+
+    @override
+    async def _aretrieve(
+        self,
+        query_bundle: QueryBundle,
+        mode: Optional[CompositeRetrievalMode] = None,
+        rerank_top_n: Optional[int] = None,
+        rerank_config: Optional[ReRankConfigParam] = None,
+    ) -> List[NodeWithScore]:
+        mode = mode if mode is not None else self._mode  # type: ignore
+
+        rerank_top_n = rerank_top_n if rerank_top_n is not None else self._rerank_top_n  # type: ignore
+        rerank_config = (
+            rerank_config if rerank_config is not None else self._rerank_config  # type: ignore
+        )
+
+        # Inject rerank_top_n into rerank_config if specified
+        if rerank_top_n is not None:
+            if rerank_config is None:
+                rerank_config = ReRankConfigParam(top_n=rerank_top_n)
+            else:
+                # Update existing rerank_config with top_n
+                rerank_config = ReRankConfigParam(top_n=rerank_top_n, type=rerank_config.get("type", "system_default"))
+
+        if self._persisted:
+            result = await self._aclient.retrievers.retriever.search(
+                retriever_id=self.retriever.id,
+                mode=mode,  # type: ignore
+                rerank_config=rerank_config,  # type: ignore
+                query=query_bundle.query_str,
+            )
+        else:
+            result = await self._aclient.retrievers.search(
+                project_id=self.project.id,
+                mode=mode,  # type: ignore
+                rerank_config=rerank_config,  # type: ignore
+                query=query_bundle.query_str,
+                pipelines=self.retriever.pipelines,  # type: ignore [arg-type]
+            )
+        node_w_scores = [self._result_nodes_to_node_with_score(node) for node in result.nodes or []]
+        image_nodes_w_scores = await apage_screenshot_nodes_to_node_with_score(
+            self._aclient, result.image_nodes, self.retriever.project_id
+        )
+        return sorted(node_w_scores + image_nodes_w_scores, key=lambda x: x.score or 1.0, reverse=True)

llama_cloud/lib/index/retriever.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+from typing_extensions import override
+
+import httpx
+from llama_index.core.schema import TextNode, QueryBundle, NodeWithScore
+from llama_index.core.constants import DEFAULT_PROJECT_NAME
+from llama_index.core.bridge.pydantic import BaseModel
+from llama_index.core.base.base_retriever import BaseRetriever
+from llama_index.core.vector_stores.types import MetadataFilters
+
+from llama_cloud import LlamaCloud, AsyncLlamaCloud, omit
+from llama_cloud.types import MetadataFiltersParam
+from llama_cloud.types.metadata_filters_param import FilterMetadataFilter
+from llama_cloud.types.pipeline_retrieve_response import RetrievalNode
+
+from .api_utils import (
+    resolve_project_and_pipeline,
+    page_figure_nodes_to_node_with_score,
+    apage_figure_nodes_to_node_with_score,
+    page_screenshot_nodes_to_node_with_score,
+    apage_screenshot_nodes_to_node_with_score,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LlamaCloudRetriever(BaseRetriever):
+    def __init__(
+        self,
+        # index identifier
+        name: Optional[str] = None,
+        index_id: Optional[str] = None,  # alias for pipeline_id
+        id: Optional[str] = None,  # alias for pipeline_id
+        pipeline_id: Optional[str] = None,
+        # project identifier
+        project_name: Optional[str] = DEFAULT_PROJECT_NAME,
+        project_id: Optional[str] = None,
+        organization_id: Optional[str] = None,
+        # connection params
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: int = 60,
+        httpx_client: Optional[httpx.Client] = None,
+        async_httpx_client: Optional[httpx.AsyncClient] = None,
+        # retrieval params
+        dense_similarity_top_k: Optional[int] = None,
+        sparse_similarity_top_k: Optional[int] = None,
+        enable_reranking: Optional[bool] = None,
+        rerank_top_n: Optional[int] = None,
+        alpha: Optional[float] = None,
+        filters: Optional[MetadataFilters] = None,
+        retrieval_mode: Optional[str] = None,
+        files_top_k: Optional[int] = None,
+        retrieve_image_nodes: Optional[bool] = None,
+        retrieve_page_screenshot_nodes: Optional[bool] = None,
+        retrieve_page_figure_nodes: Optional[bool] = None,
+        search_filters_inference_schema: Optional[BaseModel] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the Platform Retriever."""
+        if sum([bool(id), bool(index_id), bool(pipeline_id), bool(name)]) != 1:
+            raise ValueError(
+                "Exactly one of `name`, `id`, `pipeline_id` or `index_id` must be provided to identify the index."
+            )
+
+        # initialize clients
+        self._httpx_client = httpx_client
+        self._async_httpx_client = async_httpx_client
+        self._client = LlamaCloud(
+            api_key=api_key,
+            base_url=base_url,
+            timeout=timeout,
+            http_client=httpx_client,
+        )
+        self._aclient = AsyncLlamaCloud(
+            api_key=api_key,
+            base_url=base_url,
+            timeout=timeout,
+            http_client=async_httpx_client,
+        )
+
+        pipeline_id = id or index_id or pipeline_id
+        self.project, self.pipeline = resolve_project_and_pipeline(
+            self._client, name, pipeline_id, project_name, project_id, organization_id
+        )
+        self.name = self.pipeline.name
+        self.project_name = self.project.name
+
+        # retrieval params
+        self._dense_similarity_top_k = dense_similarity_top_k if dense_similarity_top_k is not None else None
+        self._sparse_similarity_top_k = sparse_similarity_top_k if sparse_similarity_top_k is not None else None
+        self._enable_reranking = enable_reranking if enable_reranking is not None else None
+        self._rerank_top_n = rerank_top_n if rerank_top_n is not None else None
+        self._alpha = alpha if alpha is not None else None
+
+        # Convert filters to MetadataFiltersParam
+        self._filters: Optional[MetadataFiltersParam] = None
+        if filters:
+            self._filters = MetadataFiltersParam(
+                filters=[FilterMetadataFilter(**f.model_dump()) for f in filters.filters]  # type: ignore[typeddict-item]
+            )
+
+        self._retrieval_mode = retrieval_mode if retrieval_mode is not None else None
+        self._files_top_k = files_top_k if files_top_k is not None else None
+        if retrieve_image_nodes is not None:
+            logger.warning(
+                "The `retrieve_image_nodes` parameter is deprecated. "
+                "Use `retrieve_page_screenshot_nodes` and `retrieve_page_figure_nodes` instead."
+            )
+        if retrieve_image_nodes:
+            if retrieve_page_screenshot_nodes is False or retrieve_page_figure_nodes is False:
+                raise ValueError(
+                    "If `retrieve_image_nodes` is set to True, "
+                    "both `retrieve_page_screenshot_nodes` and `retrieve_page_figure_nodes` must also be set to True or omitted."
+                )
+            retrieve_page_screenshot_nodes = True
+            retrieve_page_figure_nodes = True
+        self._retrieve_page_screenshot_nodes = (
+            retrieve_page_screenshot_nodes if retrieve_page_screenshot_nodes is not None else False
+        )
+        self._retrieve_page_figure_nodes = (
+            retrieve_page_figure_nodes if retrieve_page_figure_nodes is not None else False
+        )
+        self._search_filters_inference_schema = search_filters_inference_schema
+
+        super().__init__(  # type: ignore
+            callback_manager=kwargs.get("callback_manager"),
+            verbose=kwargs.get("verbose", False),
+        )
+
+    def _result_nodes_to_node_with_score(
+        self, result_nodes: List[RetrievalNode], metadata: Optional[Dict[str, str]] = None
+    ) -> List[NodeWithScore]:
+        nodes: List[NodeWithScore] = []
+        for res in result_nodes:
+            text_node = TextNode.model_validate(res.node.model_dump())
+            text_node.metadata.update(metadata or {})
+            nodes.append(NodeWithScore(node=text_node, score=res.score))
+
+        return nodes
+
+    @override
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Retrieve from the platform."""
+        search_filters_inference_schema: Optional[Dict[str, Any]] = None
+        if self._search_filters_inference_schema is not None:
+            search_filters_inference_schema = self._search_filters_inference_schema.model_json_schema()
+
+        results = self._client.pipelines.retrieve(
+            pipeline_id=self.pipeline.id,
+            query=query_bundle.query_str,
+            alpha=self._alpha or omit,
+            dense_similarity_top_k=self._dense_similarity_top_k or omit,
+            enable_reranking=self._enable_reranking or omit,
+            files_top_k=self._files_top_k or omit,
+            rerank_top_n=self._rerank_top_n or omit,
+            retrieval_mode=self._retrieval_mode or "chunks",  # type: ignore
+            retrieve_page_figure_nodes=self._retrieve_page_figure_nodes or omit,
+            retrieve_page_screenshot_nodes=self._retrieve_page_screenshot_nodes or omit,
+            search_filters=self._filters or omit,
+            search_filters_inference_schema=search_filters_inference_schema or omit,  # type: ignore
+            sparse_similarity_top_k=self._sparse_similarity_top_k or omit,
+        )
+
+        result_nodes = self._result_nodes_to_node_with_score(results.retrieval_nodes, metadata=results.metadata)
+        if self._retrieve_page_screenshot_nodes:
+            result_nodes.extend(
+                page_screenshot_nodes_to_node_with_score(
+                    self._client,
+                    results.image_nodes,
+                    self.project.id,
+                    metadata=results.metadata,
+                )
+            )
+        if self._retrieve_page_figure_nodes:
+            result_nodes.extend(
+                page_figure_nodes_to_node_with_score(
+                    self._client,
+                    results.page_figure_nodes,
+                    self.project.id,
+                    metadata=results.metadata,
+                )
+            )
+
+        return result_nodes
+
+    @override
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Asynchronously retrieve from the platform."""
+        search_filters_inference_schema: Optional[Dict[str, Any]] = None
+        if self._search_filters_inference_schema is not None:
+            search_filters_inference_schema = self._search_filters_inference_schema.model_json_schema()
+
+        results = await self._aclient.pipelines.retrieve(
+            pipeline_id=self.pipeline.id,
+            query=query_bundle.query_str,
+            alpha=self._alpha or omit,
+            dense_similarity_top_k=self._dense_similarity_top_k or omit,
+            enable_reranking=self._enable_reranking or omit,
+            files_top_k=self._files_top_k or omit,
+            rerank_top_n=self._rerank_top_n or omit,
+            retrieval_mode=self._retrieval_mode or "chunks",  # type: ignore
+            retrieve_page_figure_nodes=self._retrieve_page_figure_nodes or omit,
+            retrieve_page_screenshot_nodes=self._retrieve_page_screenshot_nodes or omit,
+            search_filters=self._filters or omit,
+            search_filters_inference_schema=search_filters_inference_schema or omit,  # type: ignore
+            sparse_similarity_top_k=self._sparse_similarity_top_k or omit,
+        )
+
+        result_nodes = self._result_nodes_to_node_with_score(results.retrieval_nodes, metadata=results.metadata)
+        if self._retrieve_page_screenshot_nodes:
+            result_nodes.extend(
+                await apage_screenshot_nodes_to_node_with_score(
+                    self._aclient,
+                    results.image_nodes,
+                    self.project.id,
+                    metadata=results.metadata,
+                )
+            )
+        if self._retrieve_page_figure_nodes:
+            result_nodes.extend(
+                await apage_figure_nodes_to_node_with_score(
+                    self._aclient,
+                    results.page_figure_nodes,
+                    self.project.id,
+                    metadata=results.metadata,
+                )
+            )
+
+        return result_nodes