MemoryBear/api/app/repositories/model_repository.py

from sqlalchemy.orm import Session, joinedload
from sqlalchemy import and_, or_, func, desc
from typing import List, Optional, Dict, Any, Tuple
import uuid

from app.models.models_model import ModelConfig, ModelApiKey, ModelType
from app.schemas.model_schema import (
    ModelConfigUpdate, ModelApiKeyCreate, ModelApiKeyUpdate,
    ModelConfigQuery
)
from app.core.logging_config import get_db_logger

# 获取数据库专用日志器
db_logger = get_db_logger()


class ModelConfigRepository:
    """模型配置Repository"""

    @staticmethod
    def get_by_id(db: Session, model_id: uuid.UUID, tenant_id: uuid.UUID | None = None) -> Optional[ModelConfig]:
        """根据ID获取模型配置"""
        db_logger.debug(f"根据ID查询模型配置: model_id={model_id}, tenant_id={tenant_id}")
        
        try:
            query = db.query(ModelConfig).options(
                joinedload(ModelConfig.api_keys)
            ).filter(ModelConfig.id == model_id)
            
            # 添加租户过滤
            if tenant_id:
                query = query.filter(
                    or_(
                        ModelConfig.tenant_id == tenant_id,
                        ModelConfig.is_public
                    )
                )
            
            model = query.first()
            
            if model:
                db_logger.debug(f"模型配置查询成功: {model.name} (ID: {model_id})")
            else:
                db_logger.debug(f"模型配置不存在: model_id={model_id}")
            return model
        except Exception as e:
            db_logger.error(f"根据ID查询模型配置失败: model_id={model_id} - {str(e)}")
            raise

    @staticmethod
    def get_by_name(db: Session, name: str, tenant_id: uuid.UUID | None = None) -> Optional[ModelConfig]:
        """根据名称获取模型配置"""
        db_logger.debug(f"根据名称查询模型配置: name={name}, tenant_id={tenant_id}")
        
        try:
            query = db.query(ModelConfig).filter(ModelConfig.name == name)
            
            # 添加租户过滤
            if tenant_id:
                query = query.filter(
                    or_(
                        ModelConfig.tenant_id == tenant_id,
                        ModelConfig.is_public
                    )
                )
            
            model = query.first()
            if model:
                db_logger.debug(f"模型配置查询成功: {model.name}")
            return model
        except Exception as e:
            db_logger.error(f"根据名称查询模型配置失败: name={name} - {str(e)}")
            raise

    @staticmethod
    def search_by_name(db: Session, name: str, tenant_id: uuid.UUID | None = None, limit: int = 10) -> List[ModelConfig]:
        """按名称模糊匹配获取模型配置列表
        
        Args:
            name: 模型名称关键词（模糊匹配）
            tenant_id: 租户ID
            limit: 返回数量上限
        Returns:
            模型配置列表
        """
        db_logger.debug(f"按名称模糊查询模型配置: name~{name}, tenant_id={tenant_id}, limit={limit}")
        try:
            query = db.query(ModelConfig).filter(ModelConfig.name.ilike(f"%{name}%"))
            
            # 添加租户过滤
            if tenant_id:
                query = query.filter(
                    or_(
                        ModelConfig.tenant_id == tenant_id,
                        ModelConfig.is_public
                    )
                )
            
            models = query.order_by(ModelConfig.name).limit(limit).all()
            db_logger.debug(f"模糊查询成功: 返回数量={len(models)}")
            return models
        except Exception as e:
            db_logger.error(f"按名称模糊查询模型配置失败: name~{name} - {str(e)}")
            raise

    @staticmethod
    def get_list(db: Session, query: ModelConfigQuery, tenant_id: uuid.UUID | None = None) -> Tuple[List[ModelConfig], int]:
        """获取模型配置列表"""
        db_logger.debug(f"查询模型配置列表: {query.dict()}, tenant_id={tenant_id}")
        
        try:
            # 构建查询条件
            filters = []
            
            # 添加租户过滤（查询本租户的模型或公开模型）
            if tenant_id:
                filters.append(
                    or_(
                        ModelConfig.tenant_id == tenant_id,
                        ModelConfig.is_public
                    )
                )
            
            # 支持多个 type 值（使用 IN 查询）
            # 兼容 chat 和 llm 类型：如果查询包含其中一个，则同时匹配两者
            if query.type:
                type_values = list(query.type)
                # 如果包含 chat 或 llm，则同时包含两者
                if ModelType.CHAT in type_values or ModelType.LLM in type_values:
                    if ModelType.CHAT not in type_values:
                        type_values.append(ModelType.CHAT)
                    if ModelType.LLM not in type_values:
                        type_values.append(ModelType.LLM)
                filters.append(ModelConfig.type.in_(type_values))
            
            if query.is_active is not None:
                filters.append(ModelConfig.is_active == query.is_active)
            
            if query.is_public is not None:
                filters.append(ModelConfig.is_public == query.is_public)
            
            if query.search:
                # 搜索逻辑需要join ModelApiKey表来搜索model_name
                search_filter = or_(
                    ModelConfig.name.ilike(f"%{query.search}%"),
                    # ModelConfig.description.ilike(f"%{query.search}%")
                )
                filters.append(search_filter)
            
            # 构建基础查询
            base_query = db.query(ModelConfig).options(
                joinedload(ModelConfig.api_keys)
            )
            
            # 如果需要按provider筛选，需要join ModelApiKey表
            if query.provider:
                base_query = base_query.join(ModelApiKey).filter(
                    ModelApiKey.provider == query.provider
                ).distinct()
            
            if filters:
                base_query = base_query.filter(and_(*filters))
            
            # 获取总数
            total = base_query.count()
            
            # 分页查询
            models = base_query.order_by(desc(ModelConfig.updated_at)).offset(
                (query.page - 1) * query.pagesize
            ).limit(query.pagesize).all()
            
            db_logger.debug(f"模型配置列表查询成功: 总数={total}, 当前页={len(models)}, type筛选={query.type}")
            return models, total
            
        except Exception as e:
            db_logger.error(f"查询模型配置列表失败: {str(e)}")
            raise

    @staticmethod
    def get_by_type(db: Session, model_type: ModelType, tenant_id: uuid.UUID | None = None, is_active: bool = True) -> List[ModelConfig]:
        """根据类型获取模型配置"""
        db_logger.debug(f"根据类型查询模型配置: type={model_type}, tenant_id={tenant_id}, is_active={is_active}")
        
        try:
            query = db.query(ModelConfig).options(
                joinedload(ModelConfig.api_keys)
            ).filter(ModelConfig.type == model_type)
            
            # 添加租户过滤
            if tenant_id:
                query = query.filter(
                    or_(
                        ModelConfig.tenant_id == tenant_id,
                        ModelConfig.is_public
                    )
                )
            
            if is_active:
                query = query.filter(ModelConfig.is_active)
            
            models = query.order_by(ModelConfig.name).all()
            db_logger.debug(f"根据类型查询模型配置成功: 数量={len(models)}")
            return models
            
        except Exception as e:
            db_logger.error(f"根据类型查询模型配置失败: type={model_type} - {str(e)}")
            raise

    @staticmethod
    def create(db: Session, model_data: dict) -> ModelConfig:
        """创建模型配置"""
        db_logger.debug(f"创建模型配置: {model_data.get('name')}")
        
        try:
            db_model = ModelConfig(**model_data)
            db.add(db_model)
            
            db_logger.info(f"模型配置已添加到会话: {db_model.name}")
            return db_model
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"创建模型配置失败: {model_data.get('name')} - {str(e)}")
            raise

    @staticmethod
    def update(db: Session, model_id: uuid.UUID, model_data: ModelConfigUpdate, tenant_id: uuid.UUID | None = None) -> Optional[ModelConfig]:
        """更新模型配置"""
        db_logger.debug(f"更新模型配置: model_id={model_id}, tenant_id={tenant_id}")
        
        try:
            query = db.query(ModelConfig).filter(ModelConfig.id == model_id)
            
            # 添加租户过滤（只能更新本租户的模型）
            if tenant_id:
                query = query.filter(ModelConfig.tenant_id == tenant_id)
            
            db_model = query.first()
            if not db_model:
                db_logger.warning(f"模型配置不存在或无权限: model_id={model_id}")
                return None
            
            # 更新字段
            update_data = model_data.dict(exclude_unset=True)
            for field, value in update_data.items():
                setattr(db_model, field, value)
            
            db.commit()
            db.refresh(db_model)
            
            db_logger.info(f"模型配置更新成功: {db_model.name} (ID: {model_id})")
            return db_model
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"更新模型配置失败: model_id={model_id} - {str(e)}")
            raise

    @staticmethod
    def delete(db: Session, model_id: uuid.UUID, tenant_id: uuid.UUID | None = None) -> bool:
        """删除模型配置"""
        db_logger.debug(f"删除模型配置: model_id={model_id}, tenant_id={tenant_id}")
        
        try:
            query = db.query(ModelConfig).filter(ModelConfig.id == model_id)
            
            # 添加租户过滤（只能删除本租户的模型）
            if tenant_id:
                query = query.filter(ModelConfig.tenant_id == tenant_id)
            
            db_model = query.first()
            if not db_model:
                db_logger.warning(f"模型配置不存在或无权限: model_id={model_id}")
                return False
            
            # 逻辑删除模型配置
            db_model.is_active = False
            db.commit()
            
            db_logger.info(f"模型配置删除成功（逻辑删除）: model_id={model_id}")
            return True
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"删除模型配置失败: model_id={model_id} - {str(e)}")
            raise

    @staticmethod
    def get_stats(db: Session) -> Dict[str, Any]:
        """获取模型统计信息"""
        db_logger.debug("获取模型统计信息")
        
        try:
            # 总数统计
            total_models = db.query(ModelConfig).count()
            active_models = db.query(ModelConfig).filter(ModelConfig.is_active).count()
            
            # 按类型统计
            llm_count = db.query(ModelConfig).filter(ModelConfig.type == ModelType.LLM).count()
            embedding_count = db.query(ModelConfig).filter(ModelConfig.type == ModelType.EMBEDDING).count()
            rerank_count = db.query(ModelConfig).filter(ModelConfig.type == ModelType.RERANK).count()
            
            # 按提供商统计 - 现在从ModelApiKey表获取
            provider_stats = {}
            provider_results = db.query(
                ModelApiKey.provider, func.count(func.distinct(ModelApiKey.model_config_id))
            ).group_by(ModelApiKey.provider).all()
            
            for provider, count in provider_results:
                provider_stats[provider.value] = count
            
            stats = {
                "total_models": total_models,
                "active_models": active_models,
                "llm_count": llm_count,
                "embedding_count": embedding_count,
                "rerank_count": rerank_count,
                "provider_stats": provider_stats
            }
            
            db_logger.debug(f"模型统计信息获取成功: {stats}")
            return stats
            
        except Exception as e:
            db_logger.error(f"获取模型统计信息失败: {str(e)}")
            raise


class ModelApiKeyRepository:
    """模型API Key Repository"""

    @staticmethod
    def get_by_id(db: Session, api_key_id: uuid.UUID) -> Optional[ModelApiKey]:
        """根据ID获取API Key"""
        db_logger.debug(f"根据ID查询API Key: api_key_id={api_key_id}")
        
        try:
            api_key = db.query(ModelApiKey).filter(ModelApiKey.id == api_key_id).first()
            if api_key:
                db_logger.debug(f"API Key查询成功: {api_key.model_name} (ID: {api_key_id})")
            return api_key
        except Exception as e:
            db_logger.error(f"根据ID查询API Key失败: api_key_id={api_key_id} - {str(e)}")
            raise

    @staticmethod
    def get_by_model_config(db: Session, model_config_id: uuid.UUID, is_active: bool = True) -> List[ModelApiKey]:
        """根据模型配置ID获取API Key列表"""
        db_logger.debug(f"根据模型配置ID查询API Key: model_config_id={model_config_id}")
        
        try:
            query = db.query(ModelApiKey).filter(ModelApiKey.model_config_id == model_config_id)
            
            if is_active:
                query = query.filter(ModelApiKey.is_active)
            
            api_keys = query.order_by(ModelApiKey.priority, ModelApiKey.created_at).all()
            db_logger.debug(f"API Key列表查询成功: 数量={len(api_keys)}")
            return api_keys
            
        except Exception as e:
            db_logger.error(f"根据模型配置ID查询API Key失败: model_config_id={model_config_id} - {str(e)}")
            raise

    @staticmethod
    def create(db: Session, api_key_data: ModelApiKeyCreate) -> ModelApiKey:
        """创建API Key"""
        db_logger.debug(f"创建API Key: {api_key_data.provider}")
        
        try:
            db_api_key = ModelApiKey(**api_key_data.dict())
            db.add(db_api_key)
            
            db_logger.info(f"API Key已添加到会话: {db_api_key.provider}")
            return db_api_key
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"创建API Key失败: {api_key_data.provider} - {str(e)}")
            raise

    @staticmethod
    def update(db: Session, api_key_id: uuid.UUID, api_key_data: ModelApiKeyUpdate) -> Optional[ModelApiKey]:
        """更新API Key"""
        db_logger.debug(f"更新API Key: api_key_id={api_key_id}")
        
        try:
            db_api_key = db.query(ModelApiKey).filter(ModelApiKey.id == api_key_id).first()
            if not db_api_key:
                db_logger.warning(f"API Key不存在: api_key_id={api_key_id}")
                return None
            
            # 更新字段
            update_data = api_key_data.dict(exclude_unset=True)
            for field, value in update_data.items():
                setattr(db_api_key, field, value)
            
            db.commit()
            db.refresh(db_api_key)
            
            db_logger.info(f"API Key更新成功: {db_api_key.model_name} (ID: {api_key_id})")
            return db_api_key
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"更新API Key失败: api_key_id={api_key_id} - {str(e)}")
            raise

    @staticmethod
    def delete(db: Session, api_key_id: uuid.UUID) -> bool:
        """删除API Key"""
        db_logger.debug(f"删除API Key: api_key_id={api_key_id}")
        
        try:
            db_api_key = db.query(ModelApiKey).filter(ModelApiKey.id == api_key_id).first()
            if not db_api_key:
                db_logger.warning(f"API Key不存在: api_key_id={api_key_id}")
                return False
            
            # 逻辑删除 API Key
            db_api_key.is_active = False
            db.commit()
            
            db_logger.info(f"API Key删除成功（逻辑删除）: api_key_id={api_key_id}")
            return True
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"删除API Key失败: api_key_id={api_key_id} - {str(e)}")
            raise

    @staticmethod
    def update_usage(db: Session, api_key_id: uuid.UUID) -> bool:
        """更新API Key使用统计"""
        db_logger.debug(f"更新API Key使用统计: api_key_id={api_key_id}")
        
        try:
            db_api_key = db.query(ModelApiKey).filter(ModelApiKey.id == api_key_id).first()
            if not db_api_key:
                return False
            
            # 更新使用次数和最后使用时间
            current_count = int(db_api_key.usage_count or "0")
            db_api_key.usage_count = str(current_count + 1)
            db_api_key.last_used_at = func.now()
            
            db.commit()
            db_logger.debug(f"API Key使用统计更新成功: api_key_id={api_key_id}")
            return True
            
        except Exception as e:
            db.rollback()
            db_logger.error(f"更新API Key使用统计失败: api_key_id={api_key_id} - {str(e)}")
            raise