Add file related fields for file tracking. Simplify API

llama_cloud_services/beta/agent_data/__init__.py

@@ -5,6 +5,7 @@ from .schema import (
     StatusType,
     ExtractedT,
     AgentDataT,
+    ComparisonOperator,
 )
 from .client import AsyncAgentDataClient
 
@@ -16,4 +17,5 @@ __all__ = [
     "StatusType",
     "ExtractedT",
     "AgentDataT",
+    "ComparisonOperator",
 ]

llama_cloud_services/beta/agent_data/client.py

@@ -1,7 +1,6 @@
 import os
-from typing import Dict, Generic, List, Optional, Type
+from typing import Any, Dict, Generic, List, Optional, Type
 
-from llama_cloud import FilterOperation
 from llama_cloud.client import AsyncLlamaCloud
 from tenacity import (
     WrappedFn,
@@ -14,6 +13,7 @@ import httpx
 
 from .schema import (
     AgentDataT,
+    ComparisonOperator,
     TypedAgentData,
     TypedAgentDataItems,
     TypedAggregateGroup,
@@ -85,7 +85,7 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         agent_client = AsyncAgentDataClient(
             client=llama_client,
             type=ExtractedPerson,
-            collection_name="extracted_people",
+            collection="extracted_people",
             agent_url_id="person-extraction-agent"
         )
 
@@ -94,8 +94,8 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         result = await agent_client.create_agent_data(person)
 
         # Search data
-        results = await agent_client.search_agent_data(
-            filter={"age": FilterOperation(gt=25)},
+        results = await agent_client.search(
+            filter={"age": {"gt": 25}},
             order_by="data.name",
             page_size=20
         )
@@ -107,19 +107,20 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
 
     def __init__(
         self,
-        client: AsyncLlamaCloud,
         type: Type[AgentDataT],
-        collection_name: str = "default",
+        collection: str = "default",
         agent_url_id: Optional[str] = None,
+        client: Optional[AsyncLlamaCloud] = None,
+        token: Optional[str] = None,
+        base_url: Optional[str] = None,
     ):
         """
         Initialize the AsyncAgentDataClient.
 
         Args:
-            client: AsyncLlamaCloud client instance for API communication
             type: Pydantic BaseModel class that defines the data structure.
                 All agent data will be validated against this type.
-            collection_name: Named collection within the agent for organizing data.
+            collection: Named collection within the agent for organizing data.
                 Defaults to "default". Collections allow logical separation of
                 different data types or workflows within the same agent.
             agent_url_id: Unique identifier for the agent. This normally appears in the
@@ -127,6 +128,11 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
                 will attempt to use the LLAMA_DEPLOY_DEPLOYMENT_NAME environment
                 variable. Data can only be added to an already existing agent in the
                 platform.
+            client: AsyncLlamaCloud client instance for API communication. If not provided, will
+                construct one from the provided api token and base url
+            token: Llama Cloud API token. Reads from LLAMA_CLOUD_API_KEY if not provided
+            base_url: Llama Cloud API token. Reads from LLAMA_CLOUD_BASE_URL if not provided, and
+                defaults to https://api.cloud.llamaindex.ai
 
         Raises:
             ValueError: If agent_url_id is not provided and the
@@ -143,28 +149,33 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
                 "Agent ID is required, or set the LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable"
             )
 
-        self.collection_name = collection_name
+        self.collection = collection
+        if not client:
+            client = AsyncLlamaCloud(
+                token=token or os.getenv("LLAMA_CLOUD_API_KEY"),
+                base_url=base_url or os.getenv("LLAMA_CLOUD_BASE_URL"),
+            )
         self.client = client
         self.type = type
 
     @agent_data_retry
-    async def get_agent_data(self, item_id: str) -> TypedAgentData[AgentDataT]:
+    async def get_item(self, item_id: str) -> TypedAgentData[AgentDataT]:
         raw_data = await self.client.beta.get_agent_data(
             item_id=item_id,
         )
         return TypedAgentData.from_raw(raw_data, validator=self.type)
 
     @agent_data_retry
-    async def create_agent_data(self, data: AgentDataT) -> TypedAgentData[AgentDataT]:
+    async def create_item(self, data: AgentDataT) -> TypedAgentData[AgentDataT]:
         raw_data = await self.client.beta.create_agent_data(
             agent_slug=self.agent_url_id,
-            collection=self.collection_name,
+            collection=self.collection,
             data=data.model_dump(),
         )
         return TypedAgentData.from_raw(raw_data, validator=self.type)
 
     @agent_data_retry
-    async def update_agent_data(
+    async def update_item(
         self, item_id: str, data: AgentDataT
     ) -> TypedAgentData[AgentDataT]:
         raw_data = await self.client.beta.update_agent_data(
@@ -174,13 +185,13 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         return TypedAgentData.from_raw(raw_data, validator=self.type)
 
     @agent_data_retry
-    async def delete_agent_data(self, item_id: str) -> None:
+    async def delete_item(self, item_id: str) -> None:
         await self.client.beta.delete_agent_data(item_id=item_id)
 
     @agent_data_retry
-    async def search_agent_data(
+    async def search(
         self,
-        filter: Optional[Dict[str, Optional[FilterOperation]]] = None,
+        filter: Optional[Dict[str, Dict[ComparisonOperator, Any]]] = None,
         order_by: Optional[str] = None,
         offset: Optional[int] = None,
         page_size: Optional[int] = None,
@@ -191,11 +202,11 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         Args:
             filter: Filter conditions to apply to the search. Dict mapping field names to FilterOperation objects. Filters only by data fields
                 Examples:
-                - {"age": FilterOperation(gt=18)} - age greater than 18
-                - {"status": FilterOperation(eq="active")} - status equals "active"
-                - {"tags": FilterOperation(includes=["python", "ml"])} - tags include "python" or "ml"
-                - {"created_at": FilterOperation(gte="2024-01-01")} - created after date
-                - {"score": FilterOperation(lt=100, gte=50)} - score between 50 and 100
+                - {"age": {"gt": 18}} - age greater than 18
+                - {"status": {"eq": "active"}} - status equals "active"
+                - {"tags": {"includes": ["python", "ml"]}} - tags include "python" or "ml"
+                - {"created_at": {"gte": "2024-01-01"}} - created after date
+                - {"score": {"lt": 100, "gte": 50}} - score between 50 and 100
             order_by: Comma delimited list of fields to sort results by. Can order by standard agent fields like created_at, or by data fields. Data fields must be prefixed with "data.". If ordering desceding, use a " desc" suffix.
                 Examples:
                 - "data.name desc, created_at" - sort by name in descending order, and then by creation date
@@ -205,7 +216,7 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         """
         raw = await self.client.beta.search_agent_data_api_v_1_beta_agent_data_search_post(
             agent_slug=self.agent_url_id,
-            collection=self.collection_name,
+            collection=self.collection,
             filter=filter,
             order_by=order_by,
             offset=offset,
@@ -221,9 +232,9 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         )
 
     @agent_data_retry
-    async def aggregate_agent_data(
+    async def aggregate(
         self,
-        filter: Optional[Dict[str, Optional[FilterOperation]]] = None,
+        filter: Optional[Dict[str, Dict[ComparisonOperator, Any]]] = None,
         group_by: Optional[List[str]] = None,
         count: Optional[bool] = None,
         first: Optional[bool] = None,
@@ -235,20 +246,20 @@ class AsyncAgentDataClient(Generic[AgentDataT]):
         Aggregate agent data into groups according to the group_by fields.
         Args:
             filter: Filter conditions to apply to the search. Dict mapping field names to FilterOperation objects. Filters only by data fields
-                See search_agent_data for more details on filtering.
+                See search for more details on filtering.
             group_by: List of fields to group by. Groups strictly by equality. Can only group by data fields.
                 Examples:
                 - ["name"] - group by name
                 - ["name", "age"] - group by name and age
             count: Whether to include the count of items in each group.
             first: Whether to include the first item in each group.
-            order_by: Comma delimited list of fields to sort results by. See search_agent_data for more details on ordering.
+            order_by: Comma delimited list of fields to sort results by. See search for more details on ordering.
             offset: Number of groups to skip from the beginning. Defaults to 0.
             page_size: Maximum number of groups to return per page.
         """
         raw = await self.client.beta.aggregate_agent_data_api_v_1_beta_agent_data_aggregate_post(
             agent_slug=self.agent_url_id,
-            collection=self.collection_name,
+            collection=self.collection,
             page_size=page_size,
             filter=filter,
             order_by=order_by,

llama_cloud_services/beta/agent_data/schema.py

@@ -60,7 +60,11 @@ AgentDataT = TypeVar("AgentDataT", bound=BaseModel)
 ExtractedT = TypeVar("ExtractedT", bound=Union[BaseModel, dict])
 
 # Status types for extracted data workflow
-StatusType = Union[Literal["error", "accepted", "rejected", "in_review"], str]
+StatusType = Union[Literal["error", "accepted", "rejected", "pending_review"], str]
+
+ComparisonOperator = Dict[
+    str, Dict[Literal["gt", "gte", "lt", "lte", "eq", "includes"], Any]
+]
 
 
 class TypedAgentData(BaseModel, Generic[AgentDataT]):
@@ -142,7 +146,7 @@ class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
     Example:
         ```python
         # Search with pagination
-        results = await client.search_agent_data(
+        results = await client.search(
             page_size=10,
             include_total=True
         )
@@ -152,7 +156,7 @@ class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
 
         if results.has_more:
             # Load next page
-            next_page = await client.search_agent_data(
+            next_page = await client.search(
                 page_size=10,
                 offset=10
             )
@@ -183,9 +187,12 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         data: The current state of the data (may differ from original after edits)
         status: Current workflow status (in_review, accepted, rejected, error)
         confidence: Confidence scores for individual fields (if available)
+        file_id: The llamacloud file ID of the file that was used to extract the data
+        file_name: The name of the file that was used to extract the data
+        file_hash: A content hash of the file that was used to extract the data, for de-duplication
 
     Status Workflow:
-        - "in_review": Initial state, awaiting human review
+        - "pending_review": Initial state, awaiting human review
         - "accepted": Data approved and ready for use
         - "rejected": Data rejected, needs re-extraction or manual fix
         - "error": Processing error occurred
@@ -194,8 +201,8 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         ```python
         # Create extracted data for review
         extracted = ExtractedData.create(
-            extracted_data=person_data,
-            status="in_review",
+            data=person_data,
+            status="pending_review",
             confidence={"name": 0.95, "age": 0.87}
         )
 
@@ -212,20 +219,35 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
     data: ExtractedT = Field(
         description="The latest state of the data. Will differ if data has been updated"
     )
-    status: Union[Literal["error", "accepted", "rejected", "in_review"], str] = Field(
-        description="The status of the extracted data"
-    )
-    confidence: Dict[str, Union[float, Dict]] = Field(
+    status: StatusType = Field(description="The status of the extracted data")
+    confidence: Dict[str, Any] = Field(
         default_factory=dict,
         description="Confidence scores, if any, for each primitive field in the original_data data",
     )
+    file_id: Optional[str] = Field(
+        None, description="The ID of the file that was used to extract the data"
+    )
+    file_name: Optional[str] = Field(
+        None, description="The name of the file that was used to extract the data"
+    )
+    file_hash: Optional[str] = Field(
+        None, description="The hash of the file that was used to extract the data"
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default_factory=dict,
+        description="Additional metadata about the extracted data, such as errors, tokens, etc.",
+    )
 
     @classmethod
     def create(
         cls,
-        extracted_data: ExtractedT,
-        status: StatusType = "in_review",
-        confidence: Optional[Dict[str, Union[float, Dict]]] = None,
+        data: ExtractedT,
+        status: StatusType = "pending_review",
+        confidence: Optional[Dict[str, Any]] = None,
+        file_id: Optional[str] = None,
+        file_name: Optional[str] = None,
+        file_hash: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> "ExtractedData[ExtractedT]":
         """
         Create a new ExtractedData instance with sensible defaults.
@@ -234,15 +256,22 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
             extracted_data: The extracted data payload
             status: Initial workflow status
             confidence: Optional confidence scores for fields
+            file_id: The llamacloud file ID of the file that was used to extract the data
+            file_name: The name of the file that was used to extract the data
+            file_hash: A content hash of the file that was used to extract the data, for de-duplication
 
         Returns:
             New ExtractedData instance ready for storage
         """
         return cls(
-            original_data=extracted_data,
-            data=extracted_data,
+            original_data=data,
+            data=data,
             status=status,
             confidence=confidence or {},
+            file_id=file_id,
+            file_name=file_name,
+            file_hash=file_hash,
+            metadata=metadata or {},
         )
 
 

llama_parse/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "llama-parse"
-version = "0.6.45"
+version = "0.6.46"
 description = "Parse files into RAG-Optimized formats."
 authors = ["Logan Markewich <logan@llamaindex.ai>"]
 license = "MIT"

pyproject.toml

@@ -8,7 +8,7 @@ python_version = "3.10"
 
 [tool.poetry]
 name = "llama-cloud-services"
-version = "0.6.45"
+version = "0.6.46"
 description = "Tailored SDK clients for LlamaCloud services."
 authors = ["Logan Markewich <logan@runllama.ai>"]
 license = "MIT"

tests/beta/agent_data/test_agent_data_client.py

@@ -67,7 +67,7 @@ async def test_agent_data_crud_operations():
     agent_data_client = AsyncAgentDataClient(
         client=client,
         type=TestData,
-        collection_name=f"test-collection-{test_id[:8]}",
+        collection=f"test-collection-{test_id[:8]}",
         agent_url_id=LLAMA_DEPLOY_DEPLOYMENT_NAME,
     )
 
@@ -77,21 +77,21 @@ async def test_agent_data_crud_operations():
     created_item = None
     try:
         # Test CREATE
-        created_item = await agent_data_client.create_agent_data(test_data)
+        created_item = await agent_data_client.create_item(test_data)
         assert created_item.data.name == "test-item"
         assert created_item.data.test_id == test_id
         assert created_item.data.value == 42
         assert created_item.id is not None
 
         # Test READ
-        retrieved_item = await agent_data_client.get_agent_data(created_item.id)
+        retrieved_item = await agent_data_client.get_item(created_item.id)
         assert retrieved_item.id == created_item.id
         assert retrieved_item.data.name == "test-item"
         assert retrieved_item.data.test_id == test_id
         assert retrieved_item.data.value == 42
 
         # Test SEARCH
-        search_results = await agent_data_client.search_agent_data(
+        search_results = await agent_data_client.search(
             filter={"test_id": {"eq": test_id}}, page_size=10, include_total=True
         )
         assert len(search_results.items) == 1
@@ -99,7 +99,7 @@ async def test_agent_data_crud_operations():
         assert search_results.total == 1
 
         # Test AGGREGATE
-        aggregate_results = await agent_data_client.aggregate_agent_data(
+        aggregate_results = await agent_data_client.aggregate(
             group_by=["test_id"], count=True
         )
         assert len(aggregate_results.items) == 1
@@ -108,7 +108,7 @@ async def test_agent_data_crud_operations():
 
         # Test UPDATE
         updated_data = TestData(name="updated-item", test_id=test_id, value=84)
-        updated_item = await agent_data_client.update_agent_data(
+        updated_item = await agent_data_client.update_item(
             created_item.id, updated_data
         )
         assert updated_item.data.name == "updated-item"
@@ -116,7 +116,7 @@ async def test_agent_data_crud_operations():
         assert updated_item.id == created_item.id
 
         # Verify update persisted
-        verified_item = await agent_data_client.get_agent_data(created_item.id)
+        verified_item = await agent_data_client.get_item(created_item.id)
         assert verified_item.data.name == "updated-item"
         assert verified_item.data.value == 84
 
@@ -124,6 +124,6 @@ async def test_agent_data_crud_operations():
         # Clean up test data
         if created_item is not None:
             try:
-                await agent_data_client.delete_agent_data(created_item.id)
+                await agent_data_client.delete_item(created_item.id)
             except Exception as e:
                 print(f"Warning: Failed to cleanup test data {created_item.id}: {e}")

tests/beta/agent_data/test_agent_data_schema.py

@@ -70,7 +70,7 @@ def test_extracted_data_create_method():
     extracted = ExtractedData.create(person)
     assert extracted.original_data == person
     assert extracted.data == person
-    assert extracted.status == "in_review"
+    assert extracted.status == "pending_review"
     assert extracted.confidence == {}
 
     # Test with custom values