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

llama_cloud/lib/index/api_utils.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import uuid
+import base64
+from typing import Any, Dict, List, Tuple, Optional
+
+from llama_index.core.schema import ImageNode, NodeWithScore
+from llama_index.core.async_utils import run_jobs
+
+from llama_cloud import LlamaCloud, AsyncLlamaCloud
+from llama_cloud.types import Retriever, PageFigureNodeWithScore, AutoTransformConfigParam, PageScreenshotNodeWithScore
+from llama_cloud.types.project import Project
+from llama_cloud.types.pipeline import Pipeline
+
+
+def default_transform_config() -> AutoTransformConfigParam:
+    return AutoTransformConfigParam()
+
+
+def resolve_retriever(
+    client: LlamaCloud,
+    project: Project,
+    retriever_name: Optional[str] = None,
+    retriever_id: Optional[str] = None,
+    persisted: Optional[bool] = True,
+) -> Optional[Retriever]:
+    if not persisted:
+        return Retriever(
+            id=str(uuid.uuid4()),
+            project_id=project.id,
+            name=retriever_name or f"retriever-{uuid.uuid4()}",
+            pipelines=[],
+        )
+    if retriever_id:
+        return client.retrievers.get(retriever_id=retriever_id, project_id=project.id)
+    elif retriever_name:
+        retrievers = client.retrievers.list(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:
+    project: Optional[Project] = None
+    if project_id is not None:
+        project = client.projects.get(project_id=project_id)
+    elif project_name is not None:
+        projects = client.projects.list(organization_id=organization_id)
+        project = next((p for p in projects if p.name == project_name), None)
+        if project is None:
+            raise ValueError(f"Project with name '{project_name}' not found.")
+    else:
+        raise ValueError("Either project_id or project_name must be provided.")
+
+    return project
+
+
+def resolve_project_and_pipeline(
+    client: LlamaCloud,
+    name: Optional[str],
+    pipeline_id: Optional[str],
+    project_name: Optional[str],
+    project_id: Optional[str],
+    organization_id: Optional[str],
+) -> Tuple[Project, Pipeline]:
+    project = resolve_project(
+        client=client,
+        project_name=project_name,
+        project_id=project_id,
+        organization_id=organization_id,
+    )
+
+    pipeline: Optional[Pipeline] = None
+    if pipeline_id is not None:
+        pipeline = client.pipelines.get(pipeline_id=pipeline_id)
+    elif name is not None:
+        pipelines = client.pipelines.list(organization_id=organization_id, project_id=project.id, pipeline_name=name)
+        pipeline = next((p for p in pipelines if p.name == name), None)
+        if pipeline is None:
+            raise ValueError(f"Pipeline with name '{name}' not found in project '{project.name}'.")
+    else:
+        raise ValueError("Either pipeline_id or name must be provided.")
+
+    return project, pipeline
+
+
+def page_screenshot_nodes_to_node_with_score(
+    client: LlamaCloud,
+    raw_image_nodes: Optional[List[PageScreenshotNodeWithScore]],
+    project_id: str,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> List[NodeWithScore]:
+    if not raw_image_nodes:
+        return []
+
+    image_nodes: List[NodeWithScore] = []
+    for raw_image_node in raw_image_nodes:
+        image_bytes_str = client.pipelines.images.get_page_screenshot(
+            page_index=raw_image_node.node.page_index,
+            id=raw_image_node.node.file_id,
+            project_id=project_id,
+        )
+        image_base64 = base64.b64encode(str(image_bytes_str).encode("utf-8")).decode("utf-8")
+        image_node_metadata: Dict[str, Any] = {
+            **(raw_image_node.node.metadata or {}),
+            **(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,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> 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,
+        metadata=metadata,
+    )
+
+
+def page_figure_nodes_to_node_with_score(
+    client: LlamaCloud,
+    raw_figure_nodes: Optional[List[PageFigureNodeWithScore]],
+    project_id: str,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> List[NodeWithScore]:
+    if not raw_figure_nodes:
+        return []
+
+    figure_nodes: List[NodeWithScore] = []
+    for raw_figure_node in raw_figure_nodes:
+        figure_bytes_str = client.pipelines.images.get_page_figure(
+            page_index=raw_figure_node.node.page_index,
+            id=raw_figure_node.node.file_id,
+            figure_name=raw_figure_node.node.figure_name,
+            project_id=project_id,
+        )
+        figure_base64 = base64.b64encode(str(figure_bytes_str).encode("utf-8")).decode("utf-8")
+        figure_node_metadata: Dict[str, Any] = {
+            **(raw_figure_node.node.metadata or {}),
+            **(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,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> List[NodeWithScore]:
+    if not raw_image_nodes:
+        return []
+
+    async def _get_page_screenshot(
+        client: AsyncLlamaCloud,
+        file_id: str,
+        page_index: int,
+        project_id: str,
+    ) -> str:
+        figure_bytes_str = await client.pipelines.images.get_page_screenshot(
+            page_index=page_index,
+            id=file_id,
+            project_id=project_id,
+        )
+        return base64.b64encode(str(figure_bytes_str).encode("utf-8")).decode("utf-8")
+
+    image_nodes: List[NodeWithScore] = []
+    tasks = [
+        _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,
+        )
+        for raw_image_node in raw_image_nodes
+    ]
+
+    image_bytes_list = await run_jobs(tasks)
+    for image_base64, raw_image_node in zip(image_bytes_list, raw_image_nodes):
+        image_node_metadata: Dict[str, Any] = {
+            **(raw_image_node.node.metadata or {}),
+            **(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,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> 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,
+        metadata=metadata,
+    )
+
+
+async def apage_figure_nodes_to_node_with_score(
+    client: AsyncLlamaCloud,
+    raw_figure_nodes: Optional[List[PageFigureNodeWithScore]],
+    project_id: str,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> List[NodeWithScore]:
+    if not raw_figure_nodes:
+        return []
+
+    async def _get_page_figure(
+        client: AsyncLlamaCloud,
+        file_id: str,
+        figure_name: str,
+        page_index: int,
+        project_id: str,
+    ) -> str:
+        figure_bytes_str = await client.pipelines.images.get_page_figure(
+            page_index=page_index,
+            figure_name=figure_name,
+            id=file_id,
+            project_id=project_id,
+        )
+        return base64.b64encode(str(figure_bytes_str).encode("utf-8")).decode("utf-8")
+
+    figure_nodes: List[NodeWithScore] = []
+    tasks = [
+        _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,
+        )
+        for raw_figure_node in raw_figure_nodes
+    ]
+
+    figure_bytes_list = await run_jobs(tasks)
+    for figure_base64, raw_figure_node in zip(figure_bytes_list, raw_figure_nodes):
+        figure_node_metadata: Dict[str, Any] = {
+            **(raw_figure_node.node.metadata or {}),
+            **(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