150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
"""嵌入服务 — Embedding 适配器外观。
|
|
|
|
用法::
|
|
|
|
from app.ai_provider.embedding import EmbeddingService
|
|
|
|
service = EmbeddingService(manager)
|
|
result = await service.embed("要嵌入的文本")
|
|
results = await service.embed(["文本1", "文本2"])
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from app.ai_provider.embedding.base import (
|
|
EmbeddingProtocol,
|
|
EmbeddingRequest,
|
|
EmbeddingResponse,
|
|
EmbeddingResult,
|
|
EmbeddingServiceError,
|
|
)
|
|
from app.ai_provider.types.errors import AIProviderError
|
|
|
|
|
|
class EmbeddingService:
|
|
"""外观类，负责解析 Embedding 模型并委托给 ProviderManager。
|
|
|
|
接受 ProviderManager 或实现了以下接口的适配器：
|
|
- ``.model_registry`` → 具有 ``list_by_capability(capability)`` 方法
|
|
- ``.create_embedding(name)`` → 返回 EmbeddingProtocol 实例
|
|
"""
|
|
|
|
def __init__(self, manager) -> None:
|
|
self._manager = manager
|
|
self._adapter: EmbeddingProtocol | None = None
|
|
|
|
# ── Registry 兼容适配 ─────────────────────────────────────────
|
|
|
|
@property
|
|
def _registry(self):
|
|
mgr = self._manager
|
|
if hasattr(mgr, "model_registry"):
|
|
return mgr.model_registry
|
|
return mgr
|
|
|
|
@property
|
|
def _factory(self):
|
|
mgr = self._manager
|
|
if hasattr(mgr, "create_embedding"):
|
|
return mgr
|
|
return None
|
|
|
|
# ── Embedding ──────────────────────────────────────────────────
|
|
|
|
@property
|
|
def adapter(self) -> EmbeddingProtocol:
|
|
"""解析并缓存 Embedding 适配器。"""
|
|
if self._adapter is None:
|
|
self._adapter = self._resolve()
|
|
return self._adapter
|
|
|
|
async def embed(
|
|
self,
|
|
input: str | list[str],
|
|
model_hint: str | None = None,
|
|
dimensions: int | None = None,
|
|
) -> EmbeddingResponse:
|
|
"""快捷方法：构造请求并调用 Embedding 适配器。"""
|
|
model_name = self._resolve_model_name(model_hint)
|
|
# 优先从模型配置读 dimensions
|
|
effective_dimensions = dimensions
|
|
if effective_dimensions is None:
|
|
effective_dimensions = self._resolve_model_dimensions(model_hint)
|
|
request = EmbeddingRequest(
|
|
input=input if isinstance(input, list) else [input],
|
|
model=model_name,
|
|
dimensions=effective_dimensions,
|
|
)
|
|
return await self.adapter.embed(request)
|
|
|
|
# ── 解析 ───────────────────────────────────────────────────────
|
|
|
|
def _resolve(self) -> EmbeddingProtocol:
|
|
"""查找第一个活跃的 Embedding 模型并通过工厂创建。"""
|
|
models = self._list_capability_models("embedding")
|
|
if not models:
|
|
raise AIProviderError(
|
|
status_code=503,
|
|
code="embedding_no_model",
|
|
message="No embedding model configured in the model registry.",
|
|
)
|
|
factory = self._factory
|
|
if factory is None:
|
|
raise AIProviderError(
|
|
status_code=501,
|
|
code="embedding_no_factory",
|
|
message="Embedding adapter factory not available (provider manager missing create_embedding).",
|
|
)
|
|
return factory.create_embedding(models[0].name)
|
|
|
|
def _resolve_model_name(self, model_hint: str | None) -> str:
|
|
models = self._list_capability_models("embedding")
|
|
if model_hint:
|
|
return model_hint
|
|
if models:
|
|
return models[0].name
|
|
return ""
|
|
|
|
def _resolve_model_dimensions(self, model_hint: str | None) -> int | None:
|
|
"""从模型配置读取向量维度。"""
|
|
models = self._list_capability_models("embedding")
|
|
if model_hint:
|
|
for m in models:
|
|
if m.name == model_hint:
|
|
return m.dimensions
|
|
if models:
|
|
return models[0].dimensions
|
|
return None
|
|
|
|
def _list_capability_models(self, tag: str):
|
|
reg = self._registry
|
|
if hasattr(reg, "list_capability_models"):
|
|
return reg.list_capability_models(tag)
|
|
if hasattr(reg, "list_by_capability"):
|
|
return reg.list_by_capability(tag)
|
|
return []
|
|
|
|
# ── 生命周期 ────────────────────────────────────────────────────
|
|
|
|
async def close(self) -> None:
|
|
if self._adapter is not None:
|
|
close = getattr(self._adapter, "close", None)
|
|
if close is not None:
|
|
await close()
|
|
|
|
async def __aenter__(self) -> EmbeddingService:
|
|
return self
|
|
|
|
async def __aexit__(self, *exc_info: object) -> None:
|
|
await self.close()
|
|
|
|
|
|
__all__ = [
|
|
"EmbeddingProtocol",
|
|
"EmbeddingRequest",
|
|
"EmbeddingResponse",
|
|
"EmbeddingResult",
|
|
"EmbeddingService",
|
|
"EmbeddingServiceError",
|
|
]
|