server.py

"""
Multi-Step AI Agent — FastAPI REST API
=======================================
Run with:  uvicorn server:app --host 0.0.0.0 --port 7860
"""

from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import os
import tempfile
import shutil

from app.agent.agent import create_agent, get_memory_manager


# ── Lifespan — warm up agent on startup ─────────────────────────

_agent = None
_memory = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize the agent once at startup."""
    global _agent, _memory
    _agent = create_agent()
    _memory = get_memory_manager()
    yield


# ── FastAPI app ─────────────────────────────────────────────────

app = FastAPI(
    title="Multi-Step AI Agent",
    description="Plan-and-Execute AI Agent with RAG Knowledge Base",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Request / response models ───────────────────────────────────

class ChatRequest(BaseModel):
    message: str


class ChatResponse(BaseModel):
    reply: str
    steps: list[dict] = []


class DocumentResponse(BaseModel):
    status: str
    detail: str


# ── Endpoints ───────────────────────────────────────────────────

@app.get("/health")
async def health():
    """Health check for deployment platforms."""
    return {"status": "ok", "agent": _agent is not None}


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """Send a message to the agent and get a response."""
    if not _agent:
        raise HTTPException(status_code=503, detail="Agent not initialized")

    steps_log = []

    def on_step(step_num, tool_name, status):
        steps_log.append({
            "step": step_num,
            "tool": tool_name,
            "status": status,
        })

    try:
        reply = _agent.chat(request.message, on_step=on_step)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    return ChatResponse(reply=reply, steps=steps_log)


@app.post("/upload", response_model=DocumentResponse)
async def upload_document(file: UploadFile = File(...)):
    """Upload a PDF to the knowledge base."""
    if not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are supported")

    # Save to a temp location, then ingest
    upload_dir = os.path.join("data", "uploads")
    os.makedirs(upload_dir, exist_ok=True)
    file_path = os.path.join(upload_dir, file.filename)

    with open(file_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    # Use the store_document tool directly
    from app.tools.knowledge_base import store_document
    result = store_document.invoke({"file_path": file_path})

    return DocumentResponse(status="ok", detail=result)


@app.get("/documents")
async def list_docs():
    """List documents in the knowledge base."""
    from app.tools.knowledge_base import list_documents
    result = list_documents.invoke({})
    return {"documents": result}


@app.post("/clear")
async def clear_memory():
    """Clear the agent's chat memory."""
    if _memory:
        _memory.clear()
    return {"status": "ok", "detail": "Memory cleared"}


# ── Run directly ────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)