Spaces:
Sleeping
Sleeping
File size: 4,749 Bytes
24a7f55 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
#!/usr/bin/env python3
"""
Main entry point for the Enhanced AI Agentic Browser Agent Architecture.
This module initializes the FastAPI application and all core components.
"""
import asyncio
import logging
import os
from typing import Dict, List, Optional
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from src.orchestrator import AgentOrchestrator
# Load environment variables
load_dotenv()
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
# Initialize FastAPI application
app = FastAPI(
title="Enhanced AI Agentic Browser Agent",
description="A robust, scalable, and intelligent system for automating complex web tasks.",
version="1.0.0",
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize the Agent Orchestrator
orchestrator = None
@app.on_event("startup")
async def startup_event():
"""Initialize components on application startup."""
global orchestrator
orchestrator = await AgentOrchestrator.initialize()
logger.info("Agent Orchestrator initialized successfully")
@app.on_event("shutdown")
async def shutdown_event():
"""Clean up resources on application shutdown."""
if orchestrator:
await orchestrator.shutdown()
logger.info("Agent Orchestrator shut down successfully")
# Model definitions
class TaskRequest(BaseModel):
"""Model for task execution requests."""
task_description: str
urls: Optional[List[str]] = None
human_assisted: bool = False
max_retries: int = 3
timeout: int = 300 # seconds
class TaskResponse(BaseModel):
"""Model for task execution responses."""
task_id: str
status: str
message: str
class TaskResult(BaseModel):
"""Model for task execution results."""
task_id: str
status: str
result: Optional[Dict] = None
error: Optional[str] = None
# API routes
@app.get("/")
async def root():
"""Health check endpoint."""
return {"status": "online", "service": "Enhanced AI Agentic Browser Agent"}
@app.post("/tasks", response_model=TaskResponse)
async def create_task(task: TaskRequest, background_tasks: BackgroundTasks):
"""
Create and start a new task.
Args:
task: TaskRequest object containing task details
background_tasks: FastAPI background tasks object
Returns:
TaskResponse: Task creation response with task ID
"""
task_id = await orchestrator.create_task(task.dict())
background_tasks.add_task(orchestrator.execute_task, task_id)
return {"task_id": task_id, "status": "started", "message": "Task created and started"}
@app.get("/tasks/{task_id}", response_model=TaskResult)
async def get_task(task_id: str):
"""
Get the status and result of a task.
Args:
task_id: The ID of the task to retrieve
Returns:
TaskResult: Task result object with status and data
"""
result = await orchestrator.get_task_status(task_id)
if not result:
raise HTTPException(status_code=404, detail="Task not found")
return result
@app.delete("/tasks/{task_id}")
async def cancel_task(task_id: str):
"""
Cancel a running task.
Args:
task_id: The ID of the task to cancel
Returns:
Dict: Cancellation status
"""
success = await orchestrator.cancel_task(task_id)
if not success:
raise HTTPException(status_code=404, detail="Task not found or already completed")
return {"status": "cancelled", "message": "Task cancelled successfully"}
@app.websocket("/ws/tasks/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str):
"""
WebSocket endpoint for real-time task updates.
Args:
websocket: WebSocket connection
task_id: Task ID to subscribe to
"""
await websocket.accept()
try:
await orchestrator.register_websocket(task_id, websocket)
while True:
# Keep the connection open and handle incoming messages
data = await websocket.receive_text()
if data == "ping":
await websocket.send_text("pong")
except WebSocketDisconnect:
await orchestrator.unregister_websocket(task_id, websocket)
if __name__ == "__main__":
uvicorn.run("src.main:app", host="0.0.0.0", port=8000, reload=True)
|