cua-agent 0.4.34__py3-none-any.whl → 0.4.36__py3-none-any.whl

agent/human_tool/server.py

@@ -1,9 +1,9 @@
 import asyncio
 import uuid
+from dataclasses import asdict, dataclass
 from datetime import datetime
-from typing import Dict, List, Any, Optional
-from dataclasses import dataclass, asdict
 from enum import Enum
+from typing import Any, Dict, List, Optional
 
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
@@ -49,7 +49,7 @@ class CompletionQueue:
         self._queue: Dict[str, CompletionCall] = {}
         self._pending_order: List[str] = []
         self._lock = asyncio.Lock()
-    
+
     async def add_completion(self, messages: List[Dict[str, Any]], model: str) -> str:
         """Add a completion call to the queue."""
         async with self._lock:
@@ -59,42 +59,47 @@ class CompletionQueue:
                 messages=messages,
                 model=model,
                 status=CompletionStatus.PENDING,
-                created_at=datetime.now()
+                created_at=datetime.now(),
             )
             self._queue[call_id] = completion_call
             self._pending_order.append(call_id)
             return call_id
-    
+
     async def get_pending_calls(self) -> List[Dict[str, Any]]:
         """Get all pending completion calls."""
         async with self._lock:
             pending_calls = []
             for call_id in self._pending_order:
-                if call_id in self._queue and self._queue[call_id].status == CompletionStatus.PENDING:
+                if (
+                    call_id in self._queue
+                    and self._queue[call_id].status == CompletionStatus.PENDING
+                ):
                     call = self._queue[call_id]
-                    pending_calls.append({
-                        "id": call.id,
-                        "model": call.model,
-                        "created_at": call.created_at.isoformat(),
-                        "messages": call.messages
-                    })
+                    pending_calls.append(
+                        {
+                            "id": call.id,
+                            "model": call.model,
+                            "created_at": call.created_at.isoformat(),
+                            "messages": call.messages,
+                        }
+                    )
             return pending_calls
-    
+
     async def get_call_status(self, call_id: str) -> Optional[Dict[str, Any]]:
         """Get the status of a specific completion call."""
         async with self._lock:
             if call_id not in self._queue:
                 return None
-            
+
             call = self._queue[call_id]
             result = {
                 "id": call.id,
                 "status": call.status.value,
                 "created_at": call.created_at.isoformat(),
                 "model": call.model,
-                "messages": call.messages
+                "messages": call.messages,
             }
-            
+
             if call.completed_at:
                 result["completed_at"] = call.completed_at.isoformat()
             if call.response:
@@ -103,69 +108,74 @@ class CompletionQueue:
                 result["tool_calls"] = call.tool_calls
             if call.error:
                 result["error"] = call.error
-                
+
             return result
-    
-    async def complete_call(self, call_id: str, response: Optional[str] = None, tool_calls: Optional[List[Dict[str, Any]]] = None) -> bool:
+
+    async def complete_call(
+        self,
+        call_id: str,
+        response: Optional[str] = None,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ) -> bool:
         """Mark a completion call as completed with a response or tool calls."""
         async with self._lock:
             if call_id not in self._queue:
                 return False
-            
+
             call = self._queue[call_id]
             if call.status != CompletionStatus.PENDING:
                 return False
-            
+
             call.status = CompletionStatus.COMPLETED
             call.completed_at = datetime.now()
             call.response = response
             call.tool_calls = tool_calls
-            
+
             # Remove from pending order
             if call_id in self._pending_order:
                 self._pending_order.remove(call_id)
-            
+
             return True
-    
+
     async def fail_call(self, call_id: str, error: str) -> bool:
         """Mark a completion call as failed with an error."""
         async with self._lock:
             if call_id not in self._queue:
                 return False
-            
+
             call = self._queue[call_id]
             if call.status != CompletionStatus.PENDING:
                 return False
-            
+
             call.status = CompletionStatus.FAILED
             call.completed_at = datetime.now()
             call.error = error
-            
+
             # Remove from pending order
             if call_id in self._pending_order:
                 self._pending_order.remove(call_id)
-            
+
             return True
-    
+
     async def wait_for_completion(self, call_id: str, timeout: float = 300.0) -> Optional[str]:
         """Wait for a completion call to be completed and return the response."""
         start_time = asyncio.get_event_loop().time()
-        
+
         while True:
             status = await self.get_call_status(call_id)
             if not status:
                 return None
-            
+
             if status["status"] == CompletionStatus.COMPLETED.value:
                 return status.get("response")
             elif status["status"] == CompletionStatus.FAILED.value:
                 raise Exception(f"Completion failed: {status.get('error', 'Unknown error')}")
-            
+
             # Check timeout
             if asyncio.get_event_loop().time() - start_time > timeout:
                 await self.fail_call(call_id, "Timeout waiting for human response")
                 raise TimeoutError("Timeout waiting for human response")
-            
+
             # Wait a bit before checking again
             await asyncio.sleep(0.5)
 
@@ -204,9 +214,7 @@ async def get_status(call_id: str):
 async def complete_call(call_id: str, response: CompletionResponse):
     """Complete a call with a human response."""
     success = await completion_queue.complete_call(
-        call_id, 
-        response=response.response, 
-        tool_calls=response.tool_calls
+        call_id, response=response.response, tool_calls=response.tool_calls
     )
     if success:
         return {"status": "success", "message": "Call completed"}
@@ -219,7 +227,9 @@ async def fail_call(call_id: str, error: Dict[str, str]):
     """Mark a call as failed."""
     success = await completion_queue.fail_call(call_id, error.get("error", "Unknown error"))
     if not success:
-        raise HTTPException(status_code=404, detail="Completion call not found or already completed")
+        raise HTTPException(
+            status_code=404, detail="Completion call not found or already completed"
+        )
     return {"status": "failed"}
 
 
@@ -231,4 +241,5 @@ async def root():
 
 if __name__ == "__main__":
     import uvicorn
+
     uvicorn.run(app, host="0.0.0.0", port=8002)