main.py

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import json
import structlog
import logging

structlog.configure(
    processors=[
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.stdlib.add_log_level,
        structlog.processors.JSONRenderer()
    ],
    wrapper_class=structlog.make_filtering_bound_logger(logging.INFO),
)

# Import your newly refactored generator
from agent.loop import run_agent_stream
from db.memory import setup_db



app = FastAPI(title="Autonomous Agent API")

@app.on_event("startup")
async def startup_event():
    await setup_db()
    # Pre-load heavy models to prevent blocking/timeouts on first request
    from tools.medical_search import get_model, get_cross_encoder
    import asyncio
    structlog.get_logger().info("startup", message="Pre-loading SentenceTransformer and CrossEncoder...")
    await asyncio.to_thread(get_model)
    await asyncio.to_thread(get_cross_encoder)
    structlog.get_logger().info("startup", message="Models loaded successfully.")

# Repository index will no longer be used for Medical Q&A
    
# Allow the Vanilla JS frontend to hit the API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], # In production, restrict this to your frontend domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class QueryRequest(BaseModel):
    query: str
    session_id: str = "default_session"
    # You can pass session IDs here later to retrieve specific memory buffers

@app.post("/agent/query")
async def agent_query_endpoint(request: QueryRequest):
    """
    Endpoint that consumes a user query and returns a real-time SSE stream 
    of the agent's reasoning and acting process.
    """
    # Pass the query and session_id into the generator
    stream = run_agent_stream(request.query, request.session_id)
    
    # Return the StreamingResponse with the correct media type for SSE
    return StreamingResponse(
        stream, 
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive"
        }
    )

if __name__ == "__main__":
    # Run the high-performance ASGI server
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)