import asyncio
from fastapi import APIRouter
from pydantic import BaseModel

from app.rag import HotelAgent

from sse_starlette.sse import EventSourceResponse
from langchain_core.messages import HumanMessage

router = APIRouter()

class History(BaseModel):
    role: str
    text: str

class request(BaseModel):
    history: list[History]
    message: str

class OutputSchema(BaseModel):
    answer: str
    confidence: float

@router.post("/ask")
async def get_price(args: request):

    async def event_generator():
        try:

            agent = HotelAgent()

            inputs = {
                "chat_history": [],
                "messages": [HumanMessage(content=args.message)]
            }

            async for chunk in agent.astream(inputs):
                yield chunk
        
        except asyncio.CancelledError:
            raise
        finally:
            pass

    return EventSourceResponse(event_generator())