cua-agent 0.4.12__py3-none-any.whl → 0.4.14__py3-none-any.whl

agent/human_tool/server.py

@@ -0,0 +1,234 @@
+import asyncio
+import uuid
+from datetime import datetime
+from typing import Dict, List, Any, Optional
+from dataclasses import dataclass, asdict
+from enum import Enum
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+
+class CompletionStatus(str, Enum):
+    PENDING = "pending"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+@dataclass
+class CompletionCall:
+    id: str
+    messages: List[Dict[str, Any]]
+    model: str
+    status: CompletionStatus
+    created_at: datetime
+    completed_at: Optional[datetime] = None
+    response: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    error: Optional[str] = None
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: str = "function"
+    function: Dict[str, Any]
+
+
+class CompletionRequest(BaseModel):
+    messages: List[Dict[str, Any]]
+    model: str
+
+
+class CompletionResponse(BaseModel):
+    response: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+
+
+class CompletionQueue:
+    def __init__(self):
+        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:
+            call_id = str(uuid.uuid4())
+            completion_call = CompletionCall(
+                id=call_id,
+                messages=messages,
+                model=model,
+                status=CompletionStatus.PENDING,
+                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:
+                    call = self._queue[call_id]
+                    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
+            }
+            
+            if call.completed_at:
+                result["completed_at"] = call.completed_at.isoformat()
+            if call.response:
+                result["response"] = call.response
+            if call.tool_calls:
+                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:
+        """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)
+
+
+# Global queue instance
+completion_queue = CompletionQueue()
+
+# FastAPI app
+app = FastAPI(title="Human Completion Server", version="1.0.0")
+
+
+@app.post("/queue", response_model=Dict[str, str])
+async def queue_completion(request: CompletionRequest):
+    """Add a completion request to the queue."""
+    call_id = await completion_queue.add_completion(request.messages, request.model)
+    return {"id": call_id, "status": "queued"}
+
+
+@app.get("/pending")
+async def list_pending():
+    """List all pending completion calls."""
+    pending_calls = await completion_queue.get_pending_calls()
+    return {"pending_calls": pending_calls}
+
+
+@app.get("/status/{call_id}")
+async def get_status(call_id: str):
+    """Get the status of a specific completion call."""
+    status = await completion_queue.get_call_status(call_id)
+    if not status:
+        raise HTTPException(status_code=404, detail="Completion call not found")
+    return status
+
+
+@app.post("/complete/{call_id}")
+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
+    )
+    if success:
+        return {"status": "success", "message": "Call completed"}
+    else:
+        raise HTTPException(status_code=404, detail="Call not found or already completed")
+
+
+@app.post("/fail/{call_id}")
+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")
+    return {"status": "failed"}
+
+
+@app.get("/")
+async def root():
+    """Root endpoint."""
+    return {"message": "Human Completion Server is running"}
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8002)