app.py

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"<h1>Error loading index.html</h1><pre>{exc}</pre>"