CodeReview/backend/app/api/v1/endpoints/embedding_config.py

"""
嵌入模型配置 API
独立于 LLM 配置，专门用于 RAG 系统的嵌入模型
使用 UserConfig.other_config 持久化存储
"""

import json
import uuid
from typing import Any, Optional, List
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified

from app.api import deps
from app.models.user import User
from app.models.user_config import UserConfig
from app.core.config import settings

router = APIRouter()


# ============ Schemas ============

class EmbeddingProvider(BaseModel):
    """嵌入模型提供商"""
    id: str
    name: str
    description: str
    models: List[str]
    requires_api_key: bool
    default_model: str


class EmbeddingConfig(BaseModel):
    """嵌入模型配置"""
    provider: str = Field(description="提供商: openai, ollama, azure, cohere, huggingface, jina, qwen")
    model: str = Field(description="模型名称")
    api_key: Optional[str] = Field(default=None, description="API Key (如需要)")
    base_url: Optional[str] = Field(default=None, description="自定义 API 端点")
    dimensions: Optional[int] = Field(default=None, description="向量维度 (某些模型支持)")
    batch_size: int = Field(default=100, description="批处理大小")


class EmbeddingConfigResponse(BaseModel):
    """配置响应"""
    provider: str
    model: str
    api_key: Optional[str] = None  # 返回 API Key
    base_url: Optional[str]
    dimensions: int
    batch_size: int


class TestEmbeddingRequest(BaseModel):
    """测试嵌入请求"""
    provider: str
    model: str
    api_key: Optional[str] = None
    base_url: Optional[str] = None
    test_text: str = "这是一段测试文本，用于验证嵌入模型是否正常工作。"


class TestEmbeddingResponse(BaseModel):
    """测试嵌入响应"""
    success: bool
    message: str
    dimensions: Optional[int] = None
    sample_embedding: Optional[List[float]] = None  # 前 5 个维度
    latency_ms: Optional[int] = None


# ============ 提供商配置 ============

EMBEDDING_PROVIDERS: List[EmbeddingProvider] = [
    EmbeddingProvider(
        id="openai",
        name="OpenAI (兼容 DeepSeek/Moonshot/智谱 等)",
        description="OpenAI 官方或兼容 API，填写自定义端点可接入其他服务商",
        models=[
            "text-embedding-3-small",
            "text-embedding-3-large",
            "text-embedding-ada-002",
        ],
        requires_api_key=True,
        default_model="text-embedding-3-small",
    ),
    EmbeddingProvider(
        id="azure",
        name="Azure OpenAI",
        description="Azure 托管的 OpenAI 嵌入模型",
        models=[
            "text-embedding-3-small",
            "text-embedding-3-large",
            "text-embedding-ada-002",
        ],
        requires_api_key=True,
        default_model="text-embedding-3-small",
    ),
    EmbeddingProvider(
        id="ollama",
        name="Ollama (本地)",
        description="本地运行的开源嵌入模型 (使用 /api/embed 端点)",
        models=[
            "nomic-embed-text",
            "mxbai-embed-large",
            "all-minilm",
            "snowflake-arctic-embed",
            "bge-m3",
            "qwen3-embedding",
        ],
        requires_api_key=False,
        default_model="nomic-embed-text",
    ),
    EmbeddingProvider(
        id="cohere",
        name="Cohere",
        description="Cohere Embed v2 API (api.cohere.com/v2)",
        models=[
            "embed-english-v3.0",
            "embed-multilingual-v3.0",
            "embed-english-light-v3.0",
            "embed-multilingual-light-v3.0",
            "embed-v4.0",
        ],
        requires_api_key=True,
        default_model="embed-multilingual-v3.0",
    ),
    EmbeddingProvider(
        id="huggingface",
        name="HuggingFace",
        description="HuggingFace Inference Providers (router.huggingface.co)",
        models=[
            "sentence-transformers/all-MiniLM-L6-v2",
            "sentence-transformers/all-mpnet-base-v2",
            "BAAI/bge-large-zh-v1.5",
            "BAAI/bge-m3",
        ],
        requires_api_key=True,
        default_model="BAAI/bge-m3",
    ),
    EmbeddingProvider(
        id="jina",
        name="Jina AI",
        description="Jina AI 嵌入模型，代码嵌入效果好",
        models=[
            "jina-embeddings-v2-base-code",
            "jina-embeddings-v2-base-en",
            "jina-embeddings-v2-base-zh",
        ],
        requires_api_key=True,
        default_model="jina-embeddings-v2-base-code",
    ),
    EmbeddingProvider(
        id="qwen",
        name="Qwen (DashScope)",
        description="阿里云 DashScope Qwen 嵌入模型，兼容 OpenAI embeddings 接口",
        models=[
            "text-embedding-v4",
            "text-embedding-v3",
            "text-embedding-v2",
        ],
        requires_api_key=True,
        default_model="text-embedding-v4",
    ),
]


# ============ 数据库持久化存储 (异步) ============

EMBEDDING_CONFIG_KEY = "embedding_config"


def mask_api_key(key: Optional[str]) -> str:
    """部分遮盖API Key，显示前3位和后4位"""
    if not key:
        return ""
    if len(key) <= 8:
        return "***"
    return f"{key[:3]}***{key[-4:]}"


async def get_embedding_config_from_db(db: AsyncSession, user_id: str) -> EmbeddingConfig:
    """从数据库获取嵌入配置（异步）"""
    # 嵌入配置始终来自系统默认（.env），不再允许用户覆盖
    print(f"[EmbeddingConfig] 返回系统默认嵌入配置（来自 .env）")
    return EmbeddingConfig(
        provider=settings.EMBEDDING_PROVIDER,
        model=settings.EMBEDDING_MODEL,
        api_key=mask_api_key(settings.EMBEDDING_API_KEY or settings.LLM_API_KEY),
        base_url=settings.EMBEDDING_BASE_URL or settings.LLM_BASE_URL,
        dimensions=settings.EMBEDDING_DIMENSION if settings.EMBEDDING_DIMENSION > 0 else None,
        batch_size=100,
    )




# ============ API Endpoints ============

@router.get("/providers", response_model=List[EmbeddingProvider])
async def list_embedding_providers(
    current_user: User = Depends(deps.get_current_user),
) -> Any:
    """
    获取可用的嵌入模型提供商列表
    """
    return EMBEDDING_PROVIDERS


@router.get("/config", response_model=EmbeddingConfigResponse)
async def get_current_config(
    db: AsyncSession = Depends(deps.get_db),
    current_user: User = Depends(deps.get_current_user),
) -> Any:
    """
    获取当前嵌入模型配置（从数据库读取）
    """
    config = await get_embedding_config_from_db(db, current_user.id)

    # 获取维度
    dimensions = config.dimensions
    if not dimensions or dimensions <= 0:
        dimensions = _get_model_dimensions(config.provider, config.model)

    return EmbeddingConfigResponse(
        provider=config.provider,
        model=config.model,
        api_key=config.api_key,
        base_url=config.base_url,
        dimensions=dimensions,
        batch_size=config.batch_size,
    )


@router.put("/config")
async def update_config(
    current_user: User = Depends(deps.get_current_user),
) -> Any:
    """
    更新嵌入模型配置（已禁用，固定从 .env 读取）
    """
    return {"message": "嵌入模型配置已锁定，请在 .env 文件中进行修改", "provider": settings.EMBEDDING_PROVIDER, "model": settings.EMBEDDING_MODEL}


@router.post("/test", response_model=TestEmbeddingResponse)
async def test_embedding(
    request: TestEmbeddingRequest,
    current_user: User = Depends(deps.get_current_user),
) -> Any:
    """测试当前系统的嵌入模型配置"""
    import time
    from app.services.rag.embeddings import EmbeddingService
    
    try:
        start_time = time.time()
        
        # 始终使用系统的真实配置进行测试
        # API Key 优先级：EMBEDDING_API_KEY > LLM_API_KEY
        api_key = getattr(settings, "EMBEDDING_API_KEY", None) or settings.LLM_API_KEY
        
        service = EmbeddingService(
            provider=settings.EMBEDDING_PROVIDER,
            model=settings.EMBEDDING_MODEL,
            api_key=api_key,
            base_url=settings.EMBEDDING_BASE_URL,
            cache_enabled=False,
        )
        
        # 执行嵌入
        embedding = await service.embed(request.test_text)
        
        latency_ms = int((time.time() - start_time) * 1000)
        
        return TestEmbeddingResponse(
            success=True,
            message=f"嵌入成功! 提供商: {settings.EMBEDDING_PROVIDER}, 维度: {len(embedding)}",
            dimensions=len(embedding),
            sample_embedding=embedding[:5],  # 返回前 5 维
            latency_ms=latency_ms,
        )
        
    except Exception as e:
        print(f"❌ 嵌入测试失败: {str(e)}")
        return TestEmbeddingResponse(
            success=False,
            message=f"嵌入失败: {str(e)}",
        )


@router.get("/models/{provider}")
async def get_provider_models(
    provider: str,
    current_user: User = Depends(deps.get_current_user),
) -> Any:
    """
    获取指定提供商的模型列表
    """
    provider_info = next((p for p in EMBEDDING_PROVIDERS if p.id == provider), None)
    
    if not provider_info:
        raise HTTPException(status_code=404, detail=f"提供商不存在: {provider}")
    
    return {
        "provider": provider,
        "models": provider_info.models,
        "default_model": provider_info.default_model,
        "requires_api_key": provider_info.requires_api_key,
    }


def _get_model_dimensions(provider: str, model: str) -> int:
    """获取模型维度"""
    dimensions_map = {
        # OpenAI
        "text-embedding-3-small": 1536,
        "text-embedding-3-large": 3072,
        "text-embedding-ada-002": 1536,
        
        # Ollama
        "nomic-embed-text": 768,
        "mxbai-embed-large": 1024,
        "all-minilm": 384,
        "snowflake-arctic-embed": 1024,
        
        # Cohere
        "embed-english-v3.0": 1024,
        "embed-multilingual-v3.0": 1024,
        "embed-english-light-v3.0": 384,
        "embed-multilingual-light-v3.0": 384,
        
        # HuggingFace
        "sentence-transformers/all-MiniLM-L6-v2": 384,
        "sentence-transformers/all-mpnet-base-v2": 768,
        "BAAI/bge-large-zh-v1.5": 1024,
        "BAAI/bge-m3": 1024,
        
        # Jina
        "jina-embeddings-v2-base-code": 768,
        "jina-embeddings-v2-base-en": 768,
        "jina-embeddings-v2-base-zh": 768,
        
        # Qwen (DashScope)
        "text-embedding-v4": 1024,  # 支持维度: 2048, 1536, 1024(默认), 768, 512, 256, 128, 64
        "text-embedding-v3": 1024,  # 支持维度: 1024(默认), 768, 512, 256, 128, 64
        "text-embedding-v2": 1536,  # 支持维度: 1536
    }
    
    return dimensions_map.get(model, 768)