Update app.py

app.py

@@ -3,6 +3,18 @@ from fastapi.responses import HTMLResponse
 from fastapi.requests import Request
 from fastapi.templating import Jinja2Templates
 from fastapi import Body, FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+from loguru import logger
+import aiohttp
+import uvicorn
+import asyncio
+import os
+import uuid
+import toml
+from datetime import datetime
+from json import dumps
 
 
 templates = Jinja2Templates(directory="templates")
@@ -12,6 +24,71 @@ import os
 
 app = FastAPI()
 
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+templates = Jinja2Templates(directory="templates")
+
+class Task(BaseModel):
+    id: str
+    prompt: str
+    created_at: datetime
+    status: str
+    steps: list = []
+
+    def model_dump(self, *args, **kwargs):
+        data = super().model_dump(*args, **kwargs)
+        data["created_at"] = self.created_at.isoformat()
+        return data
+
+class TaskManager:
+    def __init__(self):
+        self.tasks = {}
+        self.queues = {}
+
+    def create_task(self, prompt: str) -> Task:
+        task_id = str(uuid.uuid4())
+        task = Task(
+            id=task_id, prompt=prompt, created_at=datetime.now(), status="pending"
+        )
+        self.tasks[task_id] = task
+        self.queues[task_id] = asyncio.Queue()
+        return task
+
+    async def update_task_step(self, task_id: str, step: int, result: str, step_type: str = "step"):
+        if task_id in self.tasks:
+            task = self.tasks[task_id]
+            task.steps.append({"step": step, "result": result, "type": step_type})
+            await self.queues[task_id].put({"type": step_type, "step": step, "result": result})
+            await self.queues[task_id].put({"type": "status", "status": task.status, "steps": task.steps})
+
+    async def complete_task(self, task_id: str):
+        if task_id in self.tasks:
+            task = self.tasks[task_id]
+            task.status = "completed"
+            await self.queues[task_id].put({"type": "status", "status": task.status, "steps": task.steps})
+            await self.queues[task_id].put({"type": "complete"})
+
+    async def fail_task(self, task_id: str, error: str):
+        if task_id in self.tasks:
+            self.tasks[task_id].status = f"failed: {error}"
+            await self.queues[task_id].put({"type": "error", "message": error})
+
+task_manager = TaskManager()
+
+@app.post("/tasks")
+async def create_task(prompt: str = Body(..., embed=True)):
+    task = task_manager.create_task(prompt)
+    asyncio.create_task(run_task(task.id, prompt))
+    return {"task_id": task.id}
+
+
 @app.get("/", response_class=HTMLResponse)
 async def chat(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})