Team Ai
Apppublic

Multimedika/Bot_Development

sourceHugging Facemitupdated 2y agoView on Hugging Face
0likes
bot_specific.py340 linesDownload Raw Back to router
1from typing import Annotated, List, Optional2from datetime import datetime3 4from fastapi import APIRouter, Depends5from fastapi.responses import JSONResponse6from sse_starlette.sse import EventSourceResponse7from sqlalchemy import select8from sqlalchemy.orm import Session9from sqlalchemy.exc import SQLAlchemyError, NoResultFound10 11from service.dto import UserPromptRequest, BotResponse, BotCreateRequest12from core.chat.chatstore import ChatStore13# from core.chat.bot_service import ChatCompletionService14from core.chat.bot_service_multimodal import ChatCompletionService15from db.database import get_db16from db.models import Bot_Meta, Bot, Metadata17from db.models import Session as SessionModel18from langfuse.llama_index import LlamaIndexCallbackHandler19from api.auth import check_user_authentication20from api.router.user import user_dependency21from api.function import generate_streaming_completion22from utils.utils import generate_uuid23from utils.error_handlers import handle_exception24 25 26router = APIRouter(tags=["Bot_Specific"])27 28db_dependency = Annotated[Session, Depends(get_db)]29 30 31def get_chat_store():32    return ChatStore()33 34 35@router.post("/bot")36async def create_bot_id(37    user: user_dependency,38    db: db_dependency,39    bot_request: BotCreateRequest,40):41 42    auth_response = check_user_authentication(user)43    if auth_response:44        return auth_response45 46    # Create a new bot entry47    try:48        # Create a new bot entry49        new_bot = Bot(50            user_id=user.get("id"), bot_name=bot_request.name51        )  # Assuming user has an 'id' attribute52 53        db.add(new_bot)54        db.commit()  # Commit the transaction55        db.refresh(new_bot)  # Optional: Refresh the instance with the database state56 57        return {"status": "success", "bot_id": new_bot.id}58 59    except SQLAlchemyError as e:60        db.rollback()  # Roll back the transaction in case of an error61        return JSONResponse(status_code=500, content=f"Database error: {str(e)}")62    except Exception as e:63        return JSONResponse(64            status_code=500, content=f"An unexpected error occurred: {str(e)}"65        )66 67 68@router.post("/meta/{bot_id}")69async def create_bot_specific(70    user: user_dependency,71    db: db_dependency,72    bot_id: int,73    metadata_id: List[Optional[int]],  # Use the Pydantic model74):75    auth_response = check_user_authentication(user)76    if auth_response:77        return auth_response78 79    try:80        # Create BotMeta instances for each metadata_id81        bot_meta_entries = [82            Bot_Meta(bot_id=bot_id, metadata_id=mid) for mid in metadata_id83        ]84 85        # Insert all entries into the database86        db.add_all(bot_meta_entries)87        db.commit()  # Commit the transaction88 89    except SQLAlchemyError as e:90        return JSONResponse(status_code=500, content=f"Database error: {str(e)}")91    except Exception as e:92        return JSONResponse(93            status_code=500, content=f"An unexpected error occurred: {str(e)}"94        )95 96    return {"status": "success", "bot_meta": [entry.id for entry in bot_meta_entries]}97 98 99@router.put("/meta/{bot_id}")100async def update_bot_specific(101    user: user_dependency,102    db: db_dependency,103    bot_id: int,104    metadata_id: List[Optional[int]],  # Use the Pydantic model105):106    auth_response = check_user_authentication(user)107    if auth_response:108        return auth_response109 110    try:111        # Fetch existing Bot_Meta entries related to bot_id112        existing_entries = db.query(Bot_Meta).filter(Bot_Meta.bot_id == bot_id).all()113 114        # Delete existing entries115        for entry in existing_entries:116            db.delete(entry)117 118        # Commit the deletion119        db.commit()120 121        # Insert the new metadata entries122        bot_meta_entries = [123            Bot_Meta(bot_id=bot_id, metadata_id=mid) for mid in metadata_id124        ]125        db.add_all(bot_meta_entries)126        db.commit()127        return {128            "status": "success",129            "bot_meta": [entry.id for entry in bot_meta_entries],130        }131 132    except SQLAlchemyError as e:133        db.rollback()  # Rollback in case of any database error134        return JSONResponse(status_code=500, content=f"Database error: {str(e)}")135    except Exception as e:136        return handle_exception(e)137 138 139@router.delete("/meta/{bot_id}/{metadata_id}")140async def delete_bot_specific(141    user: user_dependency,142    db: db_dependency,143    bot_id: int,144    metadata_id: int,  # Changed to int to specify a single metadata_id145):146    auth_response = check_user_authentication(user)147    if auth_response:148        return auth_response149 150    try:151        # Delete the specific metadata entry for the given bot_id152        bot_meta_entry = (153            db.query(Bot_Meta)154            .filter(Bot_Meta.bot_id == bot_id, Bot_Meta.metadata_id == metadata_id)155            .first()  # Use first() to get a single entry156        )157 158        if not bot_meta_entry:159            return JSONResponse(status_code=404, content="No entry found to delete.")160 161        # Delete the found entry162        db.delete(bot_meta_entry)163        db.commit()164        return {165            "status": "success",166            "deleted_entry_id": bot_meta_entry.id,167        }168 169    except SQLAlchemyError as e:170        db.rollback()  # Rollback in case of any database error171        return JSONResponse(status_code=500, content=f"Database error: {str(e)}")172    except Exception as e:173        return handle_exception(e)174 175 176@router.delete("/bot_all/{bot_id}")177async def delete_bot_id(178    user: user_dependency,179    db: db_dependency,180    bot_id: int,181):182    auth_response = check_user_authentication(user)183    if auth_response:184        return auth_response185 186    try:187        # Fetch the bot to ensure it exists188        bot_entry = db.query(Bot).filter(Bot.id == bot_id).first()189        print("bot entry", bot_entry)190 191        if not bot_entry:192            return JSONResponse(193                status_code=404, content=f"Bot with id {bot_id} not found."194            )195 196        db.query(SessionModel).filter(SessionModel.bot_id == bot_id).delete(197            synchronize_session="fetch"198        )199        db.query(Bot_Meta).filter(Bot_Meta.bot_id == bot_id).delete(200            synchronize_session="fetch"201        )202        db.delete(bot_entry)203        db.commit()  # Commit all deletions204 205        return {206            "status": "success",207            "deleted_bot_id": bot_id,208        }209 210    except SQLAlchemyError as e:211        db.rollback()  # Rollback in case of any database error212        return JSONResponse(status_code=500, content=f"Database error: {str(e)}")213    except Exception as e:214        return handle_exception(e)215 216 217@router.post("/session/{bot_id}/new")218async def create_new_session(user: user_dependency, db: db_dependency, bot_id: int):219    # Check if user is authenticated220    auth_response = check_user_authentication(user)221    if auth_response:222        return auth_response223 224    user_id = user.get("id")225    # Ensure the bot belongs to the user226    bot_query = select(Bot).where(Bot.id == bot_id, Bot.user_id == user_id)227 228    try:229        bot = db.execute(bot_query).scalar_one()230        print(bot)231 232    except NoResultFound:233        return JSONResponse(234            status_code=404, content="Bot not found or unauthorized access."235        )236 237    # Generate a new session ID (UUID)238    try:239        session_id = generate_uuid()240 241        # Create the new session242        new_session = SessionModel(243            id=session_id,244            user_id=user.get("id"),245            bot_id=bot_id,246        )247 248        db.add(new_session)249        db.commit()  # Commit the new session to the database250 251        return {252            "session_id": session_id,253        }254 255    except Exception as e:256        return handle_exception(e)257 258 259@router.get("/bot/all/{bot_id}")260async def get_all_session_ids(user: user_dependency, db: db_dependency, bot_id: int):261    auth_response = check_user_authentication(user)262    if auth_response:263        return auth_response264 265    try:266        query = select(SessionModel.id, SessionModel.updated_at).where(267            SessionModel.user_id == user.get("id"), SessionModel.bot_id == bot_id268        )269 270        result = db.execute(query)271        sessions = result.all()272        session_data = [273            {"id": session.id, "updated_at": session.updated_at} for session in sessions274        ]275 276        return session_data277 278    except Exception as e:279        return handle_exception(e)280 281 282@router.post("/bot/{bot_id}/{session_id}")283async def bot_generator_spesific(284    user: user_dependency,285    db: db_dependency,286    bot_id: int,287    session_id: str,288    user_prompt_request: UserPromptRequest,289):290    auth_response = check_user_authentication(user)291    if auth_response:292        return auth_response293    langfuse_callback_handler = LlamaIndexCallbackHandler()294    langfuse_callback_handler.set_trace_params(295        user_id=user.get("username"), session_id=session_id296    )297 298    # Query to retrieve the titles299    try:300        query = (301            select(Metadata.title)302            .join(Bot_Meta, Metadata.id == Bot_Meta.metadata_id)303            .join(SessionModel, Bot_Meta.bot_id == bot_id)304            .where(305                SessionModel.user_id == user.get("id"), SessionModel.id == session_id306            )307        )308 309        result = db.execute(query)310        titles = result.scalars().all()311        print(titles)312 313        if user_prompt_request.streaming:314            return EventSourceResponse(315                generate_streaming_completion(316                    user_prompt_request.prompt,317                    session_id,318                )319            )320        else:321            bot_service = ChatCompletionService(session_id, user_prompt_request.prompt, titles, type_bot="specific")322            response, metadata, scores = bot_service.generate_completion()323 324            existing_session = (325                db.query(SessionModel).filter(SessionModel.id == session_id).first()326            )327            existing_session.updated_at = datetime.now()328            db.commit()329 330            return BotResponse(331                content=response,332                metadata=metadata,333                scores=scores,334            )335 336    except SQLAlchemyError as e:337        return JSONResponse(status_code=500, content=f"Database error: {str(e)}")338    except Exception as e:339        return handle_exception(e)340