from typing import Annotated, List, Optional from fastapi import APIRouter, Depends from fastapi.responses import JSONResponse from service.dto import UserPromptRequest, BotResponse, BotCreateRequest from core.chat.chatstore import ChatStore from db.database import get_db from db.models import Bot_Meta, Bot, Metadata from db.models import Session as SessionModel from sqlalchemy.orm import Session from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.exc import NoResultFound from sqlalchemy import select 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 langfuse.llama_index import LlamaIndexCallbackHandler from datetime import datetime router = APIRouter(tags=["Bot_Specific"]) db_dependency = Annotated[Session, Depends(get_db)] def get_chat_store(): return ChatStore() @router.post("/bot") async def create_bot_id( user: user_dependency, db: db_dependency, bot_request: BotCreateRequest, ): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") # Create a new bot entry try: # Create a new bot entry new_bot = Bot( user_id=user.get("id"), bot_name=bot_request.bot_name ) # Assuming user has an 'id' attribute db.add(new_bot) db.commit() # Commit the transaction db.refresh(new_bot) # Optional: Refresh the instance with the database state return {"status": "success", "bot_id": new_bot.id} except SQLAlchemyError as e: db.rollback() # Roll back the transaction in case of an error return JSONResponse(status_code=500, content=f"Database error: {str(e)}") except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected error occurred: {str(e)}" ) @router.post("/meta/{bot_id}") async def create_bot_specific( user: user_dependency, db: db_dependency, bot_id: int, metadata_id: List[Optional[int]], # Use the Pydantic model ): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") try: # Create BotMeta instances for each metadata_id bot_meta_entries = [ Bot_Meta(bot_id=bot_id, metadata_id=mid) for mid in metadata_id ] # Insert all entries into the database db.add_all(bot_meta_entries) db.commit() # Commit the transaction except SQLAlchemyError as e: return JSONResponse(status_code=500, content=f"Database error: {str(e)}") except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected error occurred: {str(e)}" ) return {"status": "success", "bot_meta": [entry.id for entry in bot_meta_entries]} @router.put("/meta/{bot_id}") async def update_bot_specific( user: user_dependency, db: db_dependency, bot_id: int, metadata_id: List[Optional[int]], # Use the Pydantic model ): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") try: # Fetch existing Bot_Meta entries related to bot_id existing_entries = db.query(Bot_Meta).filter(Bot_Meta.bot_id == bot_id).all() # Delete existing entries for entry in existing_entries: db.delete(entry) # Commit the deletion db.commit() # Insert the new metadata entries bot_meta_entries = [ Bot_Meta(bot_id=bot_id, metadata_id=mid) for mid in metadata_id ] db.add_all(bot_meta_entries) db.commit() except SQLAlchemyError as e: db.rollback() # Rollback in case of any database error return JSONResponse(status_code=500, content=f"Database error: {str(e)}") except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected error occurred: {str(e)}" ) return {"status": "success", "bot_meta": [entry.id for entry in bot_meta_entries]} @router.delete("/meta/{bot_id}/{metadata_id}") async def delete_bot_specific( user: user_dependency, db: db_dependency, bot_id: int, metadata_id: int, # Changed to int to specify a single metadata_id ): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") try: # Delete the specific metadata entry for the given bot_id bot_meta_entry = ( db.query(Bot_Meta) .filter(Bot_Meta.bot_id == bot_id, Bot_Meta.metadata_id == metadata_id) .first() # Use first() to get a single entry ) if not bot_meta_entry: return JSONResponse(status_code=404, content="No entry found to delete.") # Delete the found entry db.delete(bot_meta_entry) db.commit() except SQLAlchemyError as e: db.rollback() # Rollback in case of any database error return JSONResponse(status_code=500, content=f"Database error: {str(e)}") except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected error occurred: {str(e)}" ) return { "status": "success", "deleted_entry_id": bot_meta_entry.id, } @router.delete("/bot_all/{bot_id}") async def delete_bot_id( user: user_dependency, db: db_dependency, bot_id: int, ): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") try: # Fetch the bot to ensure it exists bot_entry = db.query(Bot).filter(Bot.id == bot_id).first() print("bot entry", bot_entry) if not bot_entry: return JSONResponse( status_code=404, content=f"Bot with id {bot_id} not found." ) db.query(SessionModel).filter(SessionModel.bot_id == bot_id).delete( synchronize_session="fetch" ) db.query(Bot_Meta).filter(Bot_Meta.bot_id == bot_id).delete( synchronize_session="fetch" ) db.delete(bot_entry) db.commit() # Commit all deletions return { "status": "success", "deleted_bot_id": bot_id, } except SQLAlchemyError as e: db.rollback() # Rollback in case of any database error return JSONResponse(status_code=500, content=f"Database error: {str(e)}") except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected error occurred: {str(e)}" ) @router.post("/session/{bot_id}/new") async def create_new_session(user: user_dependency, db: db_dependency, bot_id: int): # Check if user is authenticated if user is None: return JSONResponse(status_code=401, content="Authentication Failed") print(user.get("id")) user_id = user.get("id") # Ensure the bot belongs to the user bot_query = select(Bot).where(Bot.id == bot_id, Bot.user_id == user_id) try: bot = db.execute(bot_query).scalar_one() print(bot) except NoResultFound: return JSONResponse( status_code=404, content="Bot not found or unauthorized access." ) # Generate a new session ID (UUID) try: session_id = generate_uuid() # Create the new session new_session = SessionModel( id=session_id, user_id=user.get("id"), bot_id=bot_id, ) db.add(new_session) db.commit() # Commit the new session to the database return { "session_id": session_id, } except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected in retrieving session id {str(e)}" ) @router.get("/bot/all/{bot_id}") async def get_all_session_ids(user: user_dependency, db: db_dependency, bot_id: int): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") try: query = select(SessionModel.id, SessionModel.updated_at).where( SessionModel.user_id == user.get("id"), SessionModel.bot_id == bot_id ) result = db.execute(query) sessions = result.all() session_data = [{"id": session.id, "updated_at": session.updated_at} for session in sessions] return session_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") @router.post("/bot/{bot_id}/{session_id}") async def bot_generator_spesific( user: user_dependency, db: db_dependency, bot_id: int, session_id: str, user_prompt_request: UserPromptRequest, ): if user is None: return JSONResponse(status_code=401, content="Authentication Failed") 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: query = ( select(Metadata.title) .join(Bot_Meta, Metadata.id == Bot_Meta.metadata_id) .join(SessionModel, Bot_Meta.bot_id == bot_id) .where( SessionModel.user_id == user.get("id"), SessionModel.id == session_id ) ) result = db.execute(query) titles = result.scalars().all() print(titles) except SQLAlchemyError as e: return JSONResponse(status_code=500, content=f"Database error: {str(e)}") except Exception as e: return JSONResponse( status_code=500, content=f"An unexpected error occurred: {str(e)}" ) 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, titles, type_bot="specific" ) existing_session = ( db.query(SessionModel).filter(SessionModel.id == session_id).first() ) existing_session.updated_at = datetime.now() db.commit() return BotResponse( content=response, metadata=metadata, scores=scores, )