111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
"""Embedding 模型选择、批量调用和受控测试回退。"""
|
|
|
|
import hashlib
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
from django.conf import settings
|
|
from openai import OpenAI
|
|
|
|
from apps.modelhub.crypto import decrypt_secret
|
|
from apps.modelhub.models import ModelDeployment
|
|
|
|
|
|
class KnowledgeEmbeddingError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EmbeddingBatch:
|
|
vectors: list[list[float]]
|
|
model_name: str
|
|
deployment_id: int | None
|
|
fallback: bool
|
|
|
|
|
|
def hash_embedding(content: str, dimensions: int | None = None) -> list[float]:
|
|
size = dimensions or settings.KNOWLEDGE_EMBEDDING_DIMENSIONS
|
|
values = [0.0] * size
|
|
encoded = content.encode("utf-8")
|
|
for index in range(size):
|
|
digest = hashlib.sha256(encoded + index.to_bytes(4, "big")).digest()
|
|
values[index] = (int.from_bytes(digest[:4], "big") / 2**32) * 2 - 1
|
|
length = math.sqrt(sum(value * value for value in values)) or 1.0
|
|
return [value / length for value in values]
|
|
|
|
|
|
def get_embedding_deployment() -> ModelDeployment | None:
|
|
queryset = ModelDeployment.objects.select_related("provider").filter(
|
|
capability="embedding",
|
|
enabled=True,
|
|
provider__enabled=True,
|
|
archived_at__isnull=True,
|
|
provider__archived_at__isnull=True,
|
|
)
|
|
configured_id = settings.KNOWLEDGE_EMBEDDING_DEPLOYMENT_ID
|
|
if configured_id:
|
|
return queryset.filter(id=configured_id).first()
|
|
return queryset.order_by("id").first()
|
|
|
|
|
|
def _request_embeddings(
|
|
deployment: ModelDeployment,
|
|
contents: list[str],
|
|
dimensions: int,
|
|
) -> list[list[float]]:
|
|
provider = deployment.provider
|
|
client = OpenAI(
|
|
api_key=decrypt_secret(provider.api_key_encrypted),
|
|
base_url=provider.base_url,
|
|
timeout=min(max(provider.default_timeout_seconds, 1), 300),
|
|
max_retries=0,
|
|
)
|
|
response = client.embeddings.create(
|
|
model=deployment.model_name,
|
|
input=contents,
|
|
dimensions=dimensions,
|
|
encoding_format="float",
|
|
)
|
|
return [list(item.embedding) for item in sorted(response.data, key=lambda item: item.index)]
|
|
|
|
|
|
def embed_contents(
|
|
contents: list[str],
|
|
*,
|
|
allow_fallback: bool | None = None,
|
|
) -> EmbeddingBatch:
|
|
"""批量生成固定维度向量；哈希回退必须由配置或调用方显式允许。"""
|
|
|
|
if not contents:
|
|
return EmbeddingBatch([], "", None, False)
|
|
dimensions = settings.KNOWLEDGE_EMBEDDING_DIMENSIONS
|
|
deployment = get_embedding_deployment()
|
|
fallback_allowed = (
|
|
settings.KNOWLEDGE_EMBEDDING_ALLOW_HASH_FALLBACK
|
|
if allow_fallback is None
|
|
else allow_fallback
|
|
)
|
|
if deployment is None:
|
|
if not fallback_allowed:
|
|
raise KnowledgeEmbeddingError("未配置可用的 Embedding 模型部署")
|
|
return EmbeddingBatch(
|
|
[hash_embedding(content, dimensions) for content in contents],
|
|
f"local-hash-v2-{dimensions}",
|
|
None,
|
|
True,
|
|
)
|
|
try:
|
|
vectors = _request_embeddings(deployment, contents, dimensions)
|
|
except Exception as error:
|
|
if not fallback_allowed:
|
|
raise KnowledgeEmbeddingError("Embedding 模型调用失败") from error
|
|
return EmbeddingBatch(
|
|
[hash_embedding(content, dimensions) for content in contents],
|
|
f"local-hash-v2-{dimensions}",
|
|
None,
|
|
True,
|
|
)
|
|
if len(vectors) != len(contents) or any(len(vector) != dimensions for vector in vectors):
|
|
raise KnowledgeEmbeddingError("Embedding 模型返回的向量数量或维度不符合配置")
|
|
return EmbeddingBatch(vectors, deployment.model_name, deployment.id, False)
|