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")