mirror of
https://github.com/DrizzleTime/Foxel.git
synced 2026-09-05 07:36:53 +08:00
refactor: optimize backend module
This commit is contained in:
shiyu
@@ -0,0 +1,34 @@
|
||||
from .api import router_ai, router_vector_db
|
||||
from .service import (
|
||||
AIProviderService,
|
||||
VectorDBConfigManager,
|
||||
VectorDBService,
|
||||
DEFAULT_VECTOR_DIMENSION,
|
||||
ABILITIES,
|
||||
normalize_capabilities,
|
||||
)
|
||||
from .types import (
|
||||
AIDefaultsUpdate,
|
||||
AIModelCreate,
|
||||
AIModelUpdate,
|
||||
AIProviderCreate,
|
||||
AIProviderUpdate,
|
||||
VectorDBConfigPayload,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"router_ai",
|
||||
"router_vector_db",
|
||||
"AIProviderService",
|
||||
"VectorDBService",
|
||||
"VectorDBConfigManager",
|
||||
"DEFAULT_VECTOR_DIMENSION",
|
||||
"ABILITIES",
|
||||
"normalize_capabilities",
|
||||
"AIDefaultsUpdate",
|
||||
"AIModelCreate",
|
||||
"AIModelUpdate",
|
||||
"AIProviderCreate",
|
||||
"AIProviderUpdate",
|
||||
"VectorDBConfigPayload",
|
||||
]
|
||||
@@ -0,0 +1,305 @@
|
||||
from typing import Annotated, Dict, Optional
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Request
|
||||
|
||||
from api.response import success
|
||||
from domain.audit import AuditAction, audit
|
||||
from domain.ai.service import AIProviderService, VectorDBConfigManager, VectorDBService
|
||||
from domain.ai.types import (
|
||||
AIDefaultsUpdate,
|
||||
AIModelCreate,
|
||||
AIModelUpdate,
|
||||
AIProviderCreate,
|
||||
AIProviderUpdate,
|
||||
VectorDBConfigPayload,
|
||||
)
|
||||
from domain.ai.vector_providers import get_provider_class, get_provider_entry, list_providers
|
||||
from domain.auth.service import get_current_active_user
|
||||
from domain.auth.types import User
|
||||
|
||||
router_ai = APIRouter(prefix="/api/ai", tags=["ai"])
|
||||
router_vector_db = APIRouter(prefix="/api/vector-db", tags=["vector-db"])
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取 AI 提供商列表")
|
||||
@router_ai.get("/providers")
|
||||
async def list_providers_endpoint(
|
||||
request: Request,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)]
|
||||
):
|
||||
providers = await AIProviderService.list_providers()
|
||||
return success({"providers": providers})
|
||||
|
||||
|
||||
@audit(
|
||||
action=AuditAction.CREATE,
|
||||
description="创建 AI 提供商",
|
||||
body_fields=["name", "identifier", "provider_type", "api_format", "base_url", "logo_url"],
|
||||
redact_fields=["api_key"],
|
||||
)
|
||||
@router_ai.post("/providers")
|
||||
async def create_provider(
|
||||
request: Request,
|
||||
payload: AIProviderCreate,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)]
|
||||
):
|
||||
provider = await AIProviderService.create_provider(payload.dict())
|
||||
return success(provider)
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取 AI 提供商详情")
|
||||
@router_ai.get("/providers/{provider_id}")
|
||||
async def get_provider(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
provider = await AIProviderService.get_provider(provider_id, with_models=True)
|
||||
return success(provider)
|
||||
|
||||
|
||||
@audit(
|
||||
action=AuditAction.UPDATE,
|
||||
description="更新 AI 提供商",
|
||||
body_fields=["name", "provider_type", "api_format", "base_url", "logo_url", "api_key"],
|
||||
redact_fields=["api_key"],
|
||||
)
|
||||
@router_ai.put("/providers/{provider_id}")
|
||||
async def update_provider(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
payload: AIProviderUpdate,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
data = {k: v for k, v in payload.dict().items() if v is not None}
|
||||
if not data:
|
||||
raise HTTPException(status_code=400, detail="No fields to update")
|
||||
provider = await AIProviderService.update_provider(provider_id, data)
|
||||
return success(provider)
|
||||
|
||||
|
||||
@audit(action=AuditAction.DELETE, description="删除 AI 提供商")
|
||||
@router_ai.delete("/providers/{provider_id}")
|
||||
async def delete_provider(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
await AIProviderService.delete_provider(provider_id)
|
||||
return success({"id": provider_id})
|
||||
|
||||
|
||||
@audit(action=AuditAction.UPDATE, description="同步模型列表")
|
||||
@router_ai.post("/providers/{provider_id}/sync-models")
|
||||
async def sync_models(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
try:
|
||||
result = await AIProviderService.sync_models(provider_id)
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as exc:
|
||||
raise HTTPException(status_code=502, detail=f"Failed to synchronize models: {exc}") from exc
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
return success(result)
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取远程模型列表")
|
||||
@router_ai.get("/providers/{provider_id}/remote-models")
|
||||
async def fetch_remote_models(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
try:
|
||||
models = await AIProviderService.fetch_remote_models(provider_id)
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as exc:
|
||||
raise HTTPException(status_code=502, detail=f"Failed to pull models: {exc}") from exc
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
return success({"models": models})
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取模型列表")
|
||||
@router_ai.get("/providers/{provider_id}/models")
|
||||
async def list_models(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
models = await AIProviderService.list_models(provider_id)
|
||||
return success({"models": models})
|
||||
|
||||
|
||||
@audit(
|
||||
action=AuditAction.CREATE,
|
||||
description="创建模型",
|
||||
body_fields=["name", "display_name", "capabilities", "context_window", "embedding_dimensions"],
|
||||
)
|
||||
@router_ai.post("/providers/{provider_id}/models")
|
||||
async def create_model(
|
||||
request: Request,
|
||||
provider_id: Annotated[int, Path(..., gt=0)],
|
||||
payload: AIModelCreate,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
model = await AIProviderService.create_model(provider_id, payload.dict())
|
||||
return success(model)
|
||||
|
||||
|
||||
@audit(
|
||||
action=AuditAction.UPDATE,
|
||||
description="更新模型",
|
||||
body_fields=["display_name", "description", "capabilities", "context_window", "embedding_dimensions"],
|
||||
)
|
||||
@router_ai.put("/models/{model_id}")
|
||||
async def update_model(
|
||||
request: Request,
|
||||
model_id: Annotated[int, Path(..., gt=0)],
|
||||
payload: AIModelUpdate,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
data = {k: v for k, v in payload.dict().items() if v is not None}
|
||||
if not data:
|
||||
raise HTTPException(status_code=400, detail="No fields to update")
|
||||
model = await AIProviderService.update_model(model_id, data)
|
||||
return success(model)
|
||||
|
||||
|
||||
@audit(action=AuditAction.DELETE, description="删除模型")
|
||||
@router_ai.delete("/models/{model_id}")
|
||||
async def delete_model(
|
||||
request: Request,
|
||||
model_id: Annotated[int, Path(..., gt=0)],
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
await AIProviderService.delete_model(model_id)
|
||||
return success({"id": model_id})
|
||||
|
||||
|
||||
def _get_embedding_dimension(entry: Optional[Dict]) -> Optional[int]:
|
||||
if not entry:
|
||||
return None
|
||||
value = entry.get("embedding_dimensions")
|
||||
return int(value) if value is not None else None
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取默认模型")
|
||||
@router_ai.get("/defaults")
|
||||
async def get_defaults(
|
||||
request: Request,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
defaults = await AIProviderService.get_default_models()
|
||||
return success(defaults)
|
||||
|
||||
|
||||
@audit(
|
||||
action=AuditAction.UPDATE,
|
||||
description="更新默认模型",
|
||||
body_fields=["chat", "vision", "embedding", "rerank", "voice", "tools"],
|
||||
)
|
||||
@router_ai.put("/defaults")
|
||||
async def update_defaults(
|
||||
request: Request,
|
||||
payload: AIDefaultsUpdate,
|
||||
current_user: Annotated[User, Depends(get_current_active_user)],
|
||||
):
|
||||
previous = await AIProviderService.get_default_models()
|
||||
try:
|
||||
updated = await AIProviderService.set_default_models(payload.as_mapping())
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
prev_dim = _get_embedding_dimension(previous.get("embedding"))
|
||||
next_dim = _get_embedding_dimension(updated.get("embedding"))
|
||||
|
||||
if prev_dim and next_dim and prev_dim != next_dim:
|
||||
try:
|
||||
await VectorDBService().clear_all_data()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise HTTPException(status_code=500, detail=f"Failed to clear vector database: {exc}") from exc
|
||||
|
||||
return success(updated)
|
||||
|
||||
|
||||
@audit(action=AuditAction.UPDATE, description="清空向量数据库")
|
||||
@router_vector_db.post("/clear-all", summary="清空向量数据库")
|
||||
async def clear_vector_db(request: Request, user: User = Depends(get_current_active_user)):
|
||||
try:
|
||||
service = VectorDBService()
|
||||
await service.clear_all_data()
|
||||
return success(msg="向量数据库已清空")
|
||||
except Exception as e: # noqa: BLE001
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取向量数据库统计")
|
||||
@router_vector_db.get("/stats", summary="获取向量数据库统计")
|
||||
async def get_vector_db_stats(request: Request, user: User = Depends(get_current_active_user)):
|
||||
try:
|
||||
service = VectorDBService()
|
||||
data = await service.get_all_stats()
|
||||
return success(data=data)
|
||||
except Exception as e: # noqa: BLE001
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取向量数据库提供者列表")
|
||||
@router_vector_db.get("/providers", summary="列出可用向量数据库提供者")
|
||||
async def list_vector_providers(request: Request, user: User = Depends(get_current_active_user)):
|
||||
return success(list_providers())
|
||||
|
||||
|
||||
@audit(action=AuditAction.READ, description="获取向量数据库配置")
|
||||
@router_vector_db.get("/config", summary="获取当前向量数据库配置")
|
||||
async def get_vector_db_config(request: Request, user: User = Depends(get_current_active_user)):
|
||||
service = VectorDBService()
|
||||
data = await service.current_provider()
|
||||
return success(data)
|
||||
|
||||
|
||||
@audit(action=AuditAction.UPDATE, description="更新向量数据库配置", body_fields=["type"])
|
||||
@router_vector_db.post("/config", summary="更新向量数据库配置")
|
||||
async def update_vector_db_config(
|
||||
request: Request, payload: VectorDBConfigPayload, user: User = Depends(get_current_active_user)
|
||||
):
|
||||
entry = get_provider_entry(payload.type)
|
||||
if not entry:
|
||||
raise HTTPException(
|
||||
status_code=400, detail=f"未知的向量数据库类型: {payload.type}")
|
||||
if not entry.get("enabled", True):
|
||||
raise HTTPException(status_code=400, detail="该向量数据库类型暂不可用")
|
||||
|
||||
provider_cls = get_provider_class(payload.type)
|
||||
if not provider_cls:
|
||||
raise HTTPException(
|
||||
status_code=400, detail=f"未找到类型 {payload.type} 对应的实现")
|
||||
|
||||
test_provider = provider_cls(payload.config)
|
||||
try:
|
||||
await test_provider.initialize()
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
finally:
|
||||
client = getattr(test_provider, "client", None)
|
||||
close_fn = getattr(client, "close", None)
|
||||
if callable(close_fn):
|
||||
try:
|
||||
close_fn()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
await VectorDBConfigManager.save_config(payload.type, payload.config)
|
||||
service = VectorDBService()
|
||||
await service.reload()
|
||||
config_data = await service.current_provider()
|
||||
stats = await service.get_all_stats()
|
||||
return success({"config": config_data, "stats": stats})
|
||||
|
||||
|
||||
__all__ = ["router_ai", "router_vector_db"]
|
||||
@@ -0,0 +1,247 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
from typing import List, Sequence, Tuple
|
||||
|
||||
from models.database import AIModel, AIProvider
|
||||
from domain.ai.service import AIProviderService
|
||||
|
||||
|
||||
provider_service = AIProviderService
|
||||
|
||||
|
||||
class MissingModelError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
async def describe_image_base64(base64_image: str, detail: str = "high") -> str:
|
||||
"""
|
||||
传入 base64 图片并返回描述文本。缺省时返回错误提示。
|
||||
"""
|
||||
try:
|
||||
model, provider = await _require_model("vision")
|
||||
if provider.api_format == "openai":
|
||||
return await _describe_with_openai(provider, model, base64_image, detail)
|
||||
return await _describe_with_gemini(provider, model, base64_image, detail)
|
||||
except MissingModelError as exc:
|
||||
return str(exc)
|
||||
except httpx.ReadTimeout:
|
||||
return "请求超时，请稍后重试。"
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return f"请求失败: {exc}"
|
||||
|
||||
|
||||
async def get_text_embedding(text: str) -> List[float]:
|
||||
"""
|
||||
传入文本，返回嵌入向量。若未配置模型则抛出异常。
|
||||
"""
|
||||
model, provider = await _require_model("embedding")
|
||||
if provider.api_format == "openai":
|
||||
return await _embedding_with_openai(provider, model, text)
|
||||
return await _embedding_with_gemini(provider, model, text)
|
||||
|
||||
|
||||
async def rerank_texts(query: str, documents: Sequence[str]) -> List[float]:
|
||||
"""调用重排序模型，为一组文档返回得分。未配置时返回空列表。"""
|
||||
if not documents:
|
||||
return []
|
||||
try:
|
||||
model, provider = await _require_model("rerank")
|
||||
except MissingModelError:
|
||||
return []
|
||||
|
||||
try:
|
||||
if provider.api_format == "openai":
|
||||
return await _rerank_with_openai(provider, model, query, documents)
|
||||
return await _rerank_with_gemini(provider, model, query, documents)
|
||||
except Exception: # noqa: BLE001
|
||||
return []
|
||||
|
||||
|
||||
async def _require_model(ability: str) -> Tuple[AIModel, AIProvider]:
|
||||
model = await provider_service.get_default_model(ability)
|
||||
if not model:
|
||||
raise MissingModelError(f"未配置默认 {ability} 模型，请前往系统设置完成配置。")
|
||||
provider = getattr(model, "provider", None)
|
||||
if provider is None:
|
||||
await model.fetch_related("provider")
|
||||
provider = model.provider
|
||||
if provider is None:
|
||||
raise MissingModelError("模型缺少关联的提供商配置。")
|
||||
if not provider.base_url:
|
||||
raise MissingModelError("该提供商未设置 API 地址。")
|
||||
return model, provider
|
||||
|
||||
|
||||
def _openai_endpoint(provider: AIProvider, path: str) -> str:
|
||||
base = (provider.base_url or "").rstrip("/")
|
||||
if not base:
|
||||
raise MissingModelError("提供商 API 地址未配置。")
|
||||
return f"{base}/{path.lstrip('/')}"
|
||||
|
||||
|
||||
def _openai_headers(provider: AIProvider) -> dict:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if provider.api_key:
|
||||
headers["Authorization"] = f"Bearer {provider.api_key}"
|
||||
return headers
|
||||
|
||||
|
||||
def _gemini_endpoint(provider: AIProvider, path: str) -> str:
|
||||
base = (provider.base_url or "").rstrip("/")
|
||||
if not base:
|
||||
raise MissingModelError("提供商 API 地址未配置。")
|
||||
url = f"{base}/{path.lstrip('/')}"
|
||||
if provider.api_key:
|
||||
connector = "&" if "?" in url else "?"
|
||||
url = f"{url}{connector}key={provider.api_key}"
|
||||
return url
|
||||
|
||||
|
||||
async def _describe_with_openai(provider: AIProvider, model: AIModel, base64_image: str, detail: str) -> str:
|
||||
url = _openai_endpoint(provider, "/chat/completions")
|
||||
payload = {
|
||||
"model": model.name,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{base64_image}",
|
||||
"detail": detail,
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "描述这个图片"},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
response = await client.post(url, headers=_openai_headers(provider), json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
return body["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
async def _describe_with_gemini(provider: AIProvider, model: AIModel, base64_image: str, detail: str) -> str:
|
||||
detail_text = f"描述这个图片，细节等级：{detail}"
|
||||
model_name = model.name if model.name.startswith("models/") else f"models/{model.name}"
|
||||
url = _gemini_endpoint(provider, f"{model_name}:generateContent")
|
||||
payload = {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{
|
||||
"inline_data": {
|
||||
"mime_type": "image/jpeg",
|
||||
"data": base64_image,
|
||||
}
|
||||
},
|
||||
{"text": detail_text},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
response = await client.post(url, json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
candidates = body.get("candidates") or []
|
||||
if not candidates:
|
||||
return ""
|
||||
parts = candidates[0].get("content", {}).get("parts", [])
|
||||
text_parts = [part.get("text") for part in parts if isinstance(part, dict) and part.get("text")]
|
||||
return "\n".join(text_parts)
|
||||
|
||||
|
||||
async def _embedding_with_openai(provider: AIProvider, model: AIModel, text: str) -> List[float]:
|
||||
url = _openai_endpoint(provider, "/embeddings")
|
||||
payload = {
|
||||
"model": model.name,
|
||||
"input": text,
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, headers=_openai_headers(provider), json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
return body["data"][0]["embedding"]
|
||||
|
||||
|
||||
async def _embedding_with_gemini(provider: AIProvider, model: AIModel, text: str) -> List[float]:
|
||||
model_name = model.name if model.name.startswith("models/") else f"models/{model.name}"
|
||||
url = _gemini_endpoint(provider, f"{model_name}:embedContent")
|
||||
payload = {
|
||||
"model": model_name,
|
||||
"content": {
|
||||
"parts": [{"text": text}],
|
||||
},
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
embedding = body.get("embedding") or {}
|
||||
return embedding.get("values") or []
|
||||
|
||||
|
||||
async def _rerank_with_openai(
|
||||
provider: AIProvider,
|
||||
model: AIModel,
|
||||
query: str,
|
||||
documents: Sequence[str],
|
||||
) -> List[float]:
|
||||
url = _openai_endpoint(provider, "/rerank")
|
||||
payload = {
|
||||
"model": model.name,
|
||||
"query": query,
|
||||
"documents": [
|
||||
{"id": str(idx), "text": content}
|
||||
for idx, content in enumerate(documents)
|
||||
],
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, headers=_openai_headers(provider), json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
results = body.get("results") or body.get("data") or []
|
||||
scores: List[float] = []
|
||||
for item in results:
|
||||
try:
|
||||
scores.append(float(item.get("score", 0.0)))
|
||||
except (TypeError, ValueError):
|
||||
scores.append(0.0)
|
||||
return scores
|
||||
|
||||
|
||||
async def _rerank_with_gemini(
|
||||
provider: AIProvider,
|
||||
model: AIModel,
|
||||
query: str,
|
||||
documents: Sequence[str],
|
||||
) -> List[float]:
|
||||
model_name = model.name if model.name.startswith("models/") else f"models/{model.name}"
|
||||
url = _gemini_endpoint(provider, f"{model_name}:rankContent")
|
||||
payload = {
|
||||
"query": {"text": query},
|
||||
"documents": [
|
||||
{"id": str(idx), "content": {"parts": [{"text": content}]}}
|
||||
for idx, content in enumerate(documents)
|
||||
],
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
|
||||
scores: List[float] = []
|
||||
ranked = body.get("rankedDocuments") or body.get("results") or []
|
||||
for item in ranked:
|
||||
raw_score = item.get("relevanceScore") or item.get("score") or item.get("confidenceScore")
|
||||
try:
|
||||
scores.append(float(raw_score))
|
||||
except (TypeError, ValueError):
|
||||
scores.append(0.0)
|
||||
return scores
|
||||
@@ -0,0 +1,501 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import httpx
|
||||
from tortoise.exceptions import DoesNotExist
|
||||
from tortoise.transactions import in_transaction
|
||||
|
||||
from domain.config.service import ConfigService
|
||||
from models.database import AIDefaultModel, AIModel, AIProvider
|
||||
|
||||
from .types import ABILITIES, normalize_capabilities
|
||||
from .vector_providers import (
|
||||
BaseVectorProvider,
|
||||
get_provider_class,
|
||||
get_provider_entry,
|
||||
list_providers,
|
||||
)
|
||||
|
||||
DEFAULT_VECTOR_DIMENSION = 4096
|
||||
|
||||
OPENAI_EMBEDDING_DIMS = {
|
||||
"text-embedding-3-large": 3072,
|
||||
"text-embedding-3-small": 1536,
|
||||
"text-embedding-ada-002": 1536,
|
||||
}
|
||||
|
||||
|
||||
class VectorDBConfigManager:
|
||||
TYPE_KEY = "VECTOR_DB_TYPE"
|
||||
CONFIG_KEY = "VECTOR_DB_CONFIG"
|
||||
DEFAULT_TYPE = "milvus_lite"
|
||||
|
||||
@classmethod
|
||||
async def load_config(cls) -> Tuple[str, Dict[str, Any]]:
|
||||
raw_type = await ConfigService.get(cls.TYPE_KEY, cls.DEFAULT_TYPE)
|
||||
provider_type = str(raw_type or cls.DEFAULT_TYPE)
|
||||
|
||||
raw_config = await ConfigService.get(cls.CONFIG_KEY)
|
||||
config_dict: Dict[str, Any] = {}
|
||||
if isinstance(raw_config, str) and raw_config:
|
||||
try:
|
||||
config_dict = json.loads(raw_config)
|
||||
except json.JSONDecodeError:
|
||||
config_dict = {}
|
||||
elif isinstance(raw_config, dict):
|
||||
config_dict = raw_config
|
||||
return provider_type, config_dict
|
||||
|
||||
@classmethod
|
||||
async def save_config(cls, provider_type: str, config: Dict[str, Any]) -> None:
|
||||
await ConfigService.set(cls.TYPE_KEY, provider_type)
|
||||
await ConfigService.set(cls.CONFIG_KEY, json.dumps(config or {}))
|
||||
|
||||
@classmethod
|
||||
async def get_type(cls) -> str:
|
||||
provider_type, _ = await cls.load_config()
|
||||
return provider_type
|
||||
|
||||
@classmethod
|
||||
async def get_config(cls) -> Dict[str, Any]:
|
||||
_, config = await cls.load_config()
|
||||
return config
|
||||
|
||||
|
||||
def _normalize_embedding_dim(value: Any) -> Optional[int]:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
casted = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return casted if casted > 0 else None
|
||||
|
||||
|
||||
def _apply_embedding_dim_to_metadata(
|
||||
data: Dict[str, Any],
|
||||
embedding_dim: Optional[int],
|
||||
base_metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
source = base_metadata if isinstance(base_metadata, dict) else {}
|
||||
metadata: Dict[str, Any] = dict(source)
|
||||
override = data.get("metadata")
|
||||
if isinstance(override, dict) and override:
|
||||
metadata.update(override)
|
||||
if embedding_dim is None:
|
||||
metadata.pop("embedding_dimensions", None)
|
||||
else:
|
||||
metadata["embedding_dimensions"] = embedding_dim
|
||||
data["metadata"] = metadata or None
|
||||
return data
|
||||
|
||||
|
||||
def infer_openai_capabilities(model_id: str) -> Tuple[List[str], Optional[int]]:
|
||||
lower = model_id.lower()
|
||||
caps = set()
|
||||
|
||||
if any(keyword in lower for keyword in ["gpt", "chat", "turbo", "o1", "sonnet", "haiku", "thinking"]):
|
||||
caps.update({"chat", "tools"})
|
||||
|
||||
if any(keyword in lower for keyword in ["vision", "gpt-4o", "gpt-4.1", "o1", "vision-preview", "omni"]):
|
||||
caps.add("vision")
|
||||
|
||||
if any(keyword in lower for keyword in ["embed", "embedding"]):
|
||||
caps.add("embedding")
|
||||
|
||||
if "rerank" in lower or "re-rank" in lower:
|
||||
caps.add("rerank")
|
||||
|
||||
if any(keyword in lower for keyword in ["tts", "speech", "audio"]):
|
||||
caps.add("voice")
|
||||
|
||||
embedding_dim = OPENAI_EMBEDDING_DIMS.get(model_id)
|
||||
return normalize_capabilities(caps), embedding_dim
|
||||
|
||||
|
||||
def infer_gemini_capabilities(methods: Iterable[str]) -> List[str]:
|
||||
caps = set()
|
||||
for method in methods:
|
||||
m = method.lower()
|
||||
if m in {"generatecontent", "counttokens"}:
|
||||
caps.update({"chat", "tools", "vision"})
|
||||
if m == "embedcontent":
|
||||
caps.add("embedding")
|
||||
if m in {"generatespeech", "audiogeneration"}:
|
||||
caps.add("voice")
|
||||
if m == "rerank":
|
||||
caps.add("rerank")
|
||||
return normalize_capabilities(caps)
|
||||
|
||||
|
||||
def serialize_provider(provider: AIProvider) -> Dict[str, Any]:
|
||||
return {
|
||||
"id": provider.id,
|
||||
"name": provider.name,
|
||||
"identifier": provider.identifier,
|
||||
"provider_type": provider.provider_type,
|
||||
"api_format": provider.api_format,
|
||||
"base_url": provider.base_url,
|
||||
"api_key": provider.api_key,
|
||||
"logo_url": provider.logo_url,
|
||||
"extra_config": provider.extra_config or {},
|
||||
"created_at": provider.created_at,
|
||||
"updated_at": provider.updated_at,
|
||||
}
|
||||
|
||||
|
||||
def model_to_dict(model: AIModel, provider: Optional[AIProvider] = None) -> Dict[str, Any]:
|
||||
provider_obj = provider or getattr(model, "provider", None)
|
||||
provider_data = serialize_provider(provider_obj) if provider_obj else None
|
||||
return {
|
||||
"id": model.id,
|
||||
"provider_id": model.provider_id,
|
||||
"name": model.name,
|
||||
"display_name": model.display_name,
|
||||
"description": model.description,
|
||||
"capabilities": normalize_capabilities(model.capabilities),
|
||||
"context_window": model.context_window,
|
||||
"embedding_dimensions": model.embedding_dimensions,
|
||||
"metadata": model.metadata or {},
|
||||
"created_at": model.created_at,
|
||||
"updated_at": model.updated_at,
|
||||
"provider": provider_data,
|
||||
}
|
||||
|
||||
|
||||
def provider_to_dict(provider: AIProvider, models: Optional[List[AIModel]] = None) -> Dict[str, Any]:
|
||||
data = serialize_provider(provider)
|
||||
if models is not None:
|
||||
data["models"] = [model_to_dict(m, provider=provider) for m in models]
|
||||
return data
|
||||
|
||||
|
||||
class AIProviderService:
|
||||
@classmethod
|
||||
async def list_providers(cls) -> List[Dict[str, Any]]:
|
||||
providers = await AIProvider.all().order_by("id").prefetch_related("models")
|
||||
return [provider_to_dict(p, models=list(p.models)) for p in providers]
|
||||
|
||||
@classmethod
|
||||
async def get_provider(cls, provider_id: int, with_models: bool = False) -> Dict[str, Any]:
|
||||
if with_models:
|
||||
provider = await AIProvider.get(id=provider_id)
|
||||
models = await provider.models.all()
|
||||
return provider_to_dict(provider, models=models)
|
||||
provider = await AIProvider.get(id=provider_id)
|
||||
return provider_to_dict(provider)
|
||||
|
||||
@classmethod
|
||||
async def create_provider(cls, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
data = payload.copy()
|
||||
data.setdefault("extra_config", {})
|
||||
provider = await AIProvider.create(**data)
|
||||
return provider_to_dict(provider)
|
||||
|
||||
@classmethod
|
||||
async def update_provider(cls, provider_id: int, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
provider = await AIProvider.get(id=provider_id)
|
||||
for field, value in payload.items():
|
||||
setattr(provider, field, value)
|
||||
await provider.save()
|
||||
return provider_to_dict(provider)
|
||||
|
||||
@classmethod
|
||||
async def delete_provider(cls, provider_id: int) -> None:
|
||||
await AIProvider.filter(id=provider_id).delete()
|
||||
|
||||
@classmethod
|
||||
async def list_models(cls, provider_id: int) -> List[Dict[str, Any]]:
|
||||
models = await AIModel.filter(provider_id=provider_id).order_by("id").prefetch_related("provider")
|
||||
return [model_to_dict(m) for m in models]
|
||||
|
||||
@classmethod
|
||||
async def create_model(cls, provider_id: int, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
data = payload.copy()
|
||||
data["provider_id"] = provider_id
|
||||
data["capabilities"] = normalize_capabilities(data.get("capabilities"))
|
||||
embedding_dim = _normalize_embedding_dim(data.pop("embedding_dimensions", None))
|
||||
data = _apply_embedding_dim_to_metadata(data, embedding_dim)
|
||||
model = await AIModel.create(**data)
|
||||
await model.fetch_related("provider")
|
||||
return model_to_dict(model)
|
||||
|
||||
@classmethod
|
||||
async def update_model(cls, model_id: int, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
model = await AIModel.get(id=model_id)
|
||||
data = payload.copy()
|
||||
if "capabilities" in data:
|
||||
data["capabilities"] = normalize_capabilities(data.get("capabilities"))
|
||||
embedding_dim = None
|
||||
if "embedding_dimensions" in data:
|
||||
embedding_dim = _normalize_embedding_dim(data.pop("embedding_dimensions", None))
|
||||
_apply_embedding_dim_to_metadata(data, embedding_dim, base_metadata=model.metadata)
|
||||
for field, value in data.items():
|
||||
setattr(model, field, value)
|
||||
if embedding_dim is not None or ("embedding_dimensions" in payload and embedding_dim is None):
|
||||
model.embedding_dimensions = embedding_dim
|
||||
await model.save()
|
||||
await model.fetch_related("provider")
|
||||
return model_to_dict(model)
|
||||
|
||||
@classmethod
|
||||
async def delete_model(cls, model_id: int) -> None:
|
||||
await AIModel.filter(id=model_id).delete()
|
||||
|
||||
@classmethod
|
||||
async def fetch_remote_models(cls, provider_id: int) -> List[Dict[str, Any]]:
|
||||
provider = await AIProvider.get(id=provider_id)
|
||||
return await cls._get_remote_models(provider)
|
||||
|
||||
@classmethod
|
||||
async def _get_remote_models(cls, provider: AIProvider) -> List[Dict[str, Any]]:
|
||||
if not provider.base_url:
|
||||
raise ValueError("Provider base_url is required for syncing models")
|
||||
|
||||
fmt = (provider.api_format or "").lower()
|
||||
if fmt not in {"openai", "gemini"}:
|
||||
raise ValueError(f"Unsupported api_format '{provider.api_format}' for syncing models")
|
||||
|
||||
if fmt == "openai":
|
||||
return await cls._fetch_openai_models(provider)
|
||||
return await cls._fetch_gemini_models(provider)
|
||||
|
||||
@classmethod
|
||||
async def sync_models(cls, provider_id: int) -> Dict[str, int]:
|
||||
provider = await AIProvider.get(id=provider_id)
|
||||
remote_models = await cls._get_remote_models(provider)
|
||||
|
||||
created = 0
|
||||
updated = 0
|
||||
for entry in remote_models:
|
||||
defaults = entry.copy()
|
||||
model_id = defaults.pop("name")
|
||||
defaults["capabilities"] = normalize_capabilities(defaults.get("capabilities"))
|
||||
embedding_dim = _normalize_embedding_dim(defaults.pop("embedding_dimensions", None))
|
||||
defaults = _apply_embedding_dim_to_metadata(defaults, embedding_dim)
|
||||
obj, is_created = await AIModel.get_or_create(
|
||||
provider_id=provider.id,
|
||||
name=model_id,
|
||||
defaults=defaults,
|
||||
)
|
||||
if is_created:
|
||||
created += 1
|
||||
continue
|
||||
for field, value in defaults.items():
|
||||
setattr(obj, field, value)
|
||||
if embedding_dim is not None or ("embedding_dimensions" in entry and embedding_dim is None):
|
||||
obj.embedding_dimensions = embedding_dim
|
||||
await obj.save()
|
||||
updated += 1
|
||||
|
||||
return {"created": created, "updated": updated}
|
||||
|
||||
@classmethod
|
||||
async def get_default_models(cls) -> Dict[str, Optional[Dict[str, Any]]]:
|
||||
defaults = await AIDefaultModel.all().prefetch_related("model__provider")
|
||||
result: Dict[str, Optional[Dict[str, Any]]] = {ability: None for ability in ABILITIES}
|
||||
for item in defaults:
|
||||
result[item.ability] = model_to_dict(item.model, provider=item.model.provider) # type: ignore[attr-defined]
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
async def set_default_models(cls, mapping: Dict[str, Optional[int]]) -> Dict[str, Optional[Dict[str, Any]]]:
|
||||
normalized = {ability: mapping.get(ability) for ability in ABILITIES}
|
||||
async with in_transaction() as connection:
|
||||
for ability, model_id in normalized.items():
|
||||
record = await AIDefaultModel.get_or_none(ability=ability)
|
||||
if model_id:
|
||||
try:
|
||||
model = await AIModel.get(id=model_id)
|
||||
except DoesNotExist:
|
||||
raise ValueError(f"Model {model_id} not found")
|
||||
if record:
|
||||
record.model_id = model_id
|
||||
await record.save(using_db=connection)
|
||||
else:
|
||||
await AIDefaultModel.create(ability=ability, model_id=model_id)
|
||||
elif record:
|
||||
await record.delete(using_db=connection)
|
||||
return await cls.get_default_models()
|
||||
|
||||
@classmethod
|
||||
async def get_default_model(cls, ability: str) -> Optional[AIModel]:
|
||||
ability_key = ability.lower()
|
||||
if ability_key not in ABILITIES:
|
||||
return None
|
||||
record = await AIDefaultModel.get_or_none(ability=ability_key)
|
||||
if not record:
|
||||
return None
|
||||
model = await AIModel.get_or_none(id=record.model_id)
|
||||
if model:
|
||||
await model.fetch_related("provider")
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
async def _fetch_openai_models(cls, provider: AIProvider) -> List[Dict[str, Any]]:
|
||||
base_url = provider.base_url.rstrip("/")
|
||||
url = f"{base_url}/models"
|
||||
headers = {}
|
||||
if provider.api_key:
|
||||
headers["Authorization"] = f"Bearer {provider.api_key}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.get(url, headers=headers)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
|
||||
data = payload.get("data", [])
|
||||
entries: List[Dict[str, Any]] = []
|
||||
for item in data:
|
||||
model_id = item.get("id")
|
||||
if not model_id:
|
||||
continue
|
||||
capabilities, embedding_dim = infer_openai_capabilities(model_id)
|
||||
entries.append({
|
||||
"name": model_id,
|
||||
"display_name": item.get("display_name"),
|
||||
"description": item.get("description"),
|
||||
"capabilities": capabilities,
|
||||
"context_window": item.get("context_window"),
|
||||
"embedding_dimensions": embedding_dim,
|
||||
"metadata": item,
|
||||
})
|
||||
return entries
|
||||
|
||||
@classmethod
|
||||
async def _fetch_gemini_models(cls, provider: AIProvider) -> List[Dict[str, Any]]:
|
||||
base_url = provider.base_url.rstrip("/")
|
||||
suffix = "/models"
|
||||
if provider.api_key:
|
||||
suffix += f"?key={provider.api_key}"
|
||||
url = f"{base_url}{suffix}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
|
||||
data = payload.get("models", [])
|
||||
entries: List[Dict[str, Any]] = []
|
||||
for item in data:
|
||||
model_id = item.get("name")
|
||||
if not model_id:
|
||||
continue
|
||||
methods = item.get("supportedGenerationMethods") or []
|
||||
capabilities = infer_gemini_capabilities(methods)
|
||||
entries.append({
|
||||
"name": model_id,
|
||||
"display_name": item.get("displayName"),
|
||||
"description": item.get("description"),
|
||||
"capabilities": capabilities,
|
||||
"context_window": item.get("inputTokenLimit"),
|
||||
"embedding_dimensions": item.get("embeddingDimensions"),
|
||||
"metadata": item,
|
||||
})
|
||||
return entries
|
||||
|
||||
|
||||
class VectorDBService:
|
||||
_instance: "VectorDBService" | None = None
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not hasattr(self, "_provider"):
|
||||
self._provider: Optional[BaseVectorProvider] = None
|
||||
self._provider_type: Optional[str] = None
|
||||
self._provider_config: Dict[str, Any] | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _ensure_provider(self) -> BaseVectorProvider:
|
||||
if self._provider is None:
|
||||
await self.reload()
|
||||
assert self._provider is not None
|
||||
return self._provider
|
||||
|
||||
async def reload(self) -> BaseVectorProvider:
|
||||
async with self._lock:
|
||||
provider_type, provider_config = await VectorDBConfigManager.load_config()
|
||||
normalized_config = dict(provider_config or {})
|
||||
if (
|
||||
self._provider
|
||||
and self._provider_type == provider_type
|
||||
and self._provider_config == normalized_config
|
||||
):
|
||||
return self._provider
|
||||
|
||||
entry = get_provider_entry(provider_type)
|
||||
if not entry:
|
||||
raise RuntimeError(f"Unknown vector database provider: {provider_type}")
|
||||
if not entry.get("enabled", True):
|
||||
raise RuntimeError(f"Vector database provider '{provider_type}' is disabled")
|
||||
|
||||
provider_cls = get_provider_class(provider_type)
|
||||
if not provider_cls:
|
||||
raise RuntimeError(f"Provider class not found for '{provider_type}'")
|
||||
|
||||
provider = provider_cls(provider_config)
|
||||
await provider.initialize()
|
||||
|
||||
self._provider = provider
|
||||
self._provider_type = provider_type
|
||||
self._provider_config = normalized_config
|
||||
return provider
|
||||
|
||||
async def ensure_collection(self, collection_name: str, vector: bool = True, dim: int = DEFAULT_VECTOR_DIMENSION) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.ensure_collection(collection_name, vector, dim)
|
||||
|
||||
async def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.upsert_vector(collection_name, data)
|
||||
|
||||
async def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.delete_vector(collection_name, path)
|
||||
|
||||
async def search_vectors(self, collection_name: str, query_embedding, top_k: int = 5):
|
||||
provider = await self._ensure_provider()
|
||||
return provider.search_vectors(collection_name, query_embedding, top_k)
|
||||
|
||||
async def search_by_path(self, collection_name: str, query_path: str, top_k: int = 20):
|
||||
provider = await self._ensure_provider()
|
||||
return provider.search_by_path(collection_name, query_path, top_k)
|
||||
|
||||
async def get_all_stats(self) -> Dict[str, Any]:
|
||||
provider = await self._ensure_provider()
|
||||
return provider.get_all_stats()
|
||||
|
||||
async def clear_all_data(self) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.clear_all_data()
|
||||
|
||||
async def current_provider(self) -> Dict[str, Any]:
|
||||
provider_type, provider_config = await VectorDBConfigManager.load_config()
|
||||
entry = get_provider_entry(provider_type) or {}
|
||||
return {
|
||||
"type": provider_type,
|
||||
"config": provider_config,
|
||||
"label": entry.get("label"),
|
||||
"enabled": entry.get("enabled", True),
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AIProviderService",
|
||||
"VectorDBService",
|
||||
"VectorDBConfigManager",
|
||||
"DEFAULT_VECTOR_DIMENSION",
|
||||
"list_providers",
|
||||
"get_provider_entry",
|
||||
"get_provider_class",
|
||||
"normalize_capabilities",
|
||||
"ABILITIES",
|
||||
]
|
||||
@@ -0,0 +1,121 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
ABILITIES = ["chat", "vision", "embedding", "rerank", "voice", "tools"]
|
||||
|
||||
|
||||
def normalize_capabilities(items: Optional[Iterable[str]]) -> List[str]:
|
||||
if not items:
|
||||
return []
|
||||
normalized: List[str] = []
|
||||
for cap in items:
|
||||
key = str(cap).strip().lower()
|
||||
if key in ABILITIES and key not in normalized:
|
||||
normalized.append(key)
|
||||
return normalized
|
||||
|
||||
|
||||
class AIProviderBase(BaseModel):
|
||||
name: str
|
||||
identifier: str = Field(..., pattern=r"^[a-z0-9_\-\.]+$")
|
||||
provider_type: Optional[str] = None
|
||||
api_format: str
|
||||
base_url: Optional[str] = None
|
||||
api_key: Optional[str] = None
|
||||
logo_url: Optional[str] = None
|
||||
extra_config: Optional[dict] = None
|
||||
|
||||
@field_validator("api_format")
|
||||
@classmethod
|
||||
def normalize_format(cls, value: str) -> str:
|
||||
fmt = value.lower()
|
||||
if fmt not in {"openai", "gemini"}:
|
||||
raise ValueError("api_format must be 'openai' or 'gemini'")
|
||||
return fmt
|
||||
|
||||
|
||||
class AIProviderCreate(AIProviderBase):
|
||||
pass
|
||||
|
||||
|
||||
class AIProviderUpdate(BaseModel):
|
||||
name: Optional[str] = None
|
||||
provider_type: Optional[str] = None
|
||||
api_format: Optional[str] = None
|
||||
base_url: Optional[str] = None
|
||||
api_key: Optional[str] = None
|
||||
logo_url: Optional[str] = None
|
||||
extra_config: Optional[dict] = None
|
||||
|
||||
@field_validator("api_format")
|
||||
@classmethod
|
||||
def normalize_format(cls, value: Optional[str]) -> Optional[str]:
|
||||
if value is None:
|
||||
return value
|
||||
fmt = value.lower()
|
||||
if fmt not in {"openai", "gemini"}:
|
||||
raise ValueError("api_format must be 'openai' or 'gemini'")
|
||||
return fmt
|
||||
|
||||
|
||||
class AIModelBase(BaseModel):
|
||||
name: str
|
||||
display_name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
capabilities: Optional[List[str]] = None
|
||||
context_window: Optional[int] = None
|
||||
embedding_dimensions: Optional[int] = None
|
||||
metadata: Optional[dict] = None
|
||||
|
||||
@field_validator("capabilities")
|
||||
@classmethod
|
||||
def validate_capabilities(cls, items: Optional[List[str]]) -> Optional[List[str]]:
|
||||
if items is None:
|
||||
return None
|
||||
normalized = normalize_capabilities(items)
|
||||
invalid = set(items) - set(normalized)
|
||||
if invalid:
|
||||
raise ValueError(f"Unsupported capabilities: {', '.join(invalid)}")
|
||||
return normalized
|
||||
|
||||
|
||||
class AIModelCreate(AIModelBase):
|
||||
pass
|
||||
|
||||
|
||||
class AIModelUpdate(BaseModel):
|
||||
display_name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
capabilities: Optional[List[str]] = None
|
||||
context_window: Optional[int] = None
|
||||
embedding_dimensions: Optional[int] = None
|
||||
metadata: Optional[dict] = None
|
||||
|
||||
@field_validator("capabilities")
|
||||
@classmethod
|
||||
def validate_capabilities(cls, items: Optional[List[str]]) -> Optional[List[str]]:
|
||||
if items is None:
|
||||
return None
|
||||
normalized = normalize_capabilities(items)
|
||||
invalid = set(items) - set(normalized)
|
||||
if invalid:
|
||||
raise ValueError(f"Unsupported capabilities: {', '.join(invalid)}")
|
||||
return normalized
|
||||
|
||||
|
||||
class AIDefaultsUpdate(BaseModel):
|
||||
chat: Optional[int] = None
|
||||
vision: Optional[int] = None
|
||||
embedding: Optional[int] = None
|
||||
rerank: Optional[int] = None
|
||||
voice: Optional[int] = None
|
||||
tools: Optional[int] = None
|
||||
|
||||
def as_mapping(self) -> Dict[str, Optional[int]]:
|
||||
return {ability: getattr(self, ability) for ability in ABILITIES}
|
||||
|
||||
|
||||
class VectorDBConfigPayload(BaseModel):
|
||||
type: str = Field(..., description="向量数据库提供者类型")
|
||||
config: Dict[str, Any] = Field(default_factory=dict, description="提供者配置参数")
|
||||
@@ -0,0 +1,67 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Type
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
from .milvus_lite import MilvusLiteProvider
|
||||
from .milvus_server import MilvusServerProvider
|
||||
from .qdrant import QdrantProvider
|
||||
|
||||
_PROVIDER_REGISTRY: Dict[str, Dict[str, object]] = {
|
||||
MilvusLiteProvider.type: {
|
||||
"class": MilvusLiteProvider,
|
||||
"label": MilvusLiteProvider.label,
|
||||
"description": MilvusLiteProvider.description,
|
||||
"enabled": MilvusLiteProvider.enabled,
|
||||
"config_schema": MilvusLiteProvider.config_schema,
|
||||
},
|
||||
MilvusServerProvider.type: {
|
||||
"class": MilvusServerProvider,
|
||||
"label": MilvusServerProvider.label,
|
||||
"description": MilvusServerProvider.description,
|
||||
"enabled": MilvusServerProvider.enabled,
|
||||
"config_schema": MilvusServerProvider.config_schema,
|
||||
},
|
||||
QdrantProvider.type: {
|
||||
"class": QdrantProvider,
|
||||
"label": QdrantProvider.label,
|
||||
"description": QdrantProvider.description,
|
||||
"enabled": QdrantProvider.enabled,
|
||||
"config_schema": QdrantProvider.config_schema,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def list_providers() -> List[Dict[str, object]]:
|
||||
return [
|
||||
{
|
||||
"type": type_key,
|
||||
"label": meta["label"],
|
||||
"description": meta.get("description"),
|
||||
"enabled": meta.get("enabled", True),
|
||||
"config_schema": meta.get("config_schema", []),
|
||||
}
|
||||
for type_key, meta in _PROVIDER_REGISTRY.items()
|
||||
]
|
||||
|
||||
|
||||
def get_provider_entry(provider_type: str) -> Dict[str, object] | None:
|
||||
return _PROVIDER_REGISTRY.get(provider_type)
|
||||
|
||||
|
||||
def get_provider_class(provider_type: str) -> Type[BaseVectorProvider] | None:
|
||||
entry = get_provider_entry(provider_type)
|
||||
if not entry:
|
||||
return None
|
||||
return entry.get("class") # type: ignore[return-value]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseVectorProvider",
|
||||
"MilvusLiteProvider",
|
||||
"MilvusServerProvider",
|
||||
"QdrantProvider",
|
||||
"list_providers",
|
||||
"get_provider_entry",
|
||||
"get_provider_class",
|
||||
]
|
||||
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
|
||||
class BaseVectorProvider:
|
||||
"""向量数据库提供者基础类，所有实际实现需继承该类"""
|
||||
|
||||
type: str = ""
|
||||
label: str = ""
|
||||
description: str | None = None
|
||||
enabled: bool = True
|
||||
config_schema: List[Dict[str, Any]] = []
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
self.config = config or {}
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""执行初始化逻辑，例如建立连接"""
|
||||
raise NotImplementedError
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
raise NotImplementedError
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,278 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
|
||||
|
||||
class MilvusLiteProvider(BaseVectorProvider):
|
||||
type = "milvus_lite"
|
||||
label = "Milvus Lite"
|
||||
description = "Embedded Milvus Lite (local file storage)."
|
||||
enabled = True
|
||||
config_schema: List[Dict[str, Any]] = [
|
||||
{
|
||||
"key": "db_path",
|
||||
"label": "Database file path",
|
||||
"type": "text",
|
||||
"default": "data/db/milvus.db",
|
||||
"required": False,
|
||||
}
|
||||
]
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
super().__init__(config)
|
||||
self.db_path = Path(self.config.get("db_path") or "data/db/milvus.db")
|
||||
self.client: MilvusClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
try:
|
||||
self.client = MilvusClient(str(self.db_path))
|
||||
except Exception as exc: # pragma: no cover - depends on local environment
|
||||
raise RuntimeError(f"Failed to open Milvus Lite at {self.db_path}: {exc}") from exc
|
||||
|
||||
def _get_client(self) -> MilvusClient:
|
||||
if not self.client:
|
||||
raise RuntimeError("Milvus Lite client is not initialized")
|
||||
return self.client
|
||||
|
||||
@staticmethod
|
||||
def _extract_hit_payload(hit: Any) -> tuple[Any, Any, Dict[str, Any]]:
|
||||
hit_id = getattr(hit, "id", None)
|
||||
distance = getattr(hit, "distance", None)
|
||||
payload: Dict[str, Any] = {}
|
||||
|
||||
raw: Dict[str, Any] | None = None
|
||||
if hasattr(hit, "entity"):
|
||||
raw_entity = getattr(hit, "entity")
|
||||
if hasattr(raw_entity, "to_dict"):
|
||||
raw = dict(raw_entity.to_dict())
|
||||
else:
|
||||
raw = dict(raw_entity)
|
||||
elif isinstance(hit, dict):
|
||||
raw = dict(hit)
|
||||
|
||||
if raw:
|
||||
hit_id = hit_id or raw.get("id")
|
||||
distance = distance if distance is not None else raw.get("distance")
|
||||
inner = raw.get("entity")
|
||||
if isinstance(inner, dict):
|
||||
payload = dict(inner)
|
||||
else:
|
||||
payload = {k: v for k, v in raw.items() if k not in {"id", "distance", "entity"}}
|
||||
|
||||
payload.setdefault("path", payload.get("source_path"))
|
||||
payload.setdefault("source_path", payload.get("path"))
|
||||
return hit_id, distance, payload
|
||||
|
||||
@staticmethod
|
||||
def _to_int(value: Any) -> int:
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
client = self._get_client()
|
||||
if client.has_collection(collection_name):
|
||||
return
|
||||
common_fields = [
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=512, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="source_path", dtype=DataType.VARCHAR, max_length=512, is_primary=False, auto_id=False),
|
||||
]
|
||||
|
||||
if vector:
|
||||
vector_dim = dim if isinstance(dim, int) and dim > 0 else 0
|
||||
if vector_dim <= 0:
|
||||
vector_dim = 4096
|
||||
fields = [
|
||||
*common_fields,
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
|
||||
]
|
||||
schema = CollectionSchema(fields, description="Vector collection", enable_dynamic_field=True)
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
index_params = MilvusClient.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="embedding",
|
||||
index_type="IVF_FLAT",
|
||||
index_name="vector_index",
|
||||
metric_type="COSINE",
|
||||
params={"nlist": 64},
|
||||
)
|
||||
client.create_index(collection_name, index_params=index_params)
|
||||
else:
|
||||
schema = CollectionSchema(common_fields, description="Simple file index", enable_dynamic_field=True)
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
payload = dict(data)
|
||||
payload.setdefault("source_path", payload.get("path"))
|
||||
payload.setdefault("vector_id", payload.get("path"))
|
||||
self._get_client().upsert(collection_name, data=[payload])
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
client = self._get_client()
|
||||
escaped = path.replace('"', '\\"')
|
||||
client.delete(collection_name, filter=f'source_path == "{escaped}"')
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
search_params = {"metric_type": "COSINE"}
|
||||
output_fields = [
|
||||
"path",
|
||||
"source_path",
|
||||
"chunk_id",
|
||||
"mime",
|
||||
"text",
|
||||
"start_offset",
|
||||
"end_offset",
|
||||
"type",
|
||||
"name",
|
||||
]
|
||||
raw_results = self._get_client().search(
|
||||
collection_name,
|
||||
data=[query_embedding],
|
||||
anns_field="embedding",
|
||||
search_params=search_params,
|
||||
limit=top_k,
|
||||
output_fields=output_fields,
|
||||
)
|
||||
formatted: List[List[Dict[str, Any]]] = []
|
||||
for hits in raw_results:
|
||||
bucket: List[Dict[str, Any]] = []
|
||||
for hit in hits:
|
||||
hit_id, distance, entity = self._extract_hit_payload(hit)
|
||||
bucket.append({
|
||||
"id": hit_id,
|
||||
"distance": distance,
|
||||
"entity": entity,
|
||||
})
|
||||
formatted.append(bucket)
|
||||
return formatted
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
if query_path:
|
||||
escaped = query_path.replace('"', '\\"')
|
||||
filter_expr = f'source_path like \"%{escaped}%\"'
|
||||
else:
|
||||
filter_expr = "source_path like '%%'"
|
||||
results = self._get_client().query(
|
||||
collection_name,
|
||||
filter=filter_expr,
|
||||
limit=top_k,
|
||||
output_fields=[
|
||||
"path",
|
||||
"source_path",
|
||||
"chunk_id",
|
||||
"mime",
|
||||
"text",
|
||||
"start_offset",
|
||||
"end_offset",
|
||||
"type",
|
||||
"name",
|
||||
],
|
||||
)
|
||||
formatted = []
|
||||
for row in results:
|
||||
entity = dict(row)
|
||||
entity.setdefault("path", entity.get("source_path"))
|
||||
formatted.append({
|
||||
"id": entity.get("path"),
|
||||
"distance": 1.0,
|
||||
"entity": entity,
|
||||
})
|
||||
return [formatted]
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
client = self._get_client()
|
||||
try:
|
||||
collection_names = client.list_collections()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to list collections: {exc}") from exc
|
||||
|
||||
collections: List[Dict[str, Any]] = []
|
||||
total_vectors = 0
|
||||
total_estimated_memory = 0
|
||||
|
||||
for name in collection_names:
|
||||
try:
|
||||
stats = client.get_collection_stats(name) or {}
|
||||
except Exception:
|
||||
stats = {}
|
||||
row_count = self._to_int(stats.get("row_count"))
|
||||
total_vectors += row_count
|
||||
|
||||
dimension: Optional[int] = None
|
||||
is_vector_collection = False
|
||||
try:
|
||||
description = client.describe_collection(name)
|
||||
except Exception:
|
||||
description = None
|
||||
|
||||
if description:
|
||||
for field in description.get("fields", []):
|
||||
if field.get("type") == DataType.FLOAT_VECTOR:
|
||||
params = field.get("params") or {}
|
||||
dimension = self._to_int(params.get("dim")) or 4096
|
||||
is_vector_collection = True
|
||||
break
|
||||
|
||||
estimated_memory = 0
|
||||
if is_vector_collection and dimension:
|
||||
estimated_memory = row_count * dimension * 4
|
||||
total_estimated_memory += estimated_memory
|
||||
|
||||
indexes: List[Dict[str, Any]] = []
|
||||
try:
|
||||
index_names = client.list_indexes(name) or []
|
||||
except Exception:
|
||||
index_names = []
|
||||
|
||||
for index_name in index_names:
|
||||
try:
|
||||
detail = client.describe_index(name) or {}
|
||||
except Exception:
|
||||
detail = {}
|
||||
indexes.append(
|
||||
{
|
||||
"index_name": index_name,
|
||||
"index_type": detail.get("index_type"),
|
||||
"metric_type": detail.get("metric_type"),
|
||||
"indexed_rows": self._to_int(detail.get("indexed_rows")),
|
||||
"pending_index_rows": self._to_int(detail.get("pending_index_rows")),
|
||||
"state": detail.get("state"),
|
||||
}
|
||||
)
|
||||
|
||||
collections.append(
|
||||
{
|
||||
"name": name,
|
||||
"row_count": row_count,
|
||||
"dimension": dimension if is_vector_collection else None,
|
||||
"estimated_memory_bytes": estimated_memory,
|
||||
"is_vector_collection": is_vector_collection,
|
||||
"indexes": indexes,
|
||||
}
|
||||
)
|
||||
|
||||
db_file_size = None
|
||||
try:
|
||||
if self.db_path.exists():
|
||||
db_file_size = self.db_path.stat().st_size
|
||||
except OSError:
|
||||
db_file_size = None
|
||||
|
||||
return {
|
||||
"collections": collections,
|
||||
"collection_count": len(collections),
|
||||
"total_vectors": total_vectors,
|
||||
"estimated_total_memory_bytes": total_estimated_memory,
|
||||
"db_file_size_bytes": db_file_size,
|
||||
}
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
client = self._get_client()
|
||||
for collection_name in client.list_collections():
|
||||
client.drop_collection(collection_name)
|
||||
@@ -0,0 +1,278 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
|
||||
|
||||
class MilvusServerProvider(BaseVectorProvider):
|
||||
type = "milvus_server"
|
||||
label = "Milvus Server"
|
||||
description = "Remote Milvus instance accessed via URI."
|
||||
enabled = True
|
||||
config_schema: List[Dict[str, Any]] = [
|
||||
{
|
||||
"key": "uri",
|
||||
"label": "Server URI",
|
||||
"type": "text",
|
||||
"required": True,
|
||||
"placeholder": "http://localhost:19530",
|
||||
},
|
||||
{
|
||||
"key": "token",
|
||||
"label": "Token",
|
||||
"type": "password",
|
||||
"required": False,
|
||||
"placeholder": "user:password",
|
||||
},
|
||||
]
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
super().__init__(config)
|
||||
self.client: MilvusClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
uri = self.config.get("uri")
|
||||
if not uri:
|
||||
raise RuntimeError("Milvus Server URI is required")
|
||||
try:
|
||||
self.client = MilvusClient(uri=uri, token=self.config.get("token"))
|
||||
except Exception as exc: # pragma: no cover - depends on remote availability
|
||||
raise RuntimeError(f"Failed to connect to Milvus Server {uri}: {exc}") from exc
|
||||
|
||||
def _get_client(self) -> MilvusClient:
|
||||
if not self.client:
|
||||
raise RuntimeError("Milvus Server client is not initialized")
|
||||
return self.client
|
||||
|
||||
@staticmethod
|
||||
def _extract_hit_payload(hit: Any) -> tuple[Any, Any, Dict[str, Any]]:
|
||||
hit_id = getattr(hit, "id", None)
|
||||
distance = getattr(hit, "distance", None)
|
||||
payload: Dict[str, Any] = {}
|
||||
|
||||
raw: Dict[str, Any] | None = None
|
||||
if hasattr(hit, "entity"):
|
||||
raw_entity = getattr(hit, "entity")
|
||||
if hasattr(raw_entity, "to_dict"):
|
||||
raw = dict(raw_entity.to_dict())
|
||||
else:
|
||||
raw = dict(raw_entity)
|
||||
elif isinstance(hit, dict):
|
||||
raw = dict(hit)
|
||||
|
||||
if raw:
|
||||
hit_id = hit_id or raw.get("id")
|
||||
distance = distance if distance is not None else raw.get("distance")
|
||||
inner = raw.get("entity")
|
||||
if isinstance(inner, dict):
|
||||
payload = dict(inner)
|
||||
else:
|
||||
payload = {k: v for k, v in raw.items() if k not in {"id", "distance", "entity"}}
|
||||
|
||||
payload.setdefault("path", payload.get("source_path"))
|
||||
payload.setdefault("source_path", payload.get("path"))
|
||||
return hit_id, distance, payload
|
||||
|
||||
@staticmethod
|
||||
def _to_int(value: Any) -> int:
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
client = self._get_client()
|
||||
if client.has_collection(collection_name):
|
||||
return
|
||||
common_fields = [
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=512, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="source_path", dtype=DataType.VARCHAR, max_length=512, is_primary=False, auto_id=False),
|
||||
]
|
||||
if vector:
|
||||
vector_dim = dim if isinstance(dim, int) and dim > 0 else 0
|
||||
if vector_dim <= 0:
|
||||
vector_dim = 4096
|
||||
fields = [
|
||||
*common_fields,
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
|
||||
]
|
||||
schema = CollectionSchema(fields, description="Vector collection", enable_dynamic_field=True)
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
index_params = MilvusClient.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="embedding",
|
||||
index_type="IVF_FLAT",
|
||||
index_name="vector_index",
|
||||
metric_type="COSINE",
|
||||
params={"nlist": 64},
|
||||
)
|
||||
client.create_index(collection_name, index_params=index_params)
|
||||
else:
|
||||
schema = CollectionSchema(common_fields, description="Simple file index", enable_dynamic_field=True)
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
payload = dict(data)
|
||||
payload.setdefault("source_path", payload.get("path"))
|
||||
payload.setdefault("vector_id", payload.get("path"))
|
||||
self._get_client().upsert(collection_name, data=[payload])
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
client = self._get_client()
|
||||
escaped = path.replace('"', '\\"')
|
||||
client.delete(collection_name, filter=f'source_path == "{escaped}"')
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
search_params = {"metric_type": "COSINE"}
|
||||
output_fields = [
|
||||
"path",
|
||||
"source_path",
|
||||
"chunk_id",
|
||||
"mime",
|
||||
"text",
|
||||
"start_offset",
|
||||
"end_offset",
|
||||
"type",
|
||||
"name",
|
||||
]
|
||||
raw_results = self._get_client().search(
|
||||
collection_name,
|
||||
data=[query_embedding],
|
||||
anns_field="embedding",
|
||||
search_params=search_params,
|
||||
limit=top_k,
|
||||
output_fields=output_fields,
|
||||
)
|
||||
formatted: List[List[Dict[str, Any]]] = []
|
||||
for hits in raw_results:
|
||||
bucket: List[Dict[str, Any]] = []
|
||||
for hit in hits:
|
||||
hit_id, distance, entity = self._extract_hit_payload(hit)
|
||||
bucket.append({
|
||||
"id": hit_id,
|
||||
"distance": distance,
|
||||
"entity": entity,
|
||||
})
|
||||
formatted.append(bucket)
|
||||
return formatted
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
if query_path:
|
||||
escaped = query_path.replace('"', '\\"')
|
||||
filter_expr = f'source_path like \"%{escaped}%\"'
|
||||
else:
|
||||
filter_expr = "source_path like '%%'"
|
||||
results = self._get_client().query(
|
||||
collection_name,
|
||||
filter=filter_expr,
|
||||
limit=top_k,
|
||||
output_fields=[
|
||||
"path",
|
||||
"source_path",
|
||||
"chunk_id",
|
||||
"mime",
|
||||
"text",
|
||||
"start_offset",
|
||||
"end_offset",
|
||||
"type",
|
||||
"name",
|
||||
],
|
||||
)
|
||||
formatted = []
|
||||
for row in results:
|
||||
entity = dict(row)
|
||||
entity.setdefault("path", entity.get("source_path"))
|
||||
formatted.append({
|
||||
"id": entity.get("path"),
|
||||
"distance": 1.0,
|
||||
"entity": entity,
|
||||
})
|
||||
return [formatted]
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
client = self._get_client()
|
||||
try:
|
||||
collection_names = client.list_collections()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to list collections: {exc}") from exc
|
||||
|
||||
collections: List[Dict[str, Any]] = []
|
||||
total_vectors = 0
|
||||
total_estimated_memory = 0
|
||||
|
||||
for name in collection_names:
|
||||
try:
|
||||
stats = client.get_collection_stats(name) or {}
|
||||
except Exception:
|
||||
stats = {}
|
||||
row_count = self._to_int(stats.get("row_count"))
|
||||
total_vectors += row_count
|
||||
|
||||
dimension: Optional[int] = None
|
||||
is_vector_collection = False
|
||||
try:
|
||||
description = client.describe_collection(name)
|
||||
except Exception:
|
||||
description = None
|
||||
|
||||
if description:
|
||||
for field in description.get("fields", []):
|
||||
if field.get("type") == DataType.FLOAT_VECTOR:
|
||||
params = field.get("params") or {}
|
||||
dimension = self._to_int(params.get("dim")) or 4096
|
||||
is_vector_collection = True
|
||||
break
|
||||
|
||||
estimated_memory = 0
|
||||
if is_vector_collection and dimension:
|
||||
estimated_memory = row_count * dimension * 4
|
||||
total_estimated_memory += estimated_memory
|
||||
|
||||
indexes: List[Dict[str, Any]] = []
|
||||
try:
|
||||
index_names = client.list_indexes(name) or []
|
||||
except Exception:
|
||||
index_names = []
|
||||
|
||||
for index_name in index_names:
|
||||
try:
|
||||
detail = client.describe_index(name) or {}
|
||||
except Exception:
|
||||
detail = {}
|
||||
indexes.append(
|
||||
{
|
||||
"index_name": index_name,
|
||||
"index_type": detail.get("index_type"),
|
||||
"metric_type": detail.get("metric_type"),
|
||||
"indexed_rows": self._to_int(detail.get("indexed_rows")),
|
||||
"pending_index_rows": self._to_int(detail.get("pending_index_rows")),
|
||||
"state": detail.get("state"),
|
||||
}
|
||||
)
|
||||
|
||||
collections.append(
|
||||
{
|
||||
"name": name,
|
||||
"row_count": row_count,
|
||||
"dimension": dimension if is_vector_collection else None,
|
||||
"estimated_memory_bytes": estimated_memory,
|
||||
"is_vector_collection": is_vector_collection,
|
||||
"indexes": indexes,
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"collections": collections,
|
||||
"collection_count": len(collections),
|
||||
"total_vectors": total_vectors,
|
||||
"estimated_total_memory_bytes": total_estimated_memory,
|
||||
"db_file_size_bytes": None,
|
||||
}
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
client = self._get_client()
|
||||
for collection_name in client.list_collections():
|
||||
client.drop_collection(collection_name)
|
||||
@@ -0,0 +1,273 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qmodels
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
|
||||
|
||||
class QdrantProvider(BaseVectorProvider):
|
||||
type = "qdrant"
|
||||
label = "Qdrant"
|
||||
description = "Qdrant vector database (HTTP API)."
|
||||
enabled = True
|
||||
config_schema: List[Dict[str, Any]] = [
|
||||
{
|
||||
"key": "url",
|
||||
"label": "Server URL",
|
||||
"type": "text",
|
||||
"required": True,
|
||||
"placeholder": "http://localhost:6333",
|
||||
},
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "API Key",
|
||||
"type": "password",
|
||||
"required": False,
|
||||
},
|
||||
]
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
super().__init__(config)
|
||||
self.client: Optional[QdrantClient] = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
url = (self.config.get("url") or "").strip()
|
||||
if not url:
|
||||
raise RuntimeError("Qdrant URL is required")
|
||||
|
||||
api_key = (self.config.get("api_key") or None) or None
|
||||
try:
|
||||
client = QdrantClient(url=url, api_key=api_key)
|
||||
client.get_collections()
|
||||
self.client = client
|
||||
except Exception as exc: # pragma: no cover - 依赖外部服务
|
||||
raise RuntimeError(f"Failed to connect to Qdrant at {url}: {exc}") from exc
|
||||
|
||||
def _get_client(self) -> QdrantClient:
|
||||
if not self.client:
|
||||
raise RuntimeError("Qdrant client is not initialized")
|
||||
return self.client
|
||||
|
||||
@staticmethod
|
||||
def _vector_params(vector: bool, dim: int) -> qmodels.VectorParams:
|
||||
size = dim if vector and isinstance(dim, int) and dim > 0 else 1
|
||||
return qmodels.VectorParams(size=size, distance=qmodels.Distance.COSINE)
|
||||
|
||||
def _ensure_payload_indexes(self, client: QdrantClient, collection_name: str) -> None:
|
||||
for field in ("path", "source_path"):
|
||||
try:
|
||||
client.create_payload_index(
|
||||
collection_name=collection_name,
|
||||
field_name=field,
|
||||
field_schema="keyword",
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - 依赖外部服务
|
||||
message = str(exc).lower()
|
||||
if "already exists" in message or "index exists" in message:
|
||||
continue
|
||||
raise
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
client = self._get_client()
|
||||
try:
|
||||
exists = client.collection_exists(collection_name)
|
||||
except Exception as exc: # pragma: no cover - 依赖外部服务
|
||||
raise RuntimeError(f"Failed to check Qdrant collection '{collection_name}': {exc}") from exc
|
||||
|
||||
if exists:
|
||||
try:
|
||||
self._ensure_payload_indexes(client, collection_name)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
vectors_config = self._vector_params(vector, dim)
|
||||
try:
|
||||
client.create_collection(collection_name=collection_name, vectors_config=vectors_config)
|
||||
except Exception as exc: # pragma: no cover
|
||||
if "already exists" in str(exc).lower():
|
||||
try:
|
||||
self._ensure_payload_indexes(client, collection_name)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
raise RuntimeError(f"Failed to create Qdrant collection '{collection_name}': {exc}") from exc
|
||||
|
||||
try:
|
||||
self._ensure_payload_indexes(client, collection_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _point_id(uid: str) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, uid))
|
||||
|
||||
def _prepare_point(self, data: Dict[str, Any]) -> qmodels.PointStruct:
|
||||
uid = data.get("path")
|
||||
if not uid:
|
||||
raise ValueError("Qdrant upsert requires 'path' in data")
|
||||
|
||||
embedding = data.get("embedding")
|
||||
if embedding is None:
|
||||
vector = [0.0]
|
||||
else:
|
||||
vector = [float(x) for x in embedding]
|
||||
|
||||
payload = {k: v for k, v in data.items() if k != "embedding"}
|
||||
payload.setdefault("vector_id", uid)
|
||||
source_path = payload.get("source_path") or payload.get("path")
|
||||
payload["path"] = source_path
|
||||
return qmodels.PointStruct(id=self._point_id(str(uid)), vector=vector, payload=payload)
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
client = self._get_client()
|
||||
point = self._prepare_point(data)
|
||||
client.upsert(collection_name=collection_name, wait=True, points=[point])
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
client = self._get_client()
|
||||
condition = qmodels.FieldCondition(
|
||||
key="path",
|
||||
match=qmodels.MatchValue(value=path),
|
||||
)
|
||||
flt = qmodels.Filter(must=[condition])
|
||||
selector = qmodels.FilterSelector(filter=flt)
|
||||
client.delete(collection_name=collection_name, points_selector=selector, wait=True)
|
||||
|
||||
def _format_search_results(self, points: Sequence[qmodels.ScoredPoint]):
|
||||
return [
|
||||
{
|
||||
"id": point.id,
|
||||
"distance": point.score,
|
||||
"entity": point.payload or {},
|
||||
}
|
||||
for point in points
|
||||
]
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
client = self._get_client()
|
||||
vector = [float(x) for x in query_embedding]
|
||||
points = client.search(
|
||||
collection_name=collection_name,
|
||||
query_vector=vector,
|
||||
limit=top_k,
|
||||
with_payload=True,
|
||||
)
|
||||
return [self._format_search_results(points)]
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
client = self._get_client()
|
||||
results: List[Dict[str, Any]] = []
|
||||
offset: Optional[str | int] = None
|
||||
remaining = max(top_k, 1)
|
||||
|
||||
while len(results) < top_k:
|
||||
batch_size = min(max(remaining * 2, 10), 200)
|
||||
records, next_offset = client.scroll(
|
||||
collection_name=collection_name,
|
||||
limit=batch_size,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
)
|
||||
if not records:
|
||||
break
|
||||
|
||||
for record in records:
|
||||
payload = record.payload or {}
|
||||
path = payload.get("path")
|
||||
if query_path and path and query_path not in path:
|
||||
continue
|
||||
results.append({"id": record.id, "distance": 1.0, "entity": payload})
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
|
||||
if next_offset is None or len(results) >= top_k:
|
||||
break
|
||||
offset = next_offset
|
||||
remaining = top_k - len(results)
|
||||
|
||||
return [results]
|
||||
|
||||
def _extract_vector_config(self, vectors) -> Optional[qmodels.VectorParams]:
|
||||
if isinstance(vectors, qmodels.VectorParams):
|
||||
return vectors
|
||||
if isinstance(vectors, dict):
|
||||
for value in vectors.values():
|
||||
if isinstance(value, qmodels.VectorParams):
|
||||
return value
|
||||
return None
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
client = self._get_client()
|
||||
try:
|
||||
response = client.get_collections()
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError(f"Failed to list Qdrant collections: {exc}") from exc
|
||||
|
||||
collections: List[Dict[str, Any]] = []
|
||||
total_vectors = 0
|
||||
total_estimated_memory = 0
|
||||
|
||||
for description in response.collections or []:
|
||||
name = description.name
|
||||
try:
|
||||
info = client.get_collection(name)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
row_count = int(info.points_count or 0)
|
||||
total_vectors += row_count
|
||||
|
||||
vector_params = self._extract_vector_config(info.config.params.vectors if info.config and info.config.params else None)
|
||||
dimension = int(vector_params.size) if vector_params and vector_params.size else None
|
||||
estimated_memory = row_count * dimension * 4 if dimension else 0
|
||||
total_estimated_memory += estimated_memory
|
||||
distance = str(vector_params.distance) if vector_params and vector_params.distance else None
|
||||
|
||||
indexed_rows = int(info.indexed_vectors_count or 0)
|
||||
pending_rows = max(row_count - indexed_rows, 0)
|
||||
|
||||
collections.append(
|
||||
{
|
||||
"name": name,
|
||||
"row_count": row_count,
|
||||
"dimension": dimension,
|
||||
"estimated_memory_bytes": estimated_memory,
|
||||
"is_vector_collection": dimension is not None and dimension > 1,
|
||||
"indexes": [
|
||||
{
|
||||
"index_name": "hnsw",
|
||||
"index_type": "HNSW",
|
||||
"metric_type": distance,
|
||||
"indexed_rows": indexed_rows,
|
||||
"pending_index_rows": pending_rows,
|
||||
"state": info.status,
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"collections": collections,
|
||||
"collection_count": len(collections),
|
||||
"total_vectors": total_vectors,
|
||||
"estimated_total_memory_bytes": total_estimated_memory,
|
||||
"db_file_size_bytes": None,
|
||||
}
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
client = self._get_client()
|
||||
try:
|
||||
response = client.get_collections()
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError(f"Failed to list Qdrant collections: {exc}") from exc
|
||||
|
||||
for description in response.collections or []:
|
||||
try:
|
||||
client.delete_collection(description.name)
|
||||
except Exception:
|
||||
continue
|
||||
Reference in New Issue
Block a user