| from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
|
| from fastapi.staticfiles import StaticFiles
|
| from fastapi.responses import FileResponse
|
| from pydantic import BaseModel
|
| from agent_workflow import workflow_app
|
| import uuid
|
| import asyncio
|
| from fastapi.responses import FileResponse
|
| from pydantic import BaseModel
|
| from agent_workflow import workflow_app
|
|
|
| app = FastAPI(title="Agentic AI Document Review System", version="1.0.0")
|
|
|
| app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
| @app.get("/")
|
| async def serve_frontend():
|
| return FileResponse("static/index.html")
|
|
|
| class ConnectionManager:
|
| def __init__(self):
|
| self.active_connections: list[WebSocket] = []
|
|
|
| async def connect(self, websocket: WebSocket):
|
| await websocket.accept()
|
| self.active_connections.append(websocket)
|
|
|
| def disconnect(self, websocket: WebSocket):
|
| self.active_connections.remove(websocket)
|
|
|
| manager = ConnectionManager()
|
|
|
| @app.websocket("/ws")
|
| async def websocket_endpoint(websocket: WebSocket):
|
| await manager.connect(websocket)
|
| thread_id = None
|
| task = None
|
| config = None
|
|
|
| async def poll_status():
|
| try:
|
| while True:
|
| if config:
|
| state = workflow_app.get_state(config)
|
| if state.values:
|
| next_node = state.next
|
| values = state.values
|
| current_status = values.get("status", "unknown")
|
| is_waiting_for_human = "human_approval" in next_node
|
|
|
| await websocket.send_json({
|
| "type": "status_update",
|
| "status": current_status,
|
| "revision_count": values.get("revision_count", 0),
|
| "waiting_for_human": is_waiting_for_human,
|
| "latest_feedback": values.get("review_feedback", ""),
|
| "document": values.get("current_draft", "")
|
| })
|
| await asyncio.sleep(1)
|
| except asyncio.CancelledError:
|
| pass
|
| except Exception as e:
|
| print(f"Error in poll_status: {e}")
|
|
|
| try:
|
| while True:
|
| data = await websocket.receive_json()
|
| action = data.get("action")
|
|
|
| if action == "start":
|
| raw_input = data.get("raw_input")
|
| thread_id = str(uuid.uuid4())
|
| config = {"configurable": {"thread_id": thread_id}}
|
|
|
| initial_state = {
|
| "raw_input": raw_input,
|
| "current_draft": "",
|
| "review_feedback": "",
|
| "status": "started",
|
| "revision_count": 0
|
| }
|
|
|
| await websocket.send_json({
|
| "type": "workflow_started",
|
| "thread_id": thread_id
|
| })
|
|
|
| if task:
|
| task.cancel()
|
| task = asyncio.create_task(poll_status())
|
|
|
|
|
| await asyncio.to_thread(workflow_app.invoke, initial_state, config)
|
|
|
| elif action == "approve":
|
| if config:
|
| state = workflow_app.get_state(config)
|
| next_node = state.next
|
| if "human_approval" in next_node:
|
|
|
| await asyncio.to_thread(workflow_app.invoke, None, config)
|
| await websocket.send_json({"type": "workflow_approved"})
|
|
|
| except WebSocketDisconnect:
|
| manager.disconnect(websocket)
|
| if task:
|
| task.cancel()
|
|
|
| if __name__ == "__main__":
|
| import uvicorn
|
| uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)
|
|
|
