Merge branch 'main' into clelia/ts-refactoring

README.md

@@ -11,6 +11,7 @@ This includes:
 - [LlamaParse](./parse.md) - A GenAI-native document parser that can parse complex document data for any downstream LLM use case (Agents, RAG, data processing, etc.).
 - [LlamaReport (beta/invite-only)](./report.md) - A prebuilt agentic report builder that can be used to build reports from a variety of data sources.
 - [LlamaExtract](./extract.md) - A prebuilt agentic data extractor that can be used to transform data into a structured JSON representation.
+- [LlamaCloud Index](./index.md) - A widely customizable and fully automated document ingestion pipeline that also serves retrieval purposes.
 
 ## Getting Started
 
@@ -25,11 +26,19 @@ Then, get your API key from [LlamaCloud](https://cloud.llamaindex.ai/).
 Then, you can use the services in your code:
 
 ```python
-from llama_cloud_services import LlamaParse, LlamaReport, LlamaExtract
+from llama_cloud_services import (
+    LlamaParse,
+    LlamaReport,
+    LlamaExtract,
+    LlamaCloudIndex,
+)
 
 parser = LlamaParse(api_key="YOUR_API_KEY")
 report = LlamaReport(api_key="YOUR_API_KEY")
 extract = LlamaExtract(api_key="YOUR_API_KEY")
+index = LlamaCloudIndex(
+    "my_first_index", project_name="default", api_key="YOUR_API_KEY"
+)
 ```
 
 See the quickstart guides for each service for more information:
@@ -37,6 +46,7 @@ See the quickstart guides for each service for more information:
 - [LlamaParse](./parse.md)
 - [LlamaReport (beta/invite-only)](./report.md)
 - [LlamaExtract](./extract.md)
+- [LlamaCloud Index](./index.md)
 
 ## Switch to EU SaaS 🇪🇺
 
@@ -55,6 +65,12 @@ from llama_cloud_services import (
 parser = LlamaParse(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
 report = LlamaReport(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
 extract = LlamaExtract(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
+index = LlamaCloudIndex(
+    "my_first_index",
+    project_name="default",
+    api_key="YOUR_API_KEY",
+    base_url=EU_BASE_URL,
+)
 ```
 
 ## Documentation

index.md

@@ -0,0 +1,86 @@
+# LlamaCloud Index + Retriever
+
+LlamaCloud is a new generation of managed parsing, ingestion, and retrieval services, designed to bring production-grade context-augmentation to your LLM and RAG applications.
+
+Currently, LlamaCloud supports
+
+- Managed Ingestion API, handling parsing and document management
+- Managed Retrieval API, configuring optimal retrieval for your RAG system
+
+## Access
+
+We are opening up a private beta to a limited set of enterprise partners for the managed ingestion and retrieval API. If you’re interested in centralizing your data pipelines and spending more time working on your actual RAG use cases, come [talk to us.](https://www.llamaindex.ai/contact)
+
+If you have access to LlamaCloud, you can visit [LlamaCloud](https://cloud.llamaindex.ai) to sign in and get an API key.
+
+## Setup
+
+First, make sure you have the latest LlamaIndex version installed.
+
+```
+pip uninstall llama-index  # run this if upgrading from v0.9.x or older
+pip install -U llama-index --upgrade --no-cache-dir --force-reinstall
+```
+
+The `llama-index-indices-managed-llama-cloud` package is included with the above install, but you can also install directly
+
+```
+pip install -U llama-index-indices-managed-llama-cloud
+```
+
+## Usage
+
+You can create an index on LlamaCloud using the following code. By default, new indexes use managed embeddings (OpenAI text-embedding-3-small, 1536 dimensions, 1 credit/page):
+
+```python
+import os
+
+os.environ[
+    "LLAMA_CLOUD_API_KEY"
+] = "llx-..."  # can provide API-key in env or in the constructor later on
+
+from llama_index.core import SimpleDirectoryReader
+from llama_cloud_services import LlamaCloudIndex
+
+# create a new index (uses managed embeddings by default)
+index = LlamaCloudIndex.from_documents(
+    documents,
+    "my_first_index",
+    project_name="default",
+    api_key="llx-...",
+    verbose=True,
+)
+
+# connect to an existing index
+index = LlamaCloudIndex("my_first_index", project_name="default")
+```
+
+You can also configure a retriever for managed retrieval:
+
+```python
+# from the existing index
+index.as_retriever()
+
+# from scratch
+from llama_index.indices.managed.llama_cloud import LlamaCloudRetriever
+
+retriever = LlamaCloudRetriever("my_first_index", project_name="default")
+```
+
+And of course, you can use other index shortcuts to get use out of your new managed index:
+
+```python
+query_engine = index.as_query_engine(llm=llm)
+
+chat_engine = index.as_chat_engine(llm=llm)
+```
+
+## Retriever Settings
+
+A full list of retriever settings/kwargs is below:
+
+- `dense_similarity_top_k`: Optional[int] -- If greater than 0, retrieve `k` nodes using dense retrieval
+- `sparse_similarity_top_k`: Optional[int] -- If greater than 0, retrieve `k` nodes using sparse retrieval
+- `enable_reranking`: Optional[bool] -- Whether to enable reranking or not. Sacrifices some speed for accuracy
+- `rerank_top_n`: Optional[int] -- The number of nodes to return after reranking initial retrieval results
+- `alpha` Optional[float] -- The weighting between dense and sparse retrieval. 1 = Full dense retrieval, 0 = Full sparse retrieval.

py/README.md

@@ -11,6 +11,7 @@ This includes:
 - [LlamaParse](./parse.md) - A GenAI-native document parser that can parse complex document data for any downstream LLM use case (Agents, RAG, data processing, etc.).
 - [LlamaReport (beta/invite-only)](./report.md) - A prebuilt agentic report builder that can be used to build reports from a variety of data sources.
 - [LlamaExtract](./extract.md) - A prebuilt agentic data extractor that can be used to transform data into a structured JSON representation.
+- [LlamaCloud Index](./index.md) - A widely customizable and fully automated document ingestion pipeline that also serves retrieval purposes.
 
 ## Getting Started
 
@@ -25,11 +26,20 @@ Then, get your API key from [LlamaCloud](https://cloud.llamaindex.ai/).
 Then, you can use the services in your code:
 
 ```python
+from llama_cloud_services import (
+    LlamaParse,
+    LlamaReport,
+    LlamaExtract,
+    LlamaCloudIndex,
+)
 from llama_cloud_services import LlamaParse, LlamaReport, LlamaExtract
 
 parser = LlamaParse(api_key="YOUR_API_KEY")
 report = LlamaReport(api_key="YOUR_API_KEY")
 extract = LlamaExtract(api_key="YOUR_API_KEY")
+index = LlamaCloudIndex(
+    "my_first_index", project_name="default", api_key="YOUR_API_KEY"
+)
 ```
 
 See the quickstart guides for each service for more information:
@@ -37,6 +47,7 @@ See the quickstart guides for each service for more information:
 - [LlamaParse](./parse.md)
 - [LlamaReport (beta/invite-only)](./report.md)
 - [LlamaExtract](./extract.md)
+- [LlamaCloud Index](./index.md)
 
 ## Switch to EU SaaS 🇪🇺
 
@@ -55,6 +66,12 @@ from llama_cloud_services import (
 parser = LlamaParse(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
 report = LlamaReport(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
 extract = LlamaExtract(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
+index = LlamaCloudIndex(
+    "my_first_index",
+    project_name="default",
+    api_key="YOUR_API_KEY",
+    base_url=EU_BASE_URL,
+)
 ```
 
 ## Documentation

py/llama_cloud_services/__init__.py

@@ -2,6 +2,11 @@ from llama_cloud_services.parse import LlamaParse
 from llama_cloud_services.report import ReportClient, LlamaReport
 from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
 from llama_cloud_services.constants import EU_BASE_URL
+from llama_cloud_services.index import (
+    LlamaCloudCompositeRetriever,
+    LlamaCloudIndex,
+    LlamaCloudRetriever,
+)
 
 __all__ = [
     "LlamaParse",
@@ -10,4 +15,7 @@ __all__ = [
     "LlamaExtract",
     "ExtractionAgent",
     "EU_BASE_URL",
+    "LlamaCloudIndex",
+    "LlamaCloudRetriever",
+    "LlamaCloudCompositeRetriever",
 ]

py/llama_cloud_services/index/__init__.py

@@ -0,0 +1,11 @@
+from .base import LlamaCloudIndex
+from .retriever import LlamaCloudRetriever
+from .composite_retriever import (
+    LlamaCloudCompositeRetriever,
+)
+
+__all__ = [
+    "LlamaCloudIndex",
+    "LlamaCloudRetriever",
+    "LlamaCloudCompositeRetriever",
+]

py/llama_cloud_services/index/api_utils.py

@@ -0,0 +1,422 @@
+import base64
+import urllib.parse
+import uuid
+from httpx import Request, HTTPStatusError
+from tenacity import (
+    retry,
+    wait_exponential_jitter,
+    stop_after_attempt,
+    retry_if_exception,
+)
+from typing import Any, Optional, Tuple, List, Union, Dict, Callable
+
+from llama_index.core.async_utils import run_jobs
+from llama_index.core.schema import NodeWithScore, ImageNode
+from llama_cloud import (
+    AutoTransformConfig,
+    PageFigureNodeWithScore,
+    PageScreenshotNodeWithScore,
+    Pipeline,
+    PipelineCreateTransformConfig,
+    PipelineType,
+    Project,
+    Retriever,
+)
+from llama_cloud.core import remove_none_from_dict
+from llama_cloud.client import LlamaCloud, AsyncLlamaCloud
+from llama_cloud.core.api_error import ApiError
+
+
+def is_retryable_http_error(exception: Exception) -> bool:
+    # Retry for ApiError with 5xx status codes
+    if isinstance(exception, ApiError):
+        return 500 <= exception.status_code < 600
+    # Also retry for HTTPError with 5xx status codes
+    elif isinstance(exception, HTTPStatusError):
+        return 500 <= exception.response.status_code < 600
+    return False
+
+
+def retry_on_failure(func: Callable) -> Callable:
+    """Decorator to apply tenacity retry with exponential backoff and jitter for 5xx errors."""
+    return retry(
+        wait=wait_exponential_jitter(exp_base=2, max=10),
+        stop=stop_after_attempt(5),
+        retry=retry_if_exception(is_retryable_http_error),
+        reraise=True,
+    )(func)
+
+
+def default_transform_config() -> PipelineCreateTransformConfig:
+    return AutoTransformConfig()
+
+
+def resolve_retriever(
+    client: LlamaCloud,
+    project: Project,
+    retriever_name: Optional[str] = None,
+    retriever_id: Optional[str] = None,
+    persisted: bool = True,
+) -> Optional[Retriever]:
+    if not persisted:
+        return Retriever(
+            id=str(uuid.uuid4()),
+            project_id=project.id,
+            name=retriever_name,
+            pipelines=[],
+        )
+    if retriever_id:
+        return client.retrievers.get_retriever(
+            retriever_id=retriever_id, project_id=project.id
+        )
+    elif retriever_name:
+        retrievers = client.retrievers.list_retrievers(
+            project_id=project.id, name=retriever_name
+        )
+        return next(
+            (retriever for retriever in retrievers if retriever.name == retriever_name),
+            None,
+        )
+    else:
+        return None
+
+
+def resolve_project(
+    client: LlamaCloud,
+    project_name: Optional[str],
+    project_id: Optional[str],
+    organization_id: Optional[str],
+) -> Project:
+    if project_id is not None:
+        return client.projects.get_project(project_id=project_id)
+    else:
+        projects = client.projects.list_projects(
+            project_name=project_name, organization_id=organization_id
+        )
+        if len(projects) == 0:
+            raise ValueError(f"No project found with name {project_name}")
+        elif len(projects) > 1:
+            raise ValueError(
+                f"Multiple projects found with name {project_name}. Please specify organization_id."
+            )
+        return projects[0]
+
+
+def resolve_pipeline(
+    client: LlamaCloud,
+    pipeline_id: Optional[str],
+    project: Optional[Project],
+    pipeline_name: Optional[str],
+) -> Pipeline:
+    if pipeline_id is not None:
+        return client.pipelines.get_pipeline(pipeline_id=pipeline_id)
+    else:
+        pipelines = client.pipelines.search_pipelines(
+            project_id=project.id,  # type: ignore [union-attr]
+            pipeline_name=pipeline_name,
+            pipeline_type=PipelineType.MANAGED.value,  # type: ignore [union-attr]
+        )
+        if len(pipelines) == 0:
+            raise ValueError(
+                f"Unknown index name {pipeline_name}. Please confirm an index with this name exists."
+            )
+        elif len(pipelines) > 1:
+            raise ValueError(
+                f"Multiple pipelines found with name {pipeline_name} in project {project.name}"  # type: ignore [union-attr]
+            )
+        return pipelines[0]
+
+
+def resolve_project_and_pipeline(
+    client: LlamaCloud,
+    pipeline_name: Optional[str],
+    pipeline_id: Optional[str],
+    project_name: Optional[str],
+    project_id: Optional[str],
+    organization_id: Optional[str],
+) -> Tuple[Project, Pipeline]:
+    # resolve pipeline by ID
+    if pipeline_id is not None:
+        pipeline = resolve_pipeline(
+            client, pipeline_id=pipeline_id, project=None, pipeline_name=None
+        )
+        project_id = pipeline.project_id
+
+    # resolve project
+    project = resolve_project(client, project_name, project_id, organization_id)
+
+    # resolve pipeline by name
+    if pipeline_id is None:
+        pipeline = resolve_pipeline(
+            client, pipeline_id=None, project=project, pipeline_name=pipeline_name
+        )
+
+    return project, pipeline
+
+
+def _build_get_page_screenshot_request(
+    client: Union[LlamaCloud, AsyncLlamaCloud],
+    file_id: str,
+    page_index: int,
+    project_id: str,
+) -> Request:
+    return client._client_wrapper.httpx_client.build_request(
+        "GET",
+        urllib.parse.urljoin(
+            f"{client._client_wrapper.get_base_url()}/",
+            f"api/v1/files/{file_id}/page_screenshots/{page_index}",
+        ),
+        params=remove_none_from_dict({"project_id": project_id}),
+        headers=client._client_wrapper.get_headers(),
+        timeout=60,
+    )
+
+
+def _build_get_page_figure_request(
+    client: Union[LlamaCloud, AsyncLlamaCloud],
+    file_id: str,
+    page_index: int,
+    figure_name: str,
+    project_id: str,
+) -> Request:
+    return client._client_wrapper.httpx_client.build_request(
+        "GET",
+        urllib.parse.urljoin(
+            f"{client._client_wrapper.get_base_url()}/",
+            f"api/v1/files/{file_id}/page-figures/{page_index}/{figure_name}",
+        ),
+        params=remove_none_from_dict({"project_id": project_id}),
+        headers=client._client_wrapper.get_headers(),
+        timeout=60,
+    )
+
+
+@retry_on_failure
+def get_page_screenshot(
+    client: LlamaCloud, file_id: str, page_index: int, project_id: str
+) -> str:
+    """Get the page screenshot."""
+    # TODO: this currently uses requests, should be replaced with the client
+    request = _build_get_page_screenshot_request(
+        client, file_id, page_index, project_id
+    )
+    _response = client._client_wrapper.httpx_client.send(request)
+    if 200 <= _response.status_code < 300:
+        return _response.content
+    else:
+        raise ApiError(status_code=_response.status_code, body=_response.text)
+
+
+@retry_on_failure
+def get_page_figure(
+    client: LlamaCloud, file_id: str, page_index: int, figure_name: str, project_id: str
+) -> str:
+    request = _build_get_page_figure_request(
+        client, file_id, page_index, figure_name, project_id
+    )
+    _response = client._client_wrapper.httpx_client.send(request)
+    if 200 <= _response.status_code < 300:
+        return _response.content
+    else:
+        raise ApiError(status_code=_response.status_code, body=_response.text)
+
+
+@retry_on_failure
+async def aget_page_screenshot(
+    client: AsyncLlamaCloud, file_id: str, page_index: int, project_id: str
+) -> str:
+    """Get the page screenshot (async)."""
+    request = _build_get_page_screenshot_request(
+        client, file_id, page_index, project_id
+    )
+    _response = await client._client_wrapper.httpx_client.send(request)
+    if 200 <= _response.status_code < 300:
+        return _response.content
+    else:
+        raise ApiError(status_code=_response.status_code, body=_response.text)
+
+
+@retry_on_failure
+async def aget_page_figure(
+    client: AsyncLlamaCloud,
+    file_id: str,
+    page_index: int,
+    figure_name: str,
+    project_id: str,
+) -> str:
+    request = _build_get_page_figure_request(
+        client, file_id, page_index, figure_name, project_id
+    )
+    _response = await client._client_wrapper.httpx_client.send(request)
+    if 200 <= _response.status_code < 300:
+        return _response.content
+    else:
+        raise ApiError(status_code=_response.status_code, body=_response.text)
+
+
+def page_screenshot_nodes_to_node_with_score(
+    client: LlamaCloud,
+    raw_image_nodes: Optional[List[PageScreenshotNodeWithScore]],
+    project_id: str,
+) -> List[NodeWithScore]:
+    if not raw_image_nodes:
+        return []
+
+    image_nodes = []
+    for raw_image_node in raw_image_nodes:
+        image_bytes = get_page_screenshot(
+            client=client,
+            file_id=raw_image_node.node.file_id,
+            page_index=raw_image_node.node.page_index,
+            project_id=project_id,
+        )
+        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+        image_node_metadata: Dict[str, Any] = {
+            **(raw_image_node.node.metadata or {}),
+            "file_id": raw_image_node.node.file_id,
+            "page_index": raw_image_node.node.page_index,
+        }
+        image_node_with_score = NodeWithScore(
+            node=ImageNode(image=image_base64, metadata=image_node_metadata),
+            score=raw_image_node.score,
+        )
+        image_nodes.append(image_node_with_score)
+
+    return image_nodes
+
+
+def image_nodes_to_node_with_score(
+    client: LlamaCloud,
+    raw_image_nodes: Optional[List[PageScreenshotNodeWithScore]],
+    project_id: str,
+) -> List[NodeWithScore]:
+    """
+    Legacy method to alias page_screenshot_nodes_to_node_with_score.
+    """
+    if not raw_image_nodes:
+        return []
+
+    return page_screenshot_nodes_to_node_with_score(
+        client=client, raw_image_nodes=raw_image_nodes, project_id=project_id
+    )
+
+
+def page_figure_nodes_to_node_with_score(
+    client: LlamaCloud,
+    raw_figure_nodes: Optional[List[PageFigureNodeWithScore]],
+    project_id: str,
+) -> List[NodeWithScore]:
+    if not raw_figure_nodes:
+        return []
+
+    figure_nodes = []
+    for raw_figure_node in raw_figure_nodes:
+        figure_bytes = get_page_figure(
+            client=client,
+            file_id=raw_figure_node.node.file_id,
+            page_index=raw_figure_node.node.page_index,
+            figure_name=raw_figure_node.node.figure_name,
+            project_id=project_id,
+        )
+        figure_base64 = base64.b64encode(figure_bytes).decode("utf-8")
+        figure_node_metadata: Dict[str, Any] = {
+            **(raw_figure_node.node.metadata or {}),
+            "file_id": raw_figure_node.node.file_id,
+            "page_index": raw_figure_node.node.page_index,
+            "figure_name": raw_figure_node.node.figure_name,
+        }
+        figure_node_with_score = NodeWithScore(
+            node=ImageNode(image=figure_base64, metadata=figure_node_metadata),
+            score=raw_figure_node.score,
+        )
+        figure_nodes.append(figure_node_with_score)
+    return figure_nodes
+
+
+async def apage_screenshot_nodes_to_node_with_score(
+    client: AsyncLlamaCloud,
+    raw_image_nodes: Optional[List[PageScreenshotNodeWithScore]],
+    project_id: str,
+) -> List[NodeWithScore]:
+    if not raw_image_nodes:
+        return []
+
+    image_nodes = []
+    tasks = [
+        aget_page_screenshot(
+            client=client,
+            file_id=raw_image_node.node.file_id,
+            page_index=raw_image_node.node.page_index,
+            project_id=project_id,
+        )
+        for raw_image_node in raw_image_nodes
+    ]
+
+    image_bytes_list = await run_jobs(tasks)
+    for image_bytes, raw_image_node in zip(image_bytes_list, raw_image_nodes):
+        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+        image_node_metadata: Dict[str, Any] = {
+            **(raw_image_node.node.metadata or {}),
+            "file_id": raw_image_node.node.file_id,
+            "page_index": raw_image_node.node.page_index,
+        }
+        image_node_with_score = NodeWithScore(
+            node=ImageNode(image=image_base64, metadata=image_node_metadata),
+            score=raw_image_node.score,
+        )
+        image_nodes.append(image_node_with_score)
+    return image_nodes
+
+
+async def aimage_nodes_to_node_with_score(
+    client: AsyncLlamaCloud,
+    raw_image_nodes: Optional[List[PageScreenshotNodeWithScore]],
+    project_id: str,
+) -> List[NodeWithScore]:
+    """
+    Legacy method to alias apage_screenshot_nodes_to_node_with_score.
+    """
+    if not raw_image_nodes:
+        return []
+
+    return await apage_screenshot_nodes_to_node_with_score(
+        client=client, raw_image_nodes=raw_image_nodes, project_id=project_id
+    )
+
+
+async def apage_figure_nodes_to_node_with_score(
+    client: AsyncLlamaCloud,
+    raw_figure_nodes: Optional[List[PageFigureNodeWithScore]],
+    project_id: str,
+) -> List[NodeWithScore]:
+    if not raw_figure_nodes:
+        return []
+
+    figure_nodes = []
+    tasks = [
+        aget_page_figure(
+            client=client,
+            file_id=raw_figure_node.node.file_id,
+            page_index=raw_figure_node.node.page_index,
+            figure_name=raw_figure_node.node.figure_name,
+            project_id=project_id,
+        )
+        for raw_figure_node in raw_figure_nodes
+    ]
+
+    figure_bytes_list = await run_jobs(tasks)
+    for figure_bytes, raw_figure_node in zip(figure_bytes_list, raw_figure_nodes):
+        figure_base64 = base64.b64encode(figure_bytes).decode("utf-8")
+        figure_node_metadata: Dict[str, Any] = {
+            **(raw_figure_node.node.metadata or {}),
+            "file_id": raw_figure_node.node.file_id,
+            "page_index": raw_figure_node.node.page_index,
+            "figure_name": raw_figure_node.node.figure_name,
+        }
+        figure_node_with_score = NodeWithScore(
+            node=ImageNode(image=figure_base64, metadata=figure_node_metadata),
+            score=raw_figure_node.score,
+        )
+        figure_nodes.append(figure_node_with_score)
+
+    return figure_nodes

py/llama_cloud_services/index/composite_retriever.py

@@ -0,0 +1,291 @@
+from typing import Any, List, Optional
+
+import httpx
+from llama_cloud import (
+    CompositeRetrievalMode,
+    CompositeRetrievedTextNodeWithScore,
+    RetrieverCreate,
+    Retriever,
+    RetrieverPipeline,
+    PresetRetrievalParams,
+    ReRankConfig,
+)
+from llama_cloud.resources.pipelines.client import OMIT
+
+from llama_index.core.base.base_retriever import BaseRetriever
+from llama_index.core.constants import DEFAULT_PROJECT_NAME
+from llama_index.core.ingestion.api_utils import get_aclient, get_client
+from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
+from .base import LlamaCloudIndex
+from .api_utils import (
+    resolve_project,
+    resolve_retriever,
+    page_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,
+        app_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[ReRankConfig] = None,
+        persisted: Optional[bool] = True,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the Composite Retriever."""
+        # initialize clients
+        self._client = get_client(api_key, base_url, app_url, timeout, httpx_client)
+        self._aclient = get_aclient(
+            api_key, base_url, app_url, timeout, async_httpx_client
+        )
+
+        self.project = resolve_project(
+            self._client, project_name, project_id, organization_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  # type: ignore [arg-type]
+        )
+
+        if self.retriever is None and persisted:
+            if create_if_not_exists:
+                self.retriever = self._client.retrievers.upsert_retriever(
+                    project_id=self.project.id,
+                    request=RetrieverCreate(name=self.name, pipelines=[]),
+                )
+            else:
+                raise ValueError(
+                    f"Retriever with name '{self.name}' does not exist in project."
+                )
+
+        # 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__(
+            callback_manager=kwargs.get("callback_manager"),
+            verbose=kwargs.get("verbose", False),
+        )
+
+    @property
+    def retriever_pipelines(self) -> List[RetrieverPipeline]:
+        return self.retriever.pipelines or []  # type: ignore [union-attr]
+
+    def update_retriever_pipelines(
+        self, pipelines: List[RetrieverPipeline]
+    ) -> Retriever:
+        if self._persisted:
+            self.retriever = self._client.retrievers.update_retriever(
+                self.retriever.id, pipelines=pipelines  # type: ignore [union-attr]
+            )
+        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
+
+    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(
+                self.retriever.id, pipelines=pipelines  # type: ignore [union-attr]
+            )
+        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: CompositeRetrievedTextNodeWithScore
+    ) -> 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,
+        )
+
+    def _retrieve(
+        self,
+        query_bundle: QueryBundle,
+        mode: Optional[CompositeRetrievalMode] = None,
+        rerank_top_n: Optional[int] = None,
+        rerank_config: Optional[ReRankConfig] = None,
+    ) -> 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
+        rerank_config = (
+            rerank_config if rerank_config is not None else self._rerank_config
+        )
+
+        if self._persisted:
+            result = self._client.retrievers.retrieve(
+                self.retriever.id,  # type: ignore [union-attr]
+                mode=mode,
+                rerank_top_n=rerank_top_n,
+                rerank_config=rerank_config,
+                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,
+                rerank_config=rerank_config,
+                query=query_bundle.query_str,
+                pipelines=self.retriever.pipelines,  # type: ignore [union-attr]
+            )
+        node_w_scores = [
+            self._result_nodes_to_node_with_score(node) for node in result.nodes  # type: ignore [union-attr]
+        ]
+        image_nodes_w_scores = page_screenshot_nodes_to_node_with_score(
+            self._client, result.image_nodes, self.retriever.project_id  # type: ignore [union-attr]
+        )
+        return sorted(
+            node_w_scores + image_nodes_w_scores, key=lambda x: x.score, reverse=True
+        )
+
+    async def _aretrieve(
+        self,
+        query_bundle: QueryBundle,
+        mode: Optional[CompositeRetrievalMode] = None,
+        rerank_top_n: Optional[int] = None,
+        rerank_config: Optional[ReRankConfig] = None,
+    ) -> 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
+        rerank_config = (
+            rerank_config if rerank_config is not None else self._rerank_config
+        )
+
+        if self._persisted:
+            result = await self._aclient.retrievers.retrieve(
+                self.retriever.id,  # type: ignore [union-attr]
+                mode=mode,
+                rerank_config=rerank_config,
+                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,
+                rerank_config=rerank_config,
+                query=query_bundle.query_str,
+                pipelines=self.retriever.pipelines,  # type: ignore [union-attr]
+            )
+        node_w_scores = [
+            self._result_nodes_to_node_with_score(node) for node in result.nodes  # type: ignore [union-attr]
+        ]
+        image_nodes_w_scores = page_screenshot_nodes_to_node_with_score(
+            self._aclient, result.image_nodes, self.retriever.project_id  # type: ignore [union-attr]
+        )
+        return sorted(
+            node_w_scores + image_nodes_w_scores, key=lambda x: x.score, reverse=True
+        )

py/llama_cloud_services/index/retriever.py

@@ -0,0 +1,217 @@
+from typing import Any, List, Optional
+
+import httpx
+from llama_cloud import (
+    TextNodeWithScore,
+)
+from llama_cloud.resources.pipelines.client import OMIT
+
+from llama_index.core.base.base_retriever import BaseRetriever
+from llama_index.core.bridge.pydantic import BaseModel
+from llama_index.core.constants import DEFAULT_PROJECT_NAME
+from llama_index.core.ingestion.api_utils import get_aclient, get_client
+from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
+from llama_index.core.vector_stores.types import MetadataFilters
+from .api_utils import (
+    resolve_project_and_pipeline,
+    page_screenshot_nodes_to_node_with_score,
+    page_figure_nodes_to_node_with_score,
+    apage_screenshot_nodes_to_node_with_score,
+    apage_figure_nodes_to_node_with_score,
+)
+import logging
+
+
+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,
+        app_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 = get_client(api_key, base_url, app_url, timeout, httpx_client)
+        self._aclient = get_aclient(
+            api_key, base_url, app_url, timeout, 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 OMIT
+        )
+        self._sparse_similarity_top_k = (
+            sparse_similarity_top_k if sparse_similarity_top_k is not None else OMIT
+        )
+        self._enable_reranking = (
+            enable_reranking if enable_reranking is not None else OMIT
+        )
+        self._rerank_top_n = rerank_top_n if rerank_top_n is not None else OMIT
+        self._alpha = alpha if alpha is not None else OMIT
+        self._filters = filters if filters is not None else OMIT
+        self._retrieval_mode = retrieval_mode if retrieval_mode is not None else OMIT
+        self._files_top_k = files_top_k if files_top_k is not None else OMIT
+        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 OMIT
+        )
+        self._retrieve_page_figure_nodes = (
+            retrieve_page_figure_nodes
+            if retrieve_page_figure_nodes is not None
+            else OMIT
+        )
+        self._search_filters_inference_schema = search_filters_inference_schema
+
+        super().__init__(
+            callback_manager=kwargs.get("callback_manager"),
+            verbose=kwargs.get("verbose", False),
+        )
+
+    def _result_nodes_to_node_with_score(
+        self, result_nodes: List[TextNodeWithScore]
+    ) -> List[NodeWithScore]:
+        nodes = []
+        for res in result_nodes:
+            text_node = TextNode.parse_obj(res.node.dict())
+            nodes.append(NodeWithScore(node=text_node, score=res.score))
+
+        return nodes
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Retrieve from the platform."""
+        search_filters_inference_schema = OMIT
+        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.run_search(
+            query=query_bundle.query_str,
+            pipeline_id=self.pipeline.id,
+            dense_similarity_top_k=self._dense_similarity_top_k,
+            sparse_similarity_top_k=self._sparse_similarity_top_k,
+            enable_reranking=self._enable_reranking,
+            rerank_top_n=self._rerank_top_n,
+            alpha=self._alpha,
+            search_filters=self._filters,
+            files_top_k=self._files_top_k,
+            retrieval_mode=self._retrieval_mode,
+            retrieve_page_screenshot_nodes=self._retrieve_page_screenshot_nodes,
+            retrieve_page_figure_nodes=self._retrieve_page_figure_nodes,
+            search_filters_inference_schema=search_filters_inference_schema,
+        )
+
+        result_nodes = self._result_nodes_to_node_with_score(results.retrieval_nodes)
+        if self._retrieve_page_screenshot_nodes:
+            result_nodes.extend(
+                page_screenshot_nodes_to_node_with_score(
+                    self._client, results.image_nodes, self.project.id
+                )
+            )
+        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
+                )
+            )
+
+        return result_nodes
+
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Asynchronously retrieve from the platform."""
+        search_filters_inference_schema = OMIT
+        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.run_search(
+            query=query_bundle.query_str,
+            pipeline_id=self.pipeline.id,
+            dense_similarity_top_k=self._dense_similarity_top_k,
+            sparse_similarity_top_k=self._sparse_similarity_top_k,
+            enable_reranking=self._enable_reranking,
+            rerank_top_n=self._rerank_top_n,
+            alpha=self._alpha,
+            search_filters=self._filters,
+            files_top_k=self._files_top_k,
+            retrieval_mode=self._retrieval_mode,
+            retrieve_page_screenshot_nodes=self._retrieve_page_screenshot_nodes,
+            retrieve_page_figure_nodes=self._retrieve_page_figure_nodes,
+            search_filters_inference_schema=search_filters_inference_schema,
+        )
+
+        result_nodes = self._result_nodes_to_node_with_score(results.retrieval_nodes)
+        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
+                )
+            )
+        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
+                )
+            )
+
+        return result_nodes

py/llama_parse/pyproject.toml

@@ -17,7 +17,7 @@ authors = [{name = "Logan Markewich", email = "logan@llamaindex.ai"}]
 requires-python = ">=3.9,<4.0"
 readme = "README.md"
 license = "MIT"
-dependencies = ["llama-cloud-services>=0.6.52"]
+dependencies = ["llama-cloud-services>=0.6.53"]
 
 [project.scripts]
 llama-parse = "llama_parse.cli.main:parse"

py/tests/index/test_index.py

@@ -0,0 +1,538 @@
+import os
+import pytest
+import tempfile
+from typing import Generator, Tuple
+from uuid import uuid4
+
+from llama_cloud import (
+    AutoTransformConfig,
+    PipelineCreate,
+    PipelineFileCreate,
+    ProjectCreate,
+    CompositeRetrievalMode,
+    LlamaParseParameters,
+    ReRankConfig,
+)
+from llama_cloud.client import LlamaCloud
+from llama_index.core.bridge.pydantic import BaseModel
+from llama_index.core.constants import DEFAULT_BASE_URL
+from llama_index.core.indices.managed.base import BaseManagedIndex
+from llama_index.core.schema import Document, ImageNode
+from llama_cloud_services.index import (
+    LlamaCloudIndex,
+    LlamaCloudCompositeRetriever,
+)
+
+base_url = os.environ.get("LLAMA_CLOUD_BASE_URL", DEFAULT_BASE_URL)
+api_key = os.environ.get("LLAMA_CLOUD_API_KEY", None)
+organization_id = os.environ.get("LLAMA_CLOUD_ORGANIZATION_ID", None)
+project_name = os.environ.get("LLAMA_CLOUD_PROJECT_NAME", "framework_integration_test")
+
+
+@pytest.fixture()
+def remote_file() -> Tuple[str, str]:
+    test_file_url = "https://www.google.com/robots.txt"
+    test_file_name = "google_robots.txt"
+    return test_file_url, test_file_name
+
+
+@pytest.fixture()
+def index_name() -> Generator[str, None, None]:
+    name = f"test_index_{uuid4()}"
+    try:
+        yield name
+    finally:
+        client = LlamaCloud(token=api_key, base_url=base_url)
+        pipeline = client.pipelines.search_pipelines(project_name=name)
+        if pipeline:
+            client.pipelines.delete(pipeline_id=pipeline[0].id)
+
+
+@pytest.fixture()
+def local_file() -> str:
+    return "tests/test_files/index/Simple PDF Slides.pdf"
+
+
+@pytest.fixture()
+def local_figures_file() -> str:
+    return "tests/test_files/index/image_figure_slides.pdf"
+
+
+def _setup_index_with_file(
+    client: LlamaCloud, index_name: str, remote_file: Tuple[str, str]
+) -> LlamaCloudIndex:
+    # create project if it doesn't exist
+    project_create = ProjectCreate(name=project_name)
+    project = client.projects.upsert_project(
+        organization_id=organization_id, request=project_create
+    )
+
+    # create pipeline
+    pipeline_create = PipelineCreate(
+        name=index_name,
+        transform_config=AutoTransformConfig(),
+    )
+    pipeline = client.pipelines.upsert_pipeline(
+        project_id=project.id, request=pipeline_create
+    )
+
+    # upload file to pipeline
+    test_file_url, test_file_name = remote_file
+    file = client.files.upload_file_from_url(
+        project_id=project.id, url=test_file_url, name=test_file_name
+    )
+
+    # add file to pipeline
+    pipeline_file_create = PipelineFileCreate(file_id=file.id)
+    client.pipelines.add_files_to_pipeline_api(
+        pipeline_id=pipeline.id, request=[pipeline_file_create]
+    )
+
+    return pipeline
+
+
+def test_class():
+    names_of_base_classes = [b.__name__ for b in LlamaCloudIndex.__mro__]
+    assert BaseManagedIndex.__name__ in names_of_base_classes
+
+
+def test_conflicting_index_identifiers():
+    with pytest.raises(ValueError):
+        LlamaCloudIndex(name="test", pipeline_id="test", index_id="test")
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+def test_resolve_index_with_id(remote_file: Tuple[str, str], index_name: str):
+    """Test that we can instantiate an index with a given id."""
+    client = LlamaCloud(token=api_key, base_url=base_url)
+    pipeline = _setup_index_with_file(client, index_name, remote_file)
+
+    index = LlamaCloudIndex(
+        pipeline_id=pipeline.id,
+        api_key=api_key,
+        base_url=base_url,
+    )
+    assert index is not None
+
+    index.wait_for_completion()
+    retriever = index.as_retriever()
+
+    nodes = retriever.retrieve("Hello world.")
+    assert len(nodes) > 0
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+def test_resolve_index_with_name(remote_file: Tuple[str, str], index_name: str):
+    """Test that we can instantiate an index with a given name."""
+    client = LlamaCloud(token=api_key, base_url=base_url)
+    pipeline = _setup_index_with_file(client, index_name, remote_file)
+
+    index = LlamaCloudIndex(
+        name=pipeline.name,
+        project_name=project_name,
+        organization_id=organization_id,
+        api_key=api_key,
+        base_url=base_url,
+    )
+    assert index is not None
+
+    index.wait_for_completion()
+    retriever = index.as_retriever()
+
+    nodes = retriever.retrieve("Hello world.")
+    assert len(nodes) > 0
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+def test_upload_file(index_name: str):
+    index = LlamaCloudIndex.create_index(
+        name=index_name,
+        project_name=project_name,
+        organization_id=organization_id,
+        api_key=api_key,
+        base_url=base_url,
+    )
+
+    # Create a temporary file to upload
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as temp_file:
+        temp_file.write(b"Sample content for testing upload.")
+        temp_file_path = temp_file.name
+
+    try:
+        # Upload the file
+        file_id = index.upload_file(temp_file_path, verbose=True)
+        assert file_id is not None
+
+        # Verify the file is part of the index
+        docs = index.ref_doc_info
+        temp_file_name = os.path.basename(temp_file_path)
+        assert any(
+            temp_file_name == doc.metadata.get("file_name") for doc in docs.values()
+        )
+
+    finally:
+        # Clean up the temporary file
+        os.remove(temp_file_path)
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+def test_upload_file_from_url(remote_file: Tuple[str, str], index_name: str):
+    index = LlamaCloudIndex.create_index(
+        name=index_name,
+        project_name=project_name,
+        organization_id=organization_id,
+        api_key=api_key,
+        base_url=base_url,
+    )
+
+    # Define a URL to a file for testing
+    test_file_url, test_file_name = remote_file
+
+    # Upload the file from the URL
+    file_id = index.upload_file_from_url(
+        file_name=test_file_name, url=test_file_url, verbose=True
+    )
+    assert file_id is not None
+
+    # Verify the file is part of the index
+    docs = index.ref_doc_info
+    assert any(test_file_name == doc.metadata.get("file_name") for doc in docs.values())
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+def test_index_from_documents(index_name: str):
+    documents = [
+        Document(text="Hello world.", doc_id="1", metadata={"source": "test"}),
+    ]
+    index = LlamaCloudIndex.from_documents(
+        documents=documents,
+        name=index_name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+        organization_id=organization_id,
+        verbose=True,
+    )
+    docs = index.ref_doc_info
+    assert len(docs) == 1
+    assert docs["1"].metadata["source"] == "test"
+    nodes = index.as_retriever().retrieve("Hello world.")
+    assert len(nodes) > 0
+    assert all(n.node.ref_doc_id == "1" for n in nodes)
+    assert all(n.node.metadata["source"] == "test" for n in nodes)
+
+    index.insert(
+        Document(text="Hello world.", doc_id="2", metadata={"source": "inserted"}),
+        verbose=True,
+    )
+    docs = index.ref_doc_info
+    assert len(docs) == 2
+    assert docs["2"].metadata["source"] == "inserted"
+    nodes = index.as_retriever().retrieve("Hello world.")
+    assert len(nodes) > 0
+    assert all(n.node.ref_doc_id in ["1", "2"] for n in nodes)
+    assert any(n.node.ref_doc_id == "1" for n in nodes)
+    assert any(n.node.ref_doc_id == "2" for n in nodes)
+
+    index.update_ref_doc(
+        Document(text="Hello world.", doc_id="2", metadata={"source": "updated"}),
+        verbose=True,
+    )
+    docs = index.ref_doc_info
+    assert len(docs) == 2
+    assert docs["2"].metadata["source"] == "updated"
+
+    index.refresh_ref_docs(
+        [
+            Document(text="Hello world.", doc_id="1", metadata={"source": "refreshed"}),
+            Document(text="Hello world.", doc_id="3", metadata={"source": "refreshed"}),
+        ]
+    )
+    docs = index.ref_doc_info
+    assert len(docs) == 3
+    assert docs["3"].metadata["source"] == "refreshed"
+    assert docs["1"].metadata["source"] == "refreshed"
+
+    index.delete_ref_doc("3", verbose=True)
+    docs = index.ref_doc_info
+    assert len(docs) == 2
+    assert "3" not in docs
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_page_screenshot_retrieval(index_name: str, local_file: str):
+    index = await LlamaCloudIndex.acreate_index(
+        name=index_name,
+        project_name=project_name,
+        organization_id=organization_id,
+        api_key=api_key,
+        base_url=base_url,
+        llama_parse_parameters=LlamaParseParameters(
+            take_screenshot=True,
+        ),
+    )
+
+    file_id = index.upload_file(local_file, wait_for_ingestion=True)
+
+    retriever = index.as_retriever(retrieve_page_screenshot_nodes=True)
+    nodes = retriever.retrieve("1")
+    assert len(nodes) > 0
+
+    image_nodes = [n.node for n in nodes if isinstance(n.node, ImageNode)]
+    assert len(image_nodes) > 0
+    assert all(n.metadata["file_id"] == file_id for n in image_nodes)
+    assert all(n.metadata["page_index"] >= 0 for n in image_nodes)
+    # ensure metadata is added from the image node
+    # local_figures_file has the full absolute path, so just check the file name is in that absolute path
+    assert all(local_file.endswith(n.metadata["file_name"]) for n in image_nodes)
+
+    nodes = await retriever.aretrieve("1")
+    assert len(nodes) > 0
+
+    image_nodes = [n.node for n in nodes if isinstance(n.node, ImageNode)]
+    assert len(image_nodes) > 0
+    assert all(n.metadata["file_id"] == file_id for n in image_nodes)
+    assert all(n.metadata["page_index"] >= 0 for n in image_nodes)
+    # ensure metadata is added from the image node
+    # local_figures_file has the full absolute path, so just check the file name is in that absolute path
+    assert all(local_file.endswith(n.metadata["file_name"]) for n in image_nodes)
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_page_figure_retrieval(index_name: str, local_figures_file: str):
+    index = await LlamaCloudIndex.acreate_index(
+        name=index_name,
+        project_name=project_name,
+        organization_id=organization_id,
+        api_key=api_key,
+        base_url=base_url,
+        llama_parse_parameters=LlamaParseParameters(
+            take_screenshot=True,
+            extract_layout=True,
+        ),
+    )
+
+    file_id = index.upload_file(local_figures_file, wait_for_ingestion=True)
+
+    retriever = index.as_retriever(retrieve_page_figure_nodes=True)
+    nodes = retriever.retrieve("1")
+    assert len(nodes) > 0
+
+    image_nodes = [n.node for n in nodes if isinstance(n.node, ImageNode)]
+    assert len(image_nodes) > 0
+    assert all(n.metadata["file_id"] == file_id for n in image_nodes)
+    assert all(n.metadata["page_index"] >= 0 for n in image_nodes)
+    # ensure metadata is added from the image node
+    # local_figures_file has the full absolute path, so just check the file name is in that absolute path
+    assert all(
+        local_figures_file.endswith(n.metadata["file_name"]) for n in image_nodes
+    )
+
+    nodes = await retriever.aretrieve("1")
+    assert len(nodes) > 0
+
+    image_nodes = [n.node for n in nodes if isinstance(n.node, ImageNode)]
+    assert len(image_nodes) > 0
+    assert all(n.metadata["file_id"] == file_id for n in image_nodes)
+    assert all(n.metadata["page_index"] >= 0 for n in image_nodes)
+    # ensure metadata is added from the image node
+    # local_figures_file has the full absolute path, so just check the file name is in that absolute path
+    assert all(
+        local_figures_file.endswith(n.metadata["file_name"]) for n in image_nodes
+    )
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_composite_retriever(index_name: str):
+    """Test the LlamaCloudCompositeRetriever with multiple indices."""
+    # Create first index with documents
+    documents1 = [
+        Document(
+            text="Hello world from index 1.", doc_id="1", metadata={"source": "index1"}
+        ),
+    ]
+    index1 = LlamaCloudIndex.from_documents(
+        documents=documents1,
+        name=index_name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+        organization_id=organization_id,
+        verbose=True,
+    )
+
+    # Create second index with documents
+    documents2 = [
+        Document(
+            text="Hello world from index 2.", doc_id="2", metadata={"source": "index2"}
+        ),
+    ]
+    index2 = LlamaCloudIndex.from_documents(
+        documents=documents2,
+        name=f"test pipeline 2 {uuid4()}",
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+        organization_id=organization_id,
+        verbose=True,
+    )
+
+    # Create a composite retriever
+    retriever = LlamaCloudCompositeRetriever(
+        name="composite_retriever_test",
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+        create_if_not_exists=True,
+        mode=CompositeRetrievalMode.FULL,
+        rerank_top_n=5,
+        rerank_config=ReRankConfig(
+            top_n=5,
+        ),
+    )
+
+    # Attach indices to the composite retriever
+    retriever.add_index(index1, description="Information from index 1.")
+    retriever.add_index(index2, description="Information from index 2.")
+
+    # Retrieve nodes using the composite retriever
+    nodes = retriever.retrieve("Hello world.")
+
+    # Assertions to verify the retrieval
+    assert len(nodes) >= 2
+
+    # Retrieve nodes using the composite retriever
+    nodes = await retriever.aretrieve("Hello world.")
+
+    # Assertions to verify the retrieval
+    assert len(nodes) >= 2
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_async_index_from_documents(index_name: str):
+    documents = [
+        Document(text="Hello world.", doc_id="1", metadata={"source": "test"}),
+    ]
+    index = await LlamaCloudIndex.acreate_index(
+        name=index_name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+        organization_id=organization_id,
+        verbose=True,
+    )
+    await index.ainsert(documents[0])
+    await index.await_for_completion()
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_async_upload_file_from_url(
+    remote_file: Tuple[str, str], index_name: str
+):
+    index = await LlamaCloudIndex.acreate_index(
+        name=index_name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+    )
+
+    test_file_url, test_file_name = remote_file
+    file_id = await index.aupload_file_from_url(
+        file_name=test_file_name, url=test_file_url, verbose=True
+    )
+    assert file_id is not None
+
+    await index.await_for_completion()
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_async_index_from_file(index_name: str, local_file: str):
+    index = await LlamaCloudIndex.acreate_index(
+        name=index_name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+    )
+
+    file_id = await index.aupload_file(file_path=local_file, verbose=True)
+    assert file_id is not None
+
+    await index.await_for_completion()
+
+
+class DummySchema(BaseModel):
+    source: str
+
+
+@pytest.mark.skipif(
+    not base_url or not api_key, reason="No platform base url or api key set"
+)
+@pytest.mark.asyncio
+async def test_search_filters_inference_schema(index_name: str):
+    """Test the use of search_filters_inference_schema in retrieval."""
+    # Define a dummy schema
+    schema = DummySchema(source="test")
+
+    # Create documents
+    documents = [
+        Document(text="Hello world.", doc_id="1", metadata={"source": "test"}),
+    ]
+
+    # Create an index with documents
+    index = LlamaCloudIndex.from_documents(
+        documents=documents,
+        name=index_name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+        organization_id=organization_id,
+        verbose=True,
+    )
+
+    # Use the retriever with the schema
+    retriever = index.as_retriever(search_filters_inference_schema=schema)
+    nodes = retriever.retrieve(
+        'Search for documents where the metadata has source="test"'
+    )
+
+    # Verify that nodes are retrieved
+    assert len(nodes) > 0
+    assert all(n.node.ref_doc_id == "1" for n in nodes)
+    assert all(n.node.metadata["source"] == "test" for n in nodes)
+
+    nodes = await retriever.aretrieve(
+        'Search for documents where the metadata has source="test"'
+    )
+
+    # Verify that nodes are retrieved
+    assert len(nodes) > 0
+    assert all(n.node.ref_doc_id == "1" for n in nodes)
+    assert all(n.node.metadata["source"] == "test" for n in nodes)