feat: add AI embedding dimension configuration

This commit is contained in:
shiyu
2025-09-17 14:26:53 +08:00
parent 626ff727b3
commit 1a79e87887
6 changed files with 78 additions and 10 deletions

View File

@@ -1,10 +1,11 @@
import httpx
import time
from fastapi import APIRouter, Depends, Form
from fastapi import APIRouter, Depends, Form, HTTPException
from typing import Annotated
from services.config import ConfigCenter, VERSION
from services.auth import get_current_active_user, User, has_users
from api.response import success
from services.vector_db import VectorDBService
router = APIRouter(prefix="/api/config", tags=["config"])
@@ -23,8 +24,27 @@ async def set_config(
key: str = Form(...),
value: str = Form(...)
):
await ConfigCenter.set(key, value)
return success({"key": key, "value": value})
original_value = await ConfigCenter.get(key)
value_to_save = value
if key == "AI_EMBED_DIM":
try:
parsed_value = int(value)
except (TypeError, ValueError):
raise HTTPException(status_code=400, detail="AI_EMBED_DIM must be an integer")
if parsed_value <= 0:
raise HTTPException(status_code=400, detail="AI_EMBED_DIM must be greater than zero")
value_to_save = str(parsed_value)
await ConfigCenter.set(key, value_to_save)
if key == "AI_EMBED_DIM" and str(original_value) != value_to_save:
try:
service = VectorDBService()
service.clear_all_data()
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Failed to clear vector database: {exc}")
return success({"key": key, "value": value_to_save})
@router.get("/all")

View File

@@ -2,8 +2,9 @@ from typing import Dict, Any
from fastapi.responses import Response
import base64
from services.ai import describe_image_base64, get_text_embedding
from services.vector_db import VectorDBService
from services.vector_db import VectorDBService, DEFAULT_VECTOR_DIMENSION
from services.logging import LogService
from services.config import ConfigCenter
class VectorIndexProcessor:
@@ -71,7 +72,15 @@ class VectorIndexProcessor:
if embedding is None:
return Response(content="不支持的文件类型", status_code=400)
vector_db.ensure_collection(collection_name, vector=True)
raw_dim = await ConfigCenter.get('AI_EMBED_DIM', DEFAULT_VECTOR_DIMENSION)
try:
vector_dim = int(raw_dim)
except (TypeError, ValueError):
vector_dim = DEFAULT_VECTOR_DIMENSION
if vector_dim <= 0:
vector_dim = DEFAULT_VECTOR_DIMENSION
vector_db.ensure_collection(collection_name, vector=True, dim=vector_dim)
vector_db.upsert_vector(
collection_name, {'path': path, 'embedding': embedding})

View File

@@ -1,6 +1,9 @@
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
DEFAULT_VECTOR_DIMENSION = 4096
class VectorDBService:
_instance = None
@@ -13,15 +16,21 @@ class VectorDBService:
if not hasattr(self, 'client'):
self.client = MilvusClient("data/db/milvus.db")
def ensure_collection(self, collection_name, vector: bool = True):
def ensure_collection(self, collection_name, vector: bool = True, dim: int = DEFAULT_VECTOR_DIMENSION):
if self.client.has_collection(collection_name):
return
if vector:
try:
vector_dim = int(dim)
except (TypeError, ValueError):
vector_dim = DEFAULT_VECTOR_DIMENSION
if vector_dim <= 0:
vector_dim = DEFAULT_VECTOR_DIMENSION
fields = [
FieldSchema(name="path", dtype=DataType.VARCHAR,
max_length=512, is_primary=True, auto_id=False),
FieldSchema(name="embedding",
dtype=DataType.FLOAT_VECTOR, dim=4096)
dtype=DataType.FLOAT_VECTOR, dim=vector_dim)
]
schema = CollectionSchema(
fields, description="Image vector collection")

View File

@@ -202,9 +202,12 @@ export const en = {
'AI Settings': 'AI Settings',
'Vision Model': 'Vision Model',
'Embedding Model': 'Embedding Model',
'Embedding Dimension': 'Embedding Dimension',
'Vector Database': 'Vector Database',
'Vector Database Settings': 'Vector Database Settings',
'Database Type': 'Database Type',
'Confirm embedding dimension change': 'Confirm embedding dimension change',
'Changing the embedding dimension will clear the vector database automatically. You will need to rebuild indexes afterwards. Continue?': 'Changing the embedding dimension will clear the vector database automatically. You will need to rebuild indexes afterwards. Continue?',
'Confirm clear vector database?': 'Confirm clear vector database?',
'This will delete all collections irreversibly.': 'This will delete all collections irreversibly.',
'Confirm Clear': 'Confirm Clear',

View File

@@ -204,9 +204,12 @@ export const zh = {
'AI Settings': 'AI设置',
'Vision Model': '视觉模型',
'Embedding Model': '嵌入模型',
'Embedding Dimension': '向量维度',
'Vector Database': '向量数据库',
'Vector Database Settings': '向量数据库设置',
'Database Type': '数据库类型',
'Confirm embedding dimension change': '确认修改向量维度',
'Changing the embedding dimension will clear the vector database automatically. You will need to rebuild indexes afterwards. Continue?': '修改向量维度会自动清空向量数据库,之后需要重建索引,是否继续?',
'Confirm clear vector database?': '确认清空向量数据库?',
'This will delete all collections irreversibly.': '此操作将删除所有集合中的所有数据,且不可逆。',
'Confirm Clear': '确认清空',

View File

@@ -21,13 +21,16 @@ const VISION_CONFIG_KEYS = [
{ key: 'AI_VISION_API_KEY', label: 'Vision API Key' },
];
const DEFAULT_EMBED_DIMENSION = 4096;
const EMBED_DIM_KEY = 'AI_EMBED_DIM';
const EMBED_CONFIG_KEYS = [
{ key: 'AI_EMBED_API_URL', label: 'Embedding API URL' },
{ key: 'AI_EMBED_MODEL', label: 'Embedding Model', default: 'Qwen/Qwen3-Embedding-8B' },
{ key: 'AI_EMBED_API_KEY', label: 'Embedding API Key' },
];
const ALL_AI_KEYS = [...VISION_CONFIG_KEYS, ...EMBED_CONFIG_KEYS];
const ALL_AI_KEYS = [...VISION_CONFIG_KEYS, ...EMBED_CONFIG_KEYS, { key: EMBED_DIM_KEY, default: DEFAULT_EMBED_DIMENSION }];
// Theme related config keys
const THEME_KEYS = {
@@ -213,9 +216,27 @@ export default function SystemSettingsPage() {
<Form
layout="vertical"
initialValues={{
...Object.fromEntries(ALL_AI_KEYS.map(({ key, default: def }) => [key, config[key] ?? def ?? ''])),
...Object.fromEntries(ALL_AI_KEYS.map(({ key, default: def }) => [key, key === EMBED_DIM_KEY
? Number(config[key] ?? def ?? DEFAULT_EMBED_DIMENSION)
: config[key] ?? def ?? ''])),
}}
onFinish={async (vals) => {
const currentDim = Number(config[EMBED_DIM_KEY] ?? DEFAULT_EMBED_DIMENSION);
const nextDim = Number(vals[EMBED_DIM_KEY] ?? DEFAULT_EMBED_DIMENSION);
if (currentDim !== nextDim) {
Modal.confirm({
title: t('Confirm embedding dimension change'),
content: t('Changing the embedding dimension will clear the vector database automatically. You will need to rebuild indexes afterwards. Continue?'),
okText: t('Confirm'),
cancelText: t('Cancel'),
onOk: async () => {
await handleSave(vals);
},
});
return;
}
await handleSave(vals);
}}
onFinish={handleSave}
style={{ marginTop: 24 }}
key={JSON.stringify(config)}
>
@@ -232,6 +253,9 @@ export default function SystemSettingsPage() {
<Input size="large" />
</Form.Item>
))}
<Form.Item name={EMBED_DIM_KEY} label={t('Embedding Dimension')}>
<InputNumber min={1} max={32768} style={{ width: '100%' }} />
</Form.Item>
</Card>
<Form.Item style={{ marginTop: 24 }}>
<Button type="primary" htmlType="submit" loading={loading} block>