app/main.py

from __future__ import annotations

import os

from fastapi import FastAPI
from pydantic import BaseModel

from app.agent.health_agent import HealthAgent
from app.agent.kb_embedding import KBEmbeddingService
from app.agent.kb_retrieval import validate_similarity
from app.db.chroma_client import get_vector_backend, safe_count
from app.nlp.nlp_service import NLPService


app = FastAPI(title="AliveAI Medical RAG Chatbot")

SESSION_AGENTS: dict[str, HealthAgent] = {}
DEFAULT_MODEL = os.getenv("ALIVEAI_HEALTH_MODEL", "llama3.2:3b")


class ChatRequest(BaseModel):
    message: str
    session_id: str = "default"


def get_agent(session_id: str) -> HealthAgent:
    if session_id not in SESSION_AGENTS:
        SESSION_AGENTS[session_id] = HealthAgent(model=DEFAULT_MODEL)
    return SESSION_AGENTS[session_id]


@app.on_event("startup")
async def startup_event() -> None:
    NLPService()
    KBEmbeddingService().load_model()
    print("All models loaded. Ready.")


@app.post("/chat")
async def chat(request: ChatRequest) -> dict:
    agent = get_agent(request.session_id)
    return agent.chat(request.message)


@app.get("/health")
async def health() -> dict:
    vector_count = safe_count()
    return {
        "status": "ok",
        "backend": get_vector_backend(),
        "vector_count": vector_count,
        "chroma_count": vector_count,
    }


@app.post("/reset/{session_id}")
async def reset(session_id: str) -> dict:
    agent = get_agent(session_id)
    agent.reset()
    return {"status": "reset", "session_id": session_id}


@app.get("/validate")
async def validate(text1: str, text2: str) -> dict:
    score = validate_similarity(text1, text2)
    return {"text1": text1, "text2": text2, "similarity": score}