api/router/bot_one.py

from typing import Annotated
from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
import pytz

from service.dto import UserPromptRequest, BotResponse, BotCreateRequest
from datetime import datetime
from core.chat.chatstore import ChatStore
# from core.chat.bot_service import ChatCompletionService
from core.chat.bot_service_multimodal import ChatCompletionService
from db.database import get_db
from db.models import Session_Publisher
from db.query.query_book import BookQuery
from db.query.query_bot import BotQuery

from sqlalchemy.orm import Session
from sse_starlette.sse import EventSourceResponse
from utils.utils import generate_uuid
from utils.error_handlers import handle_exception
from langfuse.llama_index import LlamaIndexCallbackHandler

from api.auth import check_user_authentication
from api.router.user import user_dependency
from api.function import generate_streaming_completion


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

db_dependency = Annotated[Session, Depends(get_db)]


def get_chat_store():
    return ChatStore()


@router.post("/bot_one/{metadata_id}")
async def create_bot_one(
    user: user_dependency,
    db: db_dependency,
    metadata_id: int,
    bot_name: BotCreateRequest,
):
    auth_response = check_user_authentication(user)
    if auth_response:
        return auth_response
    # Generate a new session ID (UUID)
    try:
        session_id = generate_uuid()

        bot_query = BotQuery(user)
        bot_query.add_bot(db, session_id, bot_name.name, metadata_id)

        return {
            "status": "session id created successfully",
            "bot_name": bot_name.name,
            "session_id": session_id,
        }

    except Exception as e:
        return JSONResponse(
            status_code=500, content=f"An unexpected  in retrieving session id {str(e)}"
        )


@router.post("/bot_one/{metadata_id}/{session_id}")
async def generator_bot(
    user: user_dependency,
    db: db_dependency,
    metadata_id: int,
    session_id: str,
    user_prompt_request: UserPromptRequest,
):
    auth_response = check_user_authentication(user)
    if auth_response:
        return auth_response

    langfuse_callback_handler = LlamaIndexCallbackHandler()
    langfuse_callback_handler.set_trace_params(
        user_id=user.get("username"), session_id=session_id
    )

    # Query to retrieve the titles
    try:
        book_query = BookQuery(user)
        output_titles = book_query.get_title_from_session(db, metadata_id, session_id)
        titles = [item[0] for item in output_titles]
        print(titles)

    except Exception as e:
        return handle_exception(e)

    if user_prompt_request.streaming:
        return EventSourceResponse(
            generate_streaming_completion(
                user_prompt_request.prompt,
                session_id,
            )
        )
    else:
        bot_service = ChatCompletionService(
            session_id, user_prompt_request.prompt, titles, type_bot="specific"
        )
        response, metadata, scores = bot_service.generate_completion()

        # Set Jakarta timezone
        jakarta_tz = pytz.timezone("Asia/Jakarta")

        existing_session = (
            db.query(Session_Publisher)
            .filter(Session_Publisher.id == session_id)
            .first()
        )

        existing_session.updated_at = datetime.now(jakarta_tz)
        db.commit()

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


@router.get("/bot_one/{metadata_id}")
async def get_all_session_bot_one(
    user: user_dependency, db: db_dependency, metadata_id: int
):
    auth_response = check_user_authentication(user)
    if auth_response:
        return auth_response

    try:
        # Query the session IDs based on the user ID
        bot_query = BotQuery(user)
        chat_store = ChatStore()
        sessions = bot_query.get_session_ids_bot(db, metadata_id)

        session_data = [
            {
                "id": session.id,
                "bot_name": session.bot_name,
                "updated_at": str(session.updated_at),
                "last_message": chat_store.get_last_message_mongodb(session.id)
            }
            for session in sessions
        ]

        # Convert list of tuples to a simple list
        session_sorted_data = sorted(
            session_data,
            key=lambda x: datetime.fromisoformat(x["updated_at"]),
            reverse=True,
        )

        return session_sorted_data

    except Exception as e:
        # Log the error and return JSONResponse for FastAPI
        print(f"An error occurred while fetching session IDs: {e}")
        return JSONResponse(status_code=400, content="Error retrieving session IDs")