mcp-deep-research/backend/main.py

"""FastAPI service for Deep Research MCP Server."""

import time
import uuid
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from .config import FASTAPI_HOST, FASTAPI_PORT, MAX_WAIT_MINUTES
from .openai_client import get_client

app = FastAPI(
    title="Deep Research API",
    description="OpenAI Deep Research API wrapper",
    version="0.1.0",
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# In-memory task tracking (for status lookups)
_tasks: dict[str, dict[str, Any]] = {}


class ResearchRequest(BaseModel):
    """Request model for starting research."""
    query: str
    system_prompt: str | None = None
    include_code_analysis: bool = True
    max_wait_minutes: int = MAX_WAIT_MINUTES


class ResearchResponse(BaseModel):
    """Response model for research results."""
    task_id: str
    status: str
    model: str | None = None
    report_text: str | None = None
    citations: list[dict[str, Any]] | None = None
    web_searches: int | None = None
    code_executions: int | None = None
    elapsed_time: float | None = None
    error: str | None = None


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy", "service": "deep-research"}


@app.post("/api/research", response_model=ResearchResponse)
async def start_research(request: ResearchRequest) -> ResearchResponse:
    """
    Start a deep research task.

    This initiates a background research task and returns immediately.
    Use GET /api/research/{task_id} to poll for results.
    """
    client = get_client()
    task_id = str(uuid.uuid4())
    start_time = time.time()

    try:
        # Start the research
        result = await client.start_research(
            query=request.query,
            system_prompt=request.system_prompt,
            include_code_analysis=request.include_code_analysis,
        )

        # Store task info
        _tasks[task_id] = {
            "response_id": result["id"],
            "status": result["status"],
            "model": result["model"],
            "start_time": start_time,
            "max_wait_minutes": request.max_wait_minutes,
        }

        return ResearchResponse(
            task_id=task_id,
            status=result["status"],
            model=result["model"],
        )

    except Exception as e:
        return ResearchResponse(
            task_id=task_id,
            status="failed",
            error=str(e),
        )


@app.get("/api/research/{task_id}", response_model=ResearchResponse)
async def get_research_status(task_id: str) -> ResearchResponse:
    """
    Get the status/results of a research task.

    Poll this endpoint until status is 'completed', 'failed', or 'timeout'.
    """
    if task_id not in _tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = _tasks[task_id]
    client = get_client()

    try:
        result = await client.poll_research(task["response_id"])
        elapsed = time.time() - task["start_time"]

        # Check for timeout
        max_seconds = task["max_wait_minutes"] * 60
        if result["status"] not in ("completed", "failed", "cancelled") and elapsed >= max_seconds:
            result["status"] = "timeout"
            result["error"] = f"Timeout after {task['max_wait_minutes']} minutes"

        # Update stored status
        task["status"] = result["status"]

        response = ResearchResponse(
            task_id=task_id,
            status=result["status"],
            model=result.get("model"),
            elapsed_time=elapsed,
            error=result.get("error"),
        )

        # Include output if completed
        if result["status"] == "completed" and "output" in result:
            output = result["output"]
            response.report_text = output.get("report_text")
            response.citations = output.get("citations")
            response.web_searches = output.get("web_searches")
            response.code_executions = output.get("code_executions")

        return response

    except Exception as e:
        return ResearchResponse(
            task_id=task_id,
            status="failed",
            error=str(e),
        )


@app.post("/api/research/{task_id}/wait", response_model=ResearchResponse)
async def wait_for_research(task_id: str) -> ResearchResponse:
    """
    Wait for a research task to complete (blocking).

    This will poll until completion or timeout.
    """
    if task_id not in _tasks:
        raise HTTPException(status_code=404, detail="Task not found")

    task = _tasks[task_id]
    client = get_client()
    start_time = task["start_time"]

    try:
        result = await client.wait_for_completion(
            response_id=task["response_id"],
            max_wait_minutes=task["max_wait_minutes"],
        )

        elapsed = time.time() - start_time
        task["status"] = result["status"]

        response = ResearchResponse(
            task_id=task_id,
            status=result["status"],
            model=result.get("model"),
            elapsed_time=elapsed,
            error=result.get("error"),
        )

        if result["status"] == "completed" and "output" in result:
            output = result["output"]
            response.report_text = output.get("report_text")
            response.citations = output.get("citations")
            response.web_searches = output.get("web_searches")
            response.code_executions = output.get("code_executions")

        return response

    except Exception as e:
        return ResearchResponse(
            task_id=task_id,
            status="failed",
            error=str(e),
        )


def main():
    """Run the FastAPI server."""
    import uvicorn
    print(f"Starting Deep Research API on {FASTAPI_HOST}:{FASTAPI_PORT}")
    uvicorn.run(app, host=FASTAPI_HOST, port=FASTAPI_PORT)


if __name__ == "__main__":
    main()