Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
from datetime import datetime | |
import logging | |
from fastapi import APIRouter, Depends, HTTPException, Query, Body | |
from pydantic import BaseModel | |
from typing import List, Optional, Dict, Any | |
from src.db.init_db import session_factory | |
from src.db.schemas.models import ModelUsage | |
from src.managers.ai_manager import AI_Manager | |
from src.managers.chat_manager import ChatManager | |
from src.managers.user_manager import get_current_user, User | |
from src.schemas.chat_schema import * | |
from src.utils.logger import Logger | |
import os | |
from dotenv import load_dotenv | |
load_dotenv() | |
# Initialize logger with console logging disabled | |
logger = Logger("chat_routes", see_time=True, console_log=False) | |
# Initialize router | |
router = APIRouter(prefix="/chats", tags=["chats"]) | |
# Initialize chat manager | |
chat_manager = ChatManager(db_url=os.getenv("DATABASE_URL")) | |
# Initialize AI manager | |
ai_manager = AI_Manager() | |
# Routes | |
async def create_chat(chat_create: ChatCreate): | |
"""Create a new chat session""" | |
try: | |
chat = chat_manager.create_chat(chat_create.user_id) | |
return chat | |
except Exception as e: | |
logger.log_message(f"Error creating chat: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to create chat: {str(e)}") | |
async def add_message(chat_id: int, message: MessageCreate, user_id: Optional[int] = None): | |
"""Add a message to a chat""" | |
try: | |
result = chat_manager.add_message(chat_id, message.content, message.sender, user_id) | |
return result | |
except ValueError as e: | |
raise HTTPException(status_code=404, detail=str(e)) | |
except Exception as e: | |
logger.log_message(f"Error adding message: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to add message: {str(e)}") | |
async def get_chat(chat_id: int, user_id: Optional[int] = None): | |
"""Get a chat by ID with all messages""" | |
try: | |
chat = chat_manager.get_chat(chat_id, user_id) | |
return chat | |
except ValueError as e: | |
raise HTTPException(status_code=404, detail=str(e)) | |
except Exception as e: | |
logger.log_message(f"Error retrieving chat: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to retrieve chat: {str(e)}") | |
async def get_chats( | |
user_id: Optional[int] = None, | |
limit: int = Query(10, ge=1, le=100), | |
offset: int = Query(0, ge=0) | |
): | |
"""Get recent chats, optionally filtered by user_id""" | |
try: | |
chats = chat_manager.get_user_chats(user_id, limit, offset) | |
return chats | |
except Exception as e: | |
logger.log_message(f"Error retrieving chats: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to retrieve chats: {str(e)}") | |
async def delete_chat(chat_id: int, user_id: Optional[int] = None): | |
"""Delete a chat and all its messages while preserving model usage data""" | |
try: | |
# Delete the chat using the updated chat_manager method | |
# which now preserves ModelUsage records | |
success = chat_manager.delete_chat(chat_id, user_id) | |
if not success: | |
raise HTTPException(status_code=404, detail=f"Chat with ID {chat_id} not found or access denied") | |
return {"message": f"Chat {chat_id} deleted successfully", "preserved_model_usage": True} | |
except HTTPException: | |
raise | |
except Exception as e: | |
logger.log_message(f"Error deleting chat: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to delete chat: {str(e)}") | |
async def create_or_get_user(user_info: UserInfo): | |
"""Create a new user or get an existing one by email""" | |
try: | |
user = chat_manager.get_or_create_user( | |
username=user_info.username, | |
email=user_info.email | |
) | |
return user | |
except Exception as e: | |
logger.log_message(f"Error creating/getting user: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to process user: {str(e)}") | |
async def update_chat(chat_id: int, chat_update: ChatUpdate): | |
"""Update a chat's title or user_id""" | |
try: | |
chat = chat_manager.update_chat( | |
chat_id=chat_id, | |
title=chat_update.title, | |
user_id=chat_update.user_id | |
) | |
return chat | |
except ValueError as e: | |
raise HTTPException(status_code=404, detail=str(e)) | |
except Exception as e: | |
logger.log_message(f"Error updating chat: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to update chat: {str(e)}") | |
async def cleanup_empty_chats(request: ChatCreate): | |
"""Delete empty chats for a user""" | |
try: | |
deleted_count = chat_manager.delete_empty_chats(request.user_id, request.is_admin) | |
return {"message": f"Deleted {deleted_count} empty chats"} | |
except Exception as e: | |
logger.log_message(f"Error cleaning up empty chats: {str(e)}", level=logging.ERROR) | |
raise HTTPException(status_code=500, detail=f"Failed to clean up empty chats: {str(e)}") | |