server.py

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)