main.py

import json
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from functools import lru_cache
from clients.groq_client import GroqClient
from fastapi.responses import StreamingResponse

app = FastAPI()


# Allowed origins
origins = [
    "http://localhost:4200",
    "https://mosesmangabo.vercel.app",
    "https://chatbot-bridge-service.onrender.com",
    "https://eng-musa-chatbotbridge.hf.space",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],  # allow all HTTP methods (GET, POST, ...)
    allow_headers=["*"],  # allow all headers
)


class QueryRequest(BaseModel):
    query: str
    ip: str


# Lazy singleton for GroqClient
@lru_cache(maxsize=1)
def get_groq_client():
    return GroqClient()


# Non - Streaming
# ============================================
@app.post("/chat")
async def chat_endpoint(request: QueryRequest):
    try:
        groq_client = get_groq_client()
        response, model = groq_client.ask(request.query)
        if response.startswith("Error while"):
            return {"status_code": 500, "response": response, "model": model, "next_questions": []}
        return {
            "status_code": 200,
            "response": response,
            "model": model,
            "next_questions": groq_client.get_next_questions(request.ip),
        }
    except Exception as e:
        return {
            "status_code": 500,
            "response": f"Internal error: {e}",
            "model": None,
            "next_questions": [],
        }


# Streaming
# ============================================
@app.post("/chat-stream")
async def chat_stream_endpoint(request: QueryRequest):
    """Streaming endpoint that returns chunks in real-time with a consistent structure"""
    try:
        groq_client = get_groq_client()

        def generate():
            try:
                response = ""
                for chunk in groq_client.ask_stream(request.query):
                    response += chunk

                    yield f"data: {json.dumps({'status_code': 200, 'response': chunk, 'next_questions': []})}\n\n"

                if response.startswith("Error while"):
                    yield f"data: {json.dumps({'status_code': 500, 'response': response, 'next_questions': []})}\n\n"
                else:

                    yield f"data: {json.dumps({'status_code': 200, 'response': response, 'next_questions': []})}\n\n"

                    yield f"data: {json.dumps({'status_code': 200, 'response': '[DONE]', 'next_questions': groq_client.get_next_questions(request.ip)})}\n\n"

            except Exception as e:
                yield f"data: {json.dumps({'status_code': 500, 'response': f'Error: {str(e)}', 'next_questions': []})}\n\n"

        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",  # Disable nginx buffering
            },
        )

    except Exception as e:
        return {
            "status_code": 500,
            "response": f"Internal error: {e}",
            "next_questions": [],
        }


@app.get("/")
def healthcheck():
    return {"status": "ok"}


# python -m uvicorn main:app --reload