llama-deploy-appserver 0.3.27__py3-none-any.whl → 0.4.0__py3-none-any.whl

llama_deploy/appserver/workflow_store/agent_data_store.py

@@ -2,10 +2,9 @@ import asyncio
 import logging
 import os
 import sys
-from typing import Any, List, cast
+from typing import Any, List
 
-from llama_cloud.client import AsyncLlamaCloud, httpx
-from llama_cloud_services.beta.agent_data import AsyncAgentDataClient
+import httpx
 from llama_deploy.appserver.settings import ApiserverSettings
 from llama_deploy.core.client.ssl_util import get_httpx_verify_param
 from llama_deploy.core.deployment_config import DeploymentConfig
@@ -21,8 +20,12 @@ else:
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_BASE_URL = "https://api.cloud.llamaindex.ai"
+
 
 class AgentDataStore(AbstractWorkflowStore):
+    """Workflow store backed by LlamaCloud Agent Data API using raw httpx."""
+
     def __init__(
         self, settings: DeploymentConfig, server_settings: ApiserverSettings
     ) -> None:
@@ -37,71 +40,149 @@ class AgentDataStore(AbstractWorkflowStore):
             agent_url_id = settings.name
 
         self.settings = settings
-        project_id = os.getenv("LLAMA_DEPLOY_PROJECT_ID")
-        self.client = AsyncAgentDataClient(
-            type=PersistentHandler,
-            collection=collection,
-            agent_url_id=agent_url_id,
-            client=AsyncLlamaCloud(
-                base_url=os.getenv("LLAMA_CLOUD_BASE_URL"),
-                token=os.getenv("LLAMA_CLOUD_API_KEY"),
-                httpx_client=httpx.AsyncClient(
-                    headers={"Project-Id": project_id} if project_id else None,
-                    verify=get_httpx_verify_param(),
-                ),
-            ),
-        )
+        self.collection = collection
+        self.deployment_name = agent_url_id
+
+        self.base_url = os.getenv("LLAMA_CLOUD_BASE_URL", DEFAULT_BASE_URL).rstrip("/")
+        self.api_key = os.getenv("LLAMA_CLOUD_API_KEY")
+        self.project_id = os.getenv("LLAMA_DEPLOY_PROJECT_ID")
+
         self.lock = AsyncKeyedLock()
         # workflow id -> agent data id
         self.cache = LRUCache[str, str](maxsize=1024)
 
+    def _get_headers(self) -> dict[str, str]:
+        """Build HTTP headers for API requests."""
+        headers = {
+            "Content-Type": "application/json",
+        }
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        if self.project_id:
+            headers["Project-Id"] = self.project_id
+        return headers
+
+    def _get_client(self) -> httpx.AsyncClient:
+        """Create a new httpx client."""
+        return httpx.AsyncClient(
+            headers=self._get_headers(),
+            verify=get_httpx_verify_param(),
+        )
+
     @override
     async def query(self, query: HandlerQuery) -> List[PersistentHandler]:
         filters = self._build_filters(query)
-        results = await self.client.search(
-            filter=filters,
-            page_size=1000,
-        )
-        return [x.data for x in results.items]
+        search_request = {
+            "deployment_name": self.deployment_name,
+            "collection": self.collection,
+            "filter": filters,
+            "page_size": 1000,
+        }
+        async with self._get_client() as client:
+            response = await client.post(
+                f"{self.base_url}/api/v1/beta/agent-data/:search",
+                json=search_request,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        items = data.get("items", [])
+        return [PersistentHandler(**item["data"]) for item in items]
 
     @override
     async def update(self, handler: PersistentHandler) -> None:
         async with self.lock.acquire(handler.handler_id):
             id = await self._get_item_id(handler)
             if id is None:
-                item = await self.client.create_item(
-                    data=handler,
-                )
-                if item.id is None:
+                item = await self._create_item(handler)
+                item_id = item.get("id")
+                if item_id is None:
                     raise ValueError(f"Failed to create handler {handler.handler_id}")
-                self.cache.set(handler.handler_id, item.id)
+                self.cache.set(handler.handler_id, item_id)
             else:
-                await self.client.update_item(
-                    item_id=id,
-                    data=handler,
-                )
+                await self._update_item(id, handler)
 
     @override
     async def delete(self, query: HandlerQuery) -> int:
         filters = self._build_filters(query)
-        results = await self.client.search(filter=filters, page_size=1000)
+        search_request = {
+            "deployment_name": self.deployment_name,
+            "collection": self.collection,
+            "filter": filters,
+            "page_size": 1000,
+        }
+        async with self._get_client() as client:
+            response = await client.post(
+                f"{self.base_url}/api/v1/beta/agent-data/:search",
+                json=search_request,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        items = data.get("items", [])
         await asyncio.gather(
-            *[self.client.delete_item(item_id=x.id) for x in results.items if x.id]
+            *[self._delete_item(item["id"]) for item in items if item.get("id")]
         )
-        return len(results.items)
+        return len(items)
+
+    async def _create_item(self, handler: PersistentHandler) -> dict[str, Any]:
+        """Create a new agent data item."""
+        create_request = {
+            "deployment_name": self.deployment_name,
+            "collection": self.collection,
+            "data": handler.model_dump(mode="json"),
+        }
+        async with self._get_client() as client:
+            response = await client.post(
+                f"{self.base_url}/api/v1/beta/agent-data",
+                json=create_request,
+            )
+            response.raise_for_status()
+            return response.json()
+
+    async def _update_item(self, item_id: str, handler: PersistentHandler) -> None:
+        """Update an existing agent data item."""
+        update_request = {
+            "data": handler.model_dump(mode="json"),
+        }
+        async with self._get_client() as client:
+            response = await client.put(
+                f"{self.base_url}/api/v1/beta/agent-data/{item_id}",
+                json=update_request,
+            )
+            response.raise_for_status()
+
+    async def _delete_item(self, item_id: str) -> None:
+        """Delete an agent data item."""
+        async with self._get_client() as client:
+            response = await client.delete(
+                f"{self.base_url}/api/v1/beta/agent-data/{item_id}",
+            )
+            response.raise_for_status()
 
     async def _get_item_id(self, handler: PersistentHandler) -> str | None:
         cached_id = self.cache.get(handler.handler_id)
         if cached_id is not None:
             return cached_id
         search_filter = {"handler_id": {"eq": handler.handler_id}}
-        results = await self.client.search(
-            filter=cast(Any, search_filter),
-            page_size=1,
-        )
-        if not results.items:
+        search_request = {
+            "deployment_name": self.deployment_name,
+            "collection": self.collection,
+            "filter": search_filter,
+            "page_size": 1,
+        }
+        async with self._get_client() as client:
+            response = await client.post(
+                f"{self.base_url}/api/v1/beta/agent-data/:search",
+                json=search_request,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        items = data.get("items", [])
+        if not items:
             return None
-        id = results.items[0].id
+        id = items[0].get("id")
         if id is None:
             return None
         self.cache.set(handler.handler_id, id)
@@ -121,4 +202,15 @@ class AgentDataStore(AbstractWorkflowStore):
             filters["status"] = {
                 "includes": query.status_in,
             }
+        if query.is_idle is not None:
+            if query.is_idle:
+                # Filter for handlers where idle_since is set (any valid datetime)
+                filters["idle_since"] = {
+                    "gte": "1970-01-01T00:00:00Z",
+                }
+            else:
+                # Filter for handlers where idle_since is not set (null)
+                filters["idle_since"] = {
+                    "eq": None,
+                }
         return filters

{llama_deploy_appserver-0.3.27.dist-info → llama_deploy_appserver-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: llama-deploy-appserver
-Version: 0.3.27
+Version: 0.4.0
 Summary: Application server components for LlamaDeploy
 Author: Massimiliano Pippi, Adrian Lyjak
 Author-email: Massimiliano Pippi <mpippi@gmail.com>, Adrian Lyjak <adrianlyjak@gmail.com>
@@ -9,14 +9,13 @@ Requires-Dist: llama-index-workflows[server]>=2.12.0,<2.13.0
 Requires-Dist: pydantic-settings>=2.10.1
 Requires-Dist: fastapi>=0.100.0
 Requires-Dist: websockets>=12.0
-Requires-Dist: llama-deploy-core>=0.3.27,<0.4.0
+Requires-Dist: llama-deploy-core>=0.4.0,<0.5.0
 Requires-Dist: httpx>=0.24.0,<1.0.0
 Requires-Dist: prometheus-fastapi-instrumentator>=7.1.0
 Requires-Dist: packaging>=25.0
 Requires-Dist: structlog>=25.4.0
 Requires-Dist: rich>=14.1.0
 Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: llama-cloud-services>=0.6.60
 Requires-Dist: watchfiles>=1.1.0
 Requires-Dist: uvicorn>=0.35.0
 Requires-Dist: typing-extensions>=4.15.0 ; python_full_version < '3.12'

{llama_deploy_appserver-0.3.27.dist-info → llama_deploy_appserver-0.4.0.dist-info}/RECORD

@@ -16,9 +16,9 @@ llama_deploy/appserver/settings.py,sha256=aa4512d2f1f28b8ee7d3fedc8c61f77bce9f80
 llama_deploy/appserver/stats.py,sha256=1f3989f6705a6de3e4d61ee8cdd189fbe04a2c53ec5e720b2e5168acc331427f,691
 llama_deploy/appserver/types.py,sha256=4edc991aafb6b8497f068d12387455df292da3ff8440223637641ab1632553ec,2133
 llama_deploy/appserver/workflow_loader.py,sha256=c1422bda79423528f1f22fe793c9c43cb411dd883a05a80c5a161f04efd9489e,15713
-llama_deploy/appserver/workflow_store/agent_data_store.py,sha256=ecf11a1454b13b7d618e720c8ba3be41833f322fdddfacc8356dfe34858d3a30,4421
+llama_deploy/appserver/workflow_store/agent_data_store.py,sha256=816710c9957c37792ce38e97a2c57cd8079953d55489b445675646ff74b4e13e,7710
 llama_deploy/appserver/workflow_store/keyed_lock.py,sha256=72bcfafbce56d5b36d53aff764b573c2dca2b3f5bc59f2d8baa80be0e4db6e34,1037
 llama_deploy/appserver/workflow_store/lru_cache.py,sha256=10b7d69e4be7d929d9dac009b59635b20fbed4603fb004bc35cbdc3ce538af8b,1454
-llama_deploy_appserver-0.3.27.dist-info/WHEEL,sha256=66530aef82d5020ef5af27ae0123c71abb9261377c5bc519376c671346b12918,79
-llama_deploy_appserver-0.3.27.dist-info/METADATA,sha256=5548b300b61536a59918488db53f55e673dba88de8dcfc2b79cb55c2f8b02beb,1163
-llama_deploy_appserver-0.3.27.dist-info/RECORD,,
+llama_deploy_appserver-0.4.0.dist-info/WHEEL,sha256=66530aef82d5020ef5af27ae0123c71abb9261377c5bc519376c671346b12918,79
+llama_deploy_appserver-0.4.0.dist-info/METADATA,sha256=41f69121b36ccb0494337a86cf0b83ec247e69ce5c39235cc7c9af0609fff919,1117
+llama_deploy_appserver-0.4.0.dist-info/RECORD,,