-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathagent.py
More file actions
76 lines (59 loc) · 2.23 KB
/
Copy pathagent.py
File metadata and controls
76 lines (59 loc) · 2.23 KB
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
import asyncio
import json
import os
from typing import Optional
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
from agent_loop import react_agent
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, ConfigDict
app = FastAPI()
class QueryRequest(BaseModel):
model_config = ConfigDict(
extra="allow",
json_schema_extra={
"example": {"question": "Where is the capital of France?"}
},
)
question: str
chat_history: Optional[list] = None
class QueryResponse(BaseModel):
answer: str
@app.post("/")
async def query(req: QueryRequest, raw_request: Request):
"""
Research Agent API endpoint.
Supports both JSON and SSE responses based on Accept header.
When SSE is requested, sends periodic heartbeat comments to keep
the connection alive while the agent is processing.
"""
accept = raw_request.headers.get("accept", "")
if "text/event-stream" in accept:
# SSE mode: stream heartbeats while agent processes, then send answer
async def sse_response():
try:
# Run agent in background task
agent_task = asyncio.create_task(react_agent(req.question))
# Send heartbeat comments every 5s to keep connection alive
while not agent_task.done():
yield ": heartbeat\n\n"
try:
await asyncio.wait_for(asyncio.shield(agent_task), timeout=5.0)
except asyncio.TimeoutError:
continue
answer = await agent_task
except Exception as e:
print(f"[agent] SSE error: {e}")
answer = f"Error: {e}"
data = json.dumps({"answer": answer}, ensure_ascii=False)
yield f"event: Message\ndata: {data}\n\n"
return StreamingResponse(sse_response(), media_type="text/event-stream")
# Normal JSON mode
try:
answer = await react_agent(req.question)
except Exception as e:
print(f"[agent] Error: {e}")
answer = f"Error: {e}"
return JSONResponse(content={"answer": answer})