co.py

import sys
import os
from fastapi.middleware.cors import CORSMiddleware

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from rag import rag_answer, rag_answer_stream
from typing import Optional

app = FastAPI()

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

class Question(BaseModel):
    question: str
    conversation_id: Optional[int] = None

@app.post("/ask")
async def ask_question(data: Question):
    """
    Handle question, get RAG answer, and return it.
    """
    try:
        # 1. Get answer from RAG
        answer = rag_answer(data.question)
        
        return {
            "answer": answer,
            "conversation_id": data.conversation_id
        }

    except Exception as e:
        print(f"Error in ask_question: {e}")
        return {
            "answer": "Error processing your question.",
            "conversation_id": data.conversation_id,
        }

@app.post("/ask_stream")
async def ask_question_stream(data: Question):
    """Stream the answer from RAG."""
    return StreamingResponse(rag_answer_stream(data.question), media_type="text/plain")