Spaces:
Sleeping
Sleeping
File size: 4,746 Bytes
57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 3523170 57086e9 | 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 163 | import asyncio
import json
import sys
import io
import re
import threading
import logging
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
app = FastAPI(title="ResearchAgent API")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["POST", "GET", "OPTIONS"],
allow_headers=["*"],
)
# ── Thread-local stdout/stderr/logging capture ───────────────────
_ANSI_RE = re.compile(r"\x1b\[[0-9;]*[mKGHFJA-Za-z]")
_thread_local = threading.local()
_original_stdout = sys.stdout
_original_stderr = sys.stderr
def _clean(text: str) -> str:
return _ANSI_RE.sub("", text).strip()
class _RoutedStream(io.TextIOBase):
def __init__(self, original):
self._original = original
def write(self, text: str) -> int:
q = getattr(_thread_local, "log_queue", None)
if q is not None:
for line in text.splitlines():
clean = _clean(line)
if clean:
try:
q.put_nowait(clean)
except Exception:
pass
else:
try:
self._original.write(text)
except Exception:
pass
return len(text)
def flush(self):
try:
self._original.flush()
except Exception:
pass
def isatty(self):
return False
sys.stdout = _RoutedStream(_original_stdout)
sys.stderr = _RoutedStream(_original_stderr)
class _ThreadLocalLogHandler(logging.Handler):
"""Sends log records to the current thread's queue when one is active."""
def emit(self, record: logging.LogRecord):
q = getattr(_thread_local, "log_queue", None)
if q is not None:
try:
q.put_nowait(self.format(record))
except Exception:
pass
_log_handler = _ThreadLocalLogHandler()
_log_handler.setFormatter(logging.Formatter("%(name)s: %(message)s"))
logging.getLogger().addHandler(_log_handler)
logging.getLogger().setLevel(logging.INFO)
# ── Request model ─────────────────────────────────────────────────
class ResearchRequest(BaseModel):
query: str
depth: str = "standard"
# ── SSE streaming endpoint ────────────────────────────────────────
async def _stream_research(query: str, depth: str):
from agent.crew import run_research
loop = asyncio.get_event_loop()
log_queue: asyncio.Queue = asyncio.Queue(maxsize=500)
thread_done = asyncio.Event()
result_holder: dict = {}
def _worker():
_thread_local.log_queue = log_queue
try:
result = run_research(query, depth)
result_holder.update(result)
except Exception as e:
result_holder["error"] = str(e)
finally:
_thread_local.log_queue = None
loop.call_soon_threadsafe(thread_done.set)
thread = threading.Thread(target=_worker, daemon=True)
thread.start()
# Drain log queue while thread is running
while not thread_done.is_set():
try:
line = log_queue.get_nowait()
yield f"data: {json.dumps(str(line))}\n\n"
except Exception:
await asyncio.sleep(0.05)
# Drain any remaining log lines after thread finishes
while not log_queue.empty():
try:
line = log_queue.get_nowait()
yield f"data: {json.dumps(str(line))}\n\n"
except Exception:
break
if "error" in result_holder:
yield f"data: {json.dumps('[ERROR] ' + result_holder['error'])}\n\n"
else:
report = result_holder.get("report", "")
sources = result_holder.get("sources", [])
yield f"data: {json.dumps({'type': 'report', 'content': report})}\n\n"
yield f"data: {json.dumps({'type': 'sources', 'content': sources})}\n\n"
yield 'data: "__DONE__"\n\n'
@app.post("/research")
async def research(req: ResearchRequest):
if not req.query.strip():
return {"error": "Query cannot be empty"}
return StreamingResponse(
_stream_research(req.query, req.depth),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
)
@app.get("/health")
async def health():
return {"status": "ok"}
if __name__ == "__main__":
import uvicorn
uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=False)
|