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