langchain-learning-kit/app/services/model_manager.py

"""
模型管理器服务，用于LLM和Embedding模型管理。
"""
import os
import re
from typing import Optional, List, Dict, Any

from sqlalchemy.orm import Session
from langchain_core.language_models.llms import BaseLLM
from langchain_core.embeddings import Embeddings
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from app.db.models import Model
from app.config import Settings
from app.utils.exceptions import ResourceNotFoundError, InvalidConfigError, DuplicateResourceError
from app.utils.logger import get_logger

logger = get_logger(__name__)


class ModelManager:
    """
    管理LLM和Embedding模型配置。
    """

    def __init__(self, db_session: Session, settings: Settings):
        """
        初始化模型管理器。

        Args:
            db_session: 数据库会话
            settings: 应用程序配置
        """
        self.db = db_session
        self.settings = settings

    def _replace_env_vars(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """
        替换配置中的环境变量占位符。

        Args:
            config: 配置字典

        Returns:
            Dict[str, Any]: 替换环境变量后的配置
        """
        result = {}
        for key, value in config.items():
            if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
                # 提取变量名: ${VAR_NAME} -> VAR_NAME
                var_name = value[2:-1]

                # 首先尝试从os.environ获取，然后从settings获取
                env_value = os.getenv(var_name)
                if env_value is None:
                    # 尝试从settings对象获取
                    if hasattr(self.settings, var_name.lower()):
                        env_value = getattr(self.settings, var_name.lower())

                if env_value is None:
                    raise InvalidConfigError(
                        f"Environment variable {var_name} not found",
                        details={"variable": var_name}
                    )
                result[key] = env_value
            else:
                result[key] = value
        return result

    def create_model(
        self,
        name: str,
        type: str,
        config: Dict[str, Any],
        is_default: bool = False,
    ) -> Model:
        """
        创建新的模型配置。

        Args:
            name: 模型名称
            type: 模型类型（'llm' 或 'embedding'）
            config: 模型配置
            is_default: 设置为默认模型

        Returns:
            Model: 创建的模型实例

        Raises:
            DuplicateResourceError: 如果存在同名模型
            InvalidConfigError: 如果配置无效
        """
        logger.info("creating_model", name=name, type=type)

        # 检查模型是否已存在
        existing = self.db.query(Model).filter(Model.name == name).first()
        if existing:
            raise DuplicateResourceError(f"Model '{name}' already exists")

        # 验证类型
        if type not in ["llm", "embedding"]:
            raise InvalidConfigError(f"Invalid model type: {type}. Must be 'llm' or 'embedding'")

        # 如果设置为默认，取消其他默认设置
        if is_default:
            self.db.query(Model).filter(Model.type == type, Model.is_default == True).update(
                {"is_default": False}
            )

        # 创建模型
        model = Model(
            name=name,
            type=type,
            config=config,
            is_default=is_default,
            status="active",
        )

        self.db.add(model)
        self.db.commit()
        self.db.refresh(model)

        logger.info("model_created", model_id=model.id, name=name)
        return model

    def get_model(self, model_id: int) -> Model:
        """
        根据ID获取模型。

        Args:
            model_id: 模型ID

        Returns:
            Model: 模型实例

        Raises:
            ResourceNotFoundError: 如果模型未找到
        """
        model = self.db.query(Model).filter(Model.id == model_id).first()
        if not model:
            raise ResourceNotFoundError(f"Model not found: {model_id}")
        return model

    def get_model_by_name(self, name: str) -> Model:
        """
        根据名称获取模型。

        Args:
            name: 模型名称

        Returns:
            Model: 模型实例

        Raises:
            ResourceNotFoundError: 如果模型未找到
        """
        model = self.db.query(Model).filter(Model.name == name).first()
        if not model:
            raise ResourceNotFoundError(f"Model not found: {name}")
        return model

    def list_models(self, type: Optional[str] = None) -> List[Model]:
        """
        列出所有模型，可选择按类型过滤。

        Args:
            type: 可选，按模型类型过滤

        Returns:
            List[Model]: 模型列表
        """
        query = self.db.query(Model).filter(Model.status == "active")
        if type:
            query = query.filter(Model.type == type)
        return query.all()

    def update_model(
        self,
        model_id: int,
        config: Optional[Dict[str, Any]] = None,
        is_default: Optional[bool] = None,
        status: Optional[str] = None,
    ) -> Model:
        """
        更新模型配置。

        Args:
            model_id: 模型ID
            config: 新配置
            is_default: 更新默认状态
            status: 更新状态

        Returns:
            Model: 更新后的模型实例

        Raises:
            ResourceNotFoundError: 如果模型未找到
        """
        model = self.get_model(model_id)

        if config is not None:
            model.config = config

        if is_default is not None and is_default:
            # 取消其他默认设置
            self.db.query(Model).filter(
                Model.type == model.type,
                Model.is_default == True,
                Model.id != model_id
            ).update({"is_default": False})
            model.is_default = True

        if status is not None:
            model.status = status

        self.db.commit()
        self.db.refresh(model)

        logger.info("model_updated", model_id=model_id)
        return model

    def delete_model(self, model_id: int) -> bool:
        """
        删除（软删除）模型。

        Args:
            model_id: 模型ID

        Returns:
            bool: 如果删除成功返回True

        Raises:
            ResourceNotFoundError: 如果模型未找到
        """
        model = self.get_model(model_id)
        model.status = "deleted"
        self.db.commit()

        logger.info("model_deleted", model_id=model_id)
        return True

    def get_llm(self, name: Optional[str] = None) -> BaseLLM:
        """
        根据名称获取LLM实例或使用默认实例。

        Args:
            name: 模型名称（None表示使用默认）

        Returns:
            BaseLLM: LangChain LLM实例

        Raises:
            ResourceNotFoundError: 如果模型未找到
            InvalidConfigError: 如果模型类型不是'llm'
        """
        if name:
            model = self.get_model_by_name(name)
        else:
            # 获取默认LLM
            model = self.db.query(Model).filter(
                Model.type == "llm",
                Model.is_default == True,
                Model.status == "active"
            ).first()

            if not model:
                raise ResourceNotFoundError("No default LLM model configured")

        if model.type != "llm":
            raise InvalidConfigError(f"Model '{model.name}' is not an LLM (type: {model.type})")

        # 替换配置中的环境变量
        config = self._replace_env_vars(model.config)

        # 如果需要，修正API密钥参数名称
        if "api_key" in config and "openai_api_key" not in config:
            config["openai_api_key"] = config.pop("api_key")

        # 如果配置中没有API密钥，从settings添加
        if "openai_api_key" not in config:
            if self.settings.openai_api_key:
                config["openai_api_key"] = self.settings.openai_api_key
            else:
                raise InvalidConfigError("OPENAI_API_KEY not configured in settings or model config")

        # 如果已配置，从settings添加base_url
        if self.settings.openai_base_url and "base_url" not in config:
            config["base_url"] = self.settings.openai_base_url

        logger.info("getting_llm", model_name=model.name, base_url=config.get("base_url"))

        # 目前仅支持OpenAI，可以扩展
        return ChatOpenAI(**config)

    def get_embedding(self, name: Optional[str] = None) -> Embeddings:
        """
        根据名称获取Embedding实例或使用默认实例。

        Args:
            name: 模型名称（None表示使用默认）

        Returns:
            Embeddings: LangChain Embeddings实例

        Raises:
            ResourceNotFoundError: 如果模型未找到
            InvalidConfigError: 如果模型类型不是'embedding'
        """
        if name:
            model = self.get_model_by_name(name)
        else:
            # 获取默认embedding
            model = self.db.query(Model).filter(
                Model.type == "embedding",
                Model.is_default == True,
                Model.status == "active"
            ).first()

            if not model:
                raise ResourceNotFoundError("No default embedding model configured")

        if model.type != "embedding":
            raise InvalidConfigError(f"Model '{model.name}' is not an embedding model (type: {model.type})")

        # 替换配置中的环境变量
        config = self._replace_env_vars(model.config)

        # 如果需要，修正API密钥参数名称
        if "api_key" in config and "openai_api_key" not in config:
            config["openai_api_key"] = config.pop("api_key")

        # 如果配置中没有API密钥，从settings添加
        if "openai_api_key" not in config:
            if self.settings.openai_api_key:
                config["openai_api_key"] = self.settings.openai_api_key
            else:
                raise InvalidConfigError("OPENAI_API_KEY not configured in settings or model config")

        # 移除embedding的无效参数
        config.pop("temperature", None)

        # 如果已配置，从settings添加base_url
        if self.settings.openai_base_url and "base_url" not in config:
            config["base_url"] = self.settings.openai_base_url

        logger.info("getting_embedding", model_name=model.name, base_url=config.get("base_url"))

        # 目前仅支持OpenAI，可以扩展
        return OpenAIEmbeddings(**config)