from __future__ import annotations import asyncio import json import os import re from collections import defaultdict from time import time from typing import Literal from fastapi import FastAPI, HTTPException, Request from fastapi.responses import HTMLResponse, StreamingResponse from pydantic import BaseModel, Field from research_workflow import run_chat_completion, run_research_pipeline app = FastAPI(title="Deep Research") memory_store: dict[str, list[dict[str, str]]] = {} usage_store: defaultdict[str, list[float]] = defaultdict(list) MAX_RESEARCH_REQUESTS = int(os.getenv("MAX_RESEARCH_REQUESTS", "10")) WINDOW_SECONDS = 86400 MAX_MEMORY_ITEMS = 12 RESEARCH_PATTERNS = [ r"\bdeep\s+research\b", r"\bresearch\b", r"\binvestigate\b", r"\bliterature\s+review\b", r"\bmarket\s+analysis\b", r"\bcompetitive\s+analysis\b", r"\bdue\s+diligence\b", r"\bwhite\s+paper\b", r"\bfull\s+report\b", r"\bdetailed\s+report\b", r"\bsources?\b", r"\bcitations?\b", r"\blatest\b", r"\bcurrent\b", r"\bup[-\s]?to[-\s]?date\b", r"\b202[4-9]\b", ] CASUAL_PATTERNS = [ r"^(hi|hello|hey|yo|thanks|thank you)\b", r"\bwhat can you do\b", r"\bwho are you\b", r"\bhelp me\b", ] class MessageRequest(BaseModel): topic: str = Field(..., min_length=1, max_length=4000) session_id: str = "default" mode: Literal["auto", "chat", "research"] = "auto" def get_memory(session_id: str) -> list[dict[str, str]]: memory = memory_store.setdefault(session_id, []) if len(memory) > MAX_MEMORY_ITEMS: del memory[:-MAX_MEMORY_ITEMS] return memory def check_research_limit(ip: str) -> None: now = time() usage_store[ip] = [t for t in usage_store[ip] if now - t < WINDOW_SECONDS] if len(usage_store[ip]) >= MAX_RESEARCH_REQUESTS: raise HTTPException( status_code=429, detail="Daily research limit reached. Please try again later.", ) usage_store[ip].append(now) def classify_intent(text: str, mode: str = "auto") -> Literal["chat", "research"]: if mode in {"chat", "research"}: return mode normalized = re.sub(r"\s+", " ", text.lower()).strip() if not normalized: return "chat" if any(re.search(pattern, normalized) for pattern in CASUAL_PATTERNS): return "chat" if any(re.search(pattern, normalized) for pattern in RESEARCH_PATTERNS): return "research" word_count = len(normalized.split()) asks_for_depth = any( phrase in normalized for phrase in [ "pros and cons", "trade offs", "tradeoffs", "compare and contrast", "analysis of", "explain the evidence", ] ) return "research" if word_count >= 18 and asks_for_depth else "chat" async def sse_event(payload: dict) -> str: return f"data: {json.dumps(payload)}\n\n" async def chat_stream_generator(message: str, session_id: str): memory = get_memory(session_id) yield await sse_event({"status": "progress", "message": "Answering in chat mode..."}) answer = await run_chat_completion(message, memory) memory.append({"role": "user", "content": message}) memory.append({"role": "assistant", "content": answer}) yield await sse_event( { "status": "complete", "mode": "chat", "message": "Chat response", "report": answer, } ) async def research_stream_generator(topic: str, session_id: str): memory = get_memory(session_id) progress_messages = [ "Breaking the request into research subtopics...", "Searching multiple sources...", "Extracting and cross-checking findings...", "Synthesizing the final report...", ] for message in progress_messages: yield await sse_event({"status": "progress", "message": message}) await asyncio.sleep(0.15) final_report = await run_research_pipeline(topic, memory) yield await sse_event( { "status": "complete", "mode": "research", "message": "Research complete", "report": final_report, } ) @app.post("/research") async def message_endpoint(req: Request, request: MessageRequest): intent = classify_intent(request.topic, request.mode) if intent == "chat": return StreamingResponse( chat_stream_generator(request.topic, request.session_id), media_type="text/event-stream", ) check_research_limit(req.client.host if req.client else "unknown") return StreamingResponse( research_stream_generator(request.topic, request.session_id), media_type="text/event-stream", ) @app.get("/health") async def health(): return {"status": "ok"} @app.get("/", response_class=HTMLResponse) async def read_root(): try: with open("index.html", "r", encoding="utf-8") as f: return f.read() except OSError as exc: return f"
{exc}"