bump to 0.6.44

llama_cloud_services/beta/agent_data/__init__.py

@@ -0,0 +1,19 @@
+from .schema import (
+    TypedAgentData,
+    ExtractedData,
+    TypedAgentDataItems,
+    StatusType,
+    ExtractedT,
+    AgentDataT,
+)
+from .client import AsyncAgentDataClient
+
+__all__ = [
+    "TypedAgentData",
+    "AsyncAgentDataClient",
+    "ExtractedData",
+    "TypedAgentDataItems",
+    "StatusType",
+    "ExtractedT",
+    "AgentDataT",
+]

llama_cloud_services/beta/agent_data/client.py

@@ -0,0 +1,267 @@
+import os
+from typing import Dict, Generic, List, Optional, Type
+
+from llama_cloud import FilterOperation
+from llama_cloud.client import AsyncLlamaCloud
+from tenacity import (
+    WrappedFn,
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    retry_if_exception_type,
+)
+import httpx
+
+from .schema import (
+    AgentDataT,
+    TypedAgentData,
+    TypedAgentDataItems,
+    TypedAggregateGroup,
+    TypedAggregateGroupItems,
+)
+
+
+def agent_data_retry(func: WrappedFn) -> WrappedFn:
+    """
+    Decorator that adds automatic retry logic to agent data API calls.
+
+    Applies exponential backoff retry strategy for common network-related exceptions:
+    - Up to 3 retry attempts
+    - Exponential wait time between 0.5s and 10s
+    - Retries on timeout, connection, and HTTP status errors
+
+    This ensures resilient API communication in distributed environments where
+    temporary network issues or service unavailability may occur.
+    """
+    return retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(min=0.5, max=10),
+        retry=retry_if_exception_type(
+            (httpx.TimeoutException, httpx.ConnectError, httpx.HTTPStatusError)
+        ),
+    )(func)
+
+
+def get_default_agent_id() -> Optional[str]:
+    """
+    Retrieve the default agent ID from environment variables.
+
+    Returns:
+        The value of LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable,
+        or None if not set
+
+    Note:
+        This provides a convenient way to configure agent ID globally
+        via environment variables instead of passing it explicitly
+        to each client instance.
+    """
+    return os.getenv("LLAMA_DEPLOY_DEPLOYMENT_NAME")
+
+
+class AsyncAgentDataClient(Generic[AgentDataT]):
+    """
+    Async client for managing agent-generated structured data with type safety.
+
+    This client provides a high-level interface for CRUD operations, searching, and
+    aggregation of structured data created by agents. It enforces type safety by
+    validating all data against a specified Pydantic model type.
+
+    The client is generic over AgentDataT, which must be a Pydantic BaseModel that
+    defines the structure of your agent's data output.
+
+    Example:
+        ```python
+        from pydantic import BaseModel
+        from llama_cloud.client import AsyncLlamaCloud
+        from llama_cloud_services.beta.agent_data import AsyncAgentDataClient
+
+        class ExtractedPerson(BaseModel):
+            name: str
+            age: int
+            email: str
+
+        # Initialize client
+        llama_client = AsyncLlamaCloud(token="your-api-key")
+        agent_client = AsyncAgentDataClient(
+            client=llama_client,
+            type=ExtractedPerson,
+            collection_name="extracted_people",
+            agent_url_id="person-extraction-agent"
+        )
+
+        # Create data
+        person = ExtractedPerson(name="John Doe", age=30, email="john@example.com")
+        result = await agent_client.create_agent_data(person)
+
+        # Search data
+        results = await agent_client.search_agent_data(
+            filter={"age": FilterOperation(gt=25)},
+            order_by="data.name",
+            page_size=20
+        )
+        ```
+
+    Type Parameters:
+        AgentDataT: Pydantic BaseModel type that defines the structure of agent data
+    """
+
+    def __init__(
+        self,
+        client: AsyncLlamaCloud,
+        type: Type[AgentDataT],
+        collection_name: str = "default",
+        agent_url_id: 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.
+                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
+                url of an agent within the llama cloud platform. If not provided,
+                will attempt to use the LLAMA_DEPLOY_DEPLOYMENT_NAME environment
+                variable. Data can only be added to an already existing agent in the
+                platform.
+
+        Raises:
+            ValueError: If agent_url_id is not provided and the
+                LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable is not set
+
+        Note:
+            The client automatically applies retry logic to all API calls with
+            exponential backoff for timeout, connection, and HTTP status errors.
+        """
+
+        self.agent_url_id = agent_url_id or get_default_agent_id()
+        if not self.agent_url_id:
+            raise ValueError(
+                "Agent ID is required, or set the LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable"
+            )
+
+        self.collection_name = collection_name
+        self.client = client
+        self.type = type
+
+    @agent_data_retry
+    async def get_agent_data(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]:
+        raw_data = await self.client.beta.create_agent_data(
+            agent_slug=self.agent_url_id,
+            collection=self.collection_name,
+            data=data.model_dump(),
+        )
+        return TypedAgentData.from_raw(raw_data, validator=self.type)
+
+    @agent_data_retry
+    async def update_agent_data(
+        self, item_id: str, data: AgentDataT
+    ) -> TypedAgentData[AgentDataT]:
+        raw_data = await self.client.beta.update_agent_data(
+            item_id=item_id,
+            data=data.model_dump(),
+        )
+        return TypedAgentData.from_raw(raw_data, validator=self.type)
+
+    @agent_data_retry
+    async def delete_agent_data(self, item_id: str) -> None:
+        await self.client.beta.delete_agent_data(item_id=item_id)
+
+    @agent_data_retry
+    async def search_agent_data(
+        self,
+        filter: Optional[Dict[str, Optional[FilterOperation]]] = None,
+        order_by: Optional[str] = None,
+        offset: Optional[int] = None,
+        page_size: Optional[int] = None,
+        include_total: bool = False,
+    ) -> TypedAgentDataItems[AgentDataT]:
+        """
+        Search agent data with filtering, sorting, and pagination.
+        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
+            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
+            page_size: Maximum number of items to return per page. Defaults to 10.
+            offset: Number of items to skip from the beginning. Defaults to 0.
+            include_total: Whether to include the total count in the response. Defaults to False to improve performance. It's recommended to only request on the first page.
+        """
+        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,
+            filter=filter,
+            order_by=order_by,
+            offset=offset,
+            page_size=page_size,
+            include_total=include_total,
+        )
+        return TypedAgentDataItems(
+            items=[
+                TypedAgentData.from_raw(item, validator=self.type) for item in raw.items
+            ],
+            has_more=raw.next_page_token is not None,
+            total=raw.total_size,
+        )
+
+    @agent_data_retry
+    async def aggregate_agent_data(
+        self,
+        filter: Optional[Dict[str, Optional[FilterOperation]]] = None,
+        group_by: Optional[List[str]] = None,
+        count: Optional[bool] = None,
+        first: Optional[bool] = None,
+        order_by: Optional[str] = None,
+        offset: Optional[int] = None,
+        page_size: Optional[int] = None,
+    ) -> TypedAggregateGroupItems[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.
+            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.
+            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,
+            page_size=page_size,
+            filter=filter,
+            order_by=order_by,
+            group_by=group_by,
+            count=count,
+            first=first,
+            offset=offset,
+        )
+        return TypedAggregateGroupItems(
+            items=[
+                TypedAggregateGroup.from_raw(item, validator=self.type)
+                for item in raw.items
+            ],
+            has_more=raw.next_page_token is not None,
+            total=raw.total_size,
+        )

llama_cloud_services/beta/agent_data/schema.py

@@ -0,0 +1,357 @@
+"""
+Agent Data API Schema Definitions
+
+This module provides typed wrappers around the raw LlamaCloud agent data API,
+enabling type-safe interactions with agent-generated structured data.
+
+The agent data API serves as a persistent storage system for structured data
+produced by LlamaCloud agents (particularly extraction agents). It provides
+CRUD operations, search capabilities, filtering, and aggregation functionality
+for managing agent-generated data at scale.
+
+Key Concepts:
+- Agent Slug: Unique identifier for an agent instance
+- Collection: Named grouping of data within an agent (defaults to "default"). Data within a collection should be of the same type.
+- Agent Data: Individual structured data records with metadata and timestamps
+
+Example Usage:
+    ```python
+    from pydantic import BaseModel
+
+    class Person(BaseModel):
+        name: str
+        age: int
+
+    client = AsyncAgentDataClient(
+        client=async_llama_cloud,
+        type=Person,
+        collection="people",
+        agent_url_id="my-extraction-agent-xyz"
+    )
+
+    # Create typed data
+    person = Person(name="John", age=30)
+    result = await client.create_agent_data(person)
+    print(result.data.name)  # Type-safe access
+    ```
+"""
+
+from datetime import datetime
+from llama_cloud.types.agent_data import AgentData
+from llama_cloud.types.aggregate_group import AggregateGroup
+from pydantic import BaseModel, Field
+from typing import (
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Dict,
+    Type,
+    TypeVar,
+    Union,
+    Any,
+)
+
+
+# Type variable for user-defined data models
+AgentDataT = TypeVar("AgentDataT", bound=BaseModel)
+
+# Type variable for extracted data (can be dict or Pydantic model)
+ExtractedT = TypeVar("ExtractedT", bound=Union[BaseModel, dict])
+
+# Status types for extracted data workflow
+StatusType = Union[Literal["error", "accepted", "rejected", "in_review"], str]
+
+
+class TypedAgentData(BaseModel, Generic[AgentDataT]):
+    """
+    Type-safe wrapper for agent data records.
+
+    This class represents a single data record stored in the agent data API,
+    combining the structured data payload with metadata about when and where
+    it was created.
+
+    Attributes:
+        id: Unique identifier for this data record
+        agent_url_id: Identifier of the agent that created this data
+        collection: Named collection within the agent (used for organization)
+        data: The actual structured data payload (typed as AgentDataT)
+        created_at: Timestamp when the record was first created
+        updated_at: Timestamp when the record was last modified
+
+    Example:
+        ```python
+        # Access typed data
+        person_data: TypedAgentData[Person] = await client.get_agent_data(id)
+        print(person_data.data.name)  # Type-safe access to Person fields
+        print(person_data.created_at)  # Access metadata
+        ```
+    """
+
+    id: Optional[str] = Field(description="Unique identifier for this data record")
+    agent_url_id: str = Field(
+        description="Identifier of the agent that created this data"
+    )
+    collection: Optional[str] = Field(
+        description="Named collection within the agent for data organization"
+    )
+    data: AgentDataT = Field(description="The structured data payload")
+    created_at: Optional[datetime] = Field(description="When this record was created")
+    updated_at: Optional[datetime] = Field(
+        description="When this record was last modified"
+    )
+
+    @classmethod
+    def from_raw(
+        cls, raw_data: AgentData, validator: Type[AgentDataT]
+    ) -> "TypedAgentData[AgentDataT]":
+        """
+        Convert raw API response to typed agent data.
+
+        Args:
+            raw_data: Raw agent data from the API
+            validator: Pydantic model class to validate the data field
+
+        Returns:
+            TypedAgentData instance with validated data
+        """
+        data: AgentDataT = validator.model_validate(raw_data.data)
+
+        return cls(
+            id=raw_data.id,
+            agent_url_id=raw_data.agent_slug,
+            collection=raw_data.collection,
+            data=data,
+            created_at=raw_data.created_at,
+            updated_at=raw_data.updated_at,
+        )
+
+
+class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
+    """
+    Paginated collection of agent data records.
+
+    This class represents a page of search results from the agent data API,
+    providing both the data records and pagination metadata.
+
+    Attributes:
+        items: List of agent data records in this page
+        total: Total number of records matching the query (only present if requested)
+        has_more: Whether there are more records available beyond this page
+
+    Example:
+        ```python
+        # Search with pagination
+        results = await client.search_agent_data(
+            page_size=10,
+            include_total=True
+        )
+
+        for item in results.items:
+            print(item.data.name)
+
+        if results.has_more:
+            # Load next page
+            next_page = await client.search_agent_data(
+                page_size=10,
+                offset=10
+            )
+        ```
+    """
+
+    items: List[TypedAgentData[AgentDataT]] = Field(
+        description="List of agent data records in this page"
+    )
+    total: Optional[int] = Field(
+        description="Total number of records matching the query (only present if requested)"
+    )
+    has_more: bool = Field(
+        description="Whether there are more records available beyond this page"
+    )
+
+
+class ExtractedData(BaseModel, Generic[ExtractedT]):
+    """
+    Wrapper for extracted data with workflow status tracking.
+
+    This class is designed for extraction workflows where data goes through
+    review and approval stages. It maintains both the original extracted data
+    and the current state after any modifications.
+
+    Attributes:
+        original_data: The data as originally extracted from the source
+        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)
+
+    Status Workflow:
+        - "in_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
+
+    Example:
+        ```python
+        # Create extracted data for review
+        extracted = ExtractedData.create(
+            extracted_data=person_data,
+            status="in_review",
+            confidence={"name": 0.95, "age": 0.87}
+        )
+
+        # Later, after review
+        if extracted.status == "accepted":
+            # Use the data
+            process_person(extracted.data)
+        ```
+    """
+
+    original_data: ExtractedT = Field(
+        description="The original data that was extracted from the document"
+    )
+    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(
+        default_factory=dict,
+        description="Confidence scores, if any, for each primitive field in the original_data data",
+    )
+
+    @classmethod
+    def create(
+        cls,
+        extracted_data: ExtractedT,
+        status: StatusType = "in_review",
+        confidence: Optional[Dict[str, Union[float, Dict]]] = None,
+    ) -> "ExtractedData[ExtractedT]":
+        """
+        Create a new ExtractedData instance with sensible defaults.
+
+        Args:
+            extracted_data: The extracted data payload
+            status: Initial workflow status
+            confidence: Optional confidence scores for fields
+
+        Returns:
+            New ExtractedData instance ready for storage
+        """
+        return cls(
+            original_data=extracted_data,
+            data=extracted_data,
+            status=status,
+            confidence=confidence or {},
+        )
+
+
+class TypedAggregateGroup(BaseModel, Generic[AgentDataT]):
+    """
+    Represents a group of agent data records aggregated by common field values.
+
+    This class is used for grouping and analyzing agent data based on shared
+    characteristics. It's particularly useful for generating summaries and
+    statistics across large datasets.
+
+    Attributes:
+        group_key: The field values that define this group
+        count: Number of records in this group (if count aggregation was requested)
+        first_item: Representative data record from this group (if requested)
+
+    Example:
+        ```python
+        # Group by age range
+        groups = await client.aggregate_agent_data(
+            group_by=["age_range"],
+            count=True,
+            first=True
+        )
+
+        for group in groups.items:
+            print(f"Age range {group.group_key['age_range']}: {group.count} people")
+            if group.first_item:
+                print(f"Example: {group.first_item.name}")
+        ```
+    """
+
+    group_key: Dict[str, Any] = Field(
+        description="The field values that define this group"
+    )
+    count: Optional[int] = Field(
+        description="Number of records in this group (if count aggregation was requested)"
+    )
+    first_item: Optional[AgentDataT] = Field(
+        description="Representative data record from this group (if requested)"
+    )
+
+    @classmethod
+    def from_raw(
+        cls, raw_data: AggregateGroup, validator: Type[AgentDataT]
+    ) -> "TypedAggregateGroup[AgentDataT]":
+        """
+        Convert raw API response to typed aggregate group.
+
+        Args:
+            raw_data: Raw aggregate group from the API
+            validator: Pydantic model class to validate the first_item field
+
+        Returns:
+            TypedAggregateGroup instance with validated first_item
+        """
+        first_item: Optional[AgentDataT] = raw_data.first_item
+        if first_item is not None:
+            first_item = validator.model_validate(first_item)
+
+        return cls(
+            group_key=raw_data.group_key,
+            count=raw_data.count,
+            first_item=first_item,
+        )
+
+
+class TypedAggregateGroupItems(BaseModel, Generic[AgentDataT]):
+    """
+    Paginated collection of aggregate groups.
+
+    This class represents a page of aggregation results from the agent data API,
+    providing both the grouped data and pagination metadata.
+
+    Attributes:
+        items: List of aggregate groups in this page
+        total: Total number of groups matching the query (only present if requested)
+        has_more: Whether there are more groups available beyond this page
+
+    Example:
+        ```python
+        # Get first page of groups
+        results = await client.aggregate_agent_data(
+            group_by=["department"],
+            count=True,
+            page_size=20
+        )
+
+        for group in results.items:
+            dept = group.group_key["department"]
+            print(f"{dept}: {group.count} employees")
+
+        # Load more if needed
+        if results.has_more:
+            next_page = await client.aggregate_agent_data(
+                group_by=["department"],
+                count=True,
+                page_size=20,
+                offset=20
+            )
+        ```
+    """
+
+    items: List[TypedAggregateGroup[AgentDataT]] = Field(
+        description="List of aggregate groups in this page"
+    )
+    total: Optional[int] = Field(
+        description="Total number of groups matching the query (only present if requested)"
+    )
+    has_more: bool = Field(
+        description="Whether there are more groups available beyond this page"
+    )

llama_parse/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "llama-parse"
-version = "0.6.43"
+version = "0.6.44"
 description = "Parse files into RAG-Optimized formats."
 authors = ["Logan Markewich <logan@llamaindex.ai>"]
 license = "MIT"
@@ -13,7 +13,7 @@ packages = [{include = "llama_parse"}]
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
-llama-cloud-services = ">=0.6.43"
+llama-cloud-services = ">=0.6.44"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"

poetry.lock

@@ -1925,14 +1925,14 @@ rapidfuzz = ">=3.9.0,<4.0.0"
 
 [[package]]
 name = "llama-cloud"
-version = "0.1.32"
+version = "0.1.33"
 description = ""
 optional = false
 python-versions = "<4,>=3.8"
 groups = ["main"]
 files = [
-    {file = "llama_cloud-0.1.32-py3-none-any.whl", hash = "sha256:c42b2d5fb24acc8595bcc3626fb84c872909a16ab6d6879a1cb1101b21c238bd"},
-    {file = "llama_cloud-0.1.32.tar.gz", hash = "sha256:cea98241127311ea91f191c3c006aa6558f01d16f9539ed93b24d716b888f10e"},
+    {file = "llama_cloud-0.1.33-py3-none-any.whl", hash = "sha256:35b7d4a30b013f0a343f7e09126b531c697d65bffd4eb4d2d79bf7d65f256178"},
+    {file = "llama_cloud-0.1.33.tar.gz", hash = "sha256:a0bb900d5a6e86f8c767b48686c5253679ad7ca1b57612dc39b0767e57ad3d78"},
 ]
 
 [package.dependencies]
@@ -4623,4 +4623,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<4.0"
-content-hash = "112c1ccc4187a295dc07bebb14753f27f2e153dfdab895f763297f5391d3fb10"
+content-hash = "487717a0bbe7ff67360e3e9f187e15a33c77e4942a7c909cb37b95425a81544c"

pyproject.toml

@@ -8,7 +8,7 @@ python_version = "3.10"
 
 [tool.poetry]
 name = "llama-cloud-services"
-version = "0.6.43"
+version = "0.6.44"
 description = "Tailored SDK clients for LlamaCloud services."
 authors = ["Logan Markewich <logan@runllama.ai>"]
 license = "MIT"
@@ -18,7 +18,7 @@ packages = [{include = "llama_cloud_services"}]
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
 llama-index-core = ">=0.12.0"
-llama-cloud = "==0.1.32"
+llama-cloud = "==0.1.33"
 pydantic = ">=2.8,!=2.10"
 click = "^8.1.7"
 python-dotenv = "^1.0.1"

tests/beta/agent_data/test_agent_data_client.py

@@ -0,0 +1,129 @@
+import os
+import httpx
+import pytest
+import uuid
+from pydantic import BaseModel
+from dotenv import load_dotenv
+from pathlib import Path
+
+from llama_cloud.client import AsyncLlamaCloud
+from llama_cloud_services.beta.agent_data import AsyncAgentDataClient
+
+
+class TrailingSlashHttpxClient(httpx.AsyncClient):
+    """Custom httpx client that ensures all URLs have trailing slashes"""
+
+    async def request(self, method, url, **kwargs):
+        # Convert URL to string and ensure trailing slash
+        url_str = str(url)
+        if not url_str.endswith("/") and "?" not in url_str:
+            url_str += "/"
+        self.headers["Authorization"] = f"Bearer {LLAMA_CLOUD_API_KEY}"
+        kwargs.pop("headers", None)
+        return await super().request(method, url_str, headers=self.headers, **kwargs)
+
+
+# Load environment variables
+def load_test_dotenv():
+    dotenv_path = Path(__file__).parent.parent.parent.parent / ".env.dev"
+    load_dotenv(dotenv_path, override=True)
+
+
+load_test_dotenv()
+
+# Get configuration from environment
+LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
+LLAMA_CLOUD_BASE_URL = os.getenv("LLAMA_CLOUD_BASE_URL")
+LLAMA_DEPLOY_DEPLOYMENT_NAME = os.getenv("LLAMA_DEPLOY_DEPLOYMENT_NAME")
+
+
+class TestData(BaseModel):
+    """Simple test data model for agent data testing"""
+
+    name: str
+    test_id: str
+    value: int
+
+
+# Skip all tests if API key is not set
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not LLAMA_CLOUD_API_KEY or not LLAMA_DEPLOY_DEPLOYMENT_NAME,
+    reason="LLAMA_CLOUD_API_KEY or LLAMA_DEPLOY_DEPLOYMENT_NAME not set",
+)
+async def test_agent_data_crud_operations():
+    """Test basic CRUD operations for agent data with automatic cleanup"""
+    # Create unique test identifier to avoid conflicts
+    test_id = str(uuid.uuid4())
+
+    # Set up client
+    client = AsyncLlamaCloud(
+        token=LLAMA_CLOUD_API_KEY,
+        base_url=LLAMA_CLOUD_BASE_URL,
+        httpx_client=TrailingSlashHttpxClient(timeout=60, follow_redirects=True),
+    )
+
+    # Create agent data client with unique collection name
+    agent_data_client = AsyncAgentDataClient(
+        client=client,
+        type=TestData,
+        collection_name=f"test-collection-{test_id[:8]}",
+        agent_url_id=LLAMA_DEPLOY_DEPLOYMENT_NAME,
+    )
+
+    # Create test data
+    test_data = TestData(name="test-item", test_id=test_id, value=42)
+
+    created_item = None
+    try:
+        # Test CREATE
+        created_item = await agent_data_client.create_agent_data(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)
+        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(
+            filter={"test_id": {"eq": test_id}}, page_size=10, include_total=True
+        )
+        assert len(search_results.items) == 1
+        assert search_results.items[0].data.test_id == test_id
+        assert search_results.total == 1
+
+        # Test AGGREGATE
+        aggregate_results = await agent_data_client.aggregate_agent_data(
+            group_by=["test_id"], count=True
+        )
+        assert len(aggregate_results.items) == 1
+        assert aggregate_results.items[0].group_key["test_id"] == test_id
+        assert aggregate_results.items[0].count == 1
+
+        # Test UPDATE
+        updated_data = TestData(name="updated-item", test_id=test_id, value=84)
+        updated_item = await agent_data_client.update_agent_data(
+            created_item.id, updated_data
+        )
+        assert updated_item.data.name == "updated-item"
+        assert updated_item.data.value == 84
+        assert updated_item.id == created_item.id
+
+        # Verify update persisted
+        verified_item = await agent_data_client.get_agent_data(created_item.id)
+        assert verified_item.data.name == "updated-item"
+        assert verified_item.data.value == 84
+
+    finally:
+        # Clean up test data
+        if created_item is not None:
+            try:
+                await agent_data_client.delete_agent_data(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

@@ -0,0 +1,109 @@
+from datetime import datetime
+from typing import Any, Dict
+
+import pytest
+from llama_cloud.types.agent_data import AgentData
+from llama_cloud.types.aggregate_group import AggregateGroup
+from pydantic import BaseModel, ValidationError
+
+from llama_cloud_services.beta.agent_data.schema import (
+    ExtractedData,
+    TypedAgentData,
+    TypedAggregateGroup,
+)
+
+
+# Test data models
+class Person(BaseModel):
+    name: str
+    age: int
+    email: str
+
+
+class Company(BaseModel):
+    name: str
+    industry: str
+    employees: int
+
+
+def test_typed_agent_data_from_raw():
+    """Test TypedAgentData.from_raw class method."""
+    raw_data = AgentData(
+        id="456",
+        agent_slug="extraction-agent",
+        collection="employees",
+        data={"name": "Jane Smith", "age": 25, "email": "jane@company.com"},
+        created_at=datetime.now(),
+        updated_at=datetime.now(),
+    )
+
+    typed_data = TypedAgentData.from_raw(raw_data, Person)
+
+    assert typed_data.id == "456"
+    assert typed_data.agent_url_id == "extraction-agent"
+    assert typed_data.collection == "employees"
+    assert typed_data.data.name == "Jane Smith"
+    assert typed_data.data.age == 25
+    assert typed_data.data.email == "jane@company.com"
+
+
+def test_typed_agent_data_from_raw_validation_error():
+    """Test TypedAgentData.from_raw with invalid data."""
+    raw_data = AgentData(
+        id="789",
+        agent_slug="test-agent",
+        collection="people",
+        data={"name": "Invalid Person", "age": "not_a_number"},  # Invalid age
+        created_at=datetime.now(),
+        updated_at=datetime.now(),
+    )
+
+    with pytest.raises(ValidationError):
+        TypedAgentData.from_raw(raw_data, Person)
+
+
+def test_extracted_data_create_method():
+    """Test ExtractedData.create class method."""
+    person = Person(name="Created Person", age=35, email="created@example.com")
+
+    # Test with defaults
+    extracted = ExtractedData.create(person)
+    assert extracted.original_data == person
+    assert extracted.data == person
+    assert extracted.status == "in_review"
+    assert extracted.confidence == {}
+
+    # Test with custom values
+    extracted_custom = ExtractedData.create(
+        person, status="accepted", confidence={"name": 0.99}
+    )
+    assert extracted_custom.status == "accepted"
+    assert extracted_custom.confidence["name"] == 0.99
+
+
+def test_extracted_data_with_dict():
+    """Test ExtractedData with dict data instead of Pydantic model."""
+    data_dict = {"name": "Dict Person", "age": 45, "email": "dict@example.com"}
+
+    extracted = ExtractedData[Dict[str, Any]](
+        original_data=data_dict, data=data_dict, status="accepted", confidence={}
+    )
+
+    assert extracted.original_data["name"] == "Dict Person"
+    assert extracted.data["age"] == 45
+
+
+def test_typed_aggregate_group_from_raw():
+    """Test TypedAggregateGroup.from_raw class method."""
+    raw_group = AggregateGroup(
+        group_key={"industry": "Technology"},
+        count=25,
+        first_item={"name": "Tech Corp", "industry": "Technology", "employees": 500},
+    )
+
+    typed_group = TypedAggregateGroup.from_raw(raw_group, Company)
+
+    assert typed_group.group_key["industry"] == "Technology"
+    assert typed_group.count == 25
+    assert typed_group.first_item.name == "Tech Corp"
+    assert typed_group.first_item.employees == 500