api/router/bot_general.py

from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from service.dto import UserPromptRequest, BotResponse
from core.chat.chatstore import ChatStore
from db.database import get_db
from db.models import Session_Publisher
from sqlalchemy.orm import Session
from api.function import (
    generate_streaming_completion,
    generate_completion_non_streaming,
)
from api.router.user import user_dependency
from sse_starlette.sse import EventSourceResponse
from utils.utils import generate_uuid
from typing import Annotated
from langfuse.llama_index import LlamaIndexCallbackHandler


router = APIRouter(tags=["Bot_General"])

db_dependency = Annotated[Session, Depends(get_db)]


def get_chat_store():
    return ChatStore()


@router.post("/bot_general/new")
async def create_session_general():
    session_id = generate_uuid()
    return {"session_id": session_id}


@router.get("/bot/{session_id}")
async def get_session_id(
    user: user_dependency,
    session_id: str,
    chat_store: ChatStore = Depends(get_chat_store),
):
    if user is None:
        return JSONResponse(status_code=401, content="Authentication Failed")

    chat_history = chat_store.get_all_messages_mongodb(session_id)
    # if not chat_history:
    #     return JSONResponse(status_code=404, content="Session not found or empty.")

    return chat_history


@router.post("/bot/{session_id}")
async def bot_generator_general(
    session_id: str, user_prompt_request: UserPromptRequest
):

    langfuse_callback_handler = LlamaIndexCallbackHandler()
    langfuse_callback_handler.set_trace_params(user_id="guest", session_id=session_id)

    if user_prompt_request.streaming:
        return EventSourceResponse(
            generate_streaming_completion(
                user_prompt_request.prompt,
                session_id
            )
        )
    else:
        response, metadata, scores = generate_completion_non_streaming(
            session_id,
            user_prompt_request.prompt,
        )

        return BotResponse(
            content=response,
            metadata=metadata,
            scores=scores,
        )


@router.delete("/bot/{session_id}")
async def delete_bot(db:db_dependency, session_id: str, chat_store: ChatStore = Depends(get_chat_store)):
    try:
        chat_store.delete_messages(session_id)
        
        # Delete session from database
        session = db.query(Session_Publisher).filter(Session_Publisher.id == session_id).first()
        if session:
            db.delete(session)
            db.commit()  # Commit the transaction
        else:
            return JSONResponse(status_code=404, content="Session not found")
        return {"info": f"Delete {session_id} successful"}
    except Exception as e:
        # Log the error and return JSONResponse for FastAPI
        print(f"An error occurred in update data.: {e}")
        return JSONResponse(status_code=400, content="the error when deleting message")