feat(flywheel): 抽公共 EmbeddingProvider 并新增 FewShotSample 模型
- 从 knowledge_rag_runtime 抽出 embedding 调用逻辑为独立 EmbeddingProvider, 复用现有 HTTP 纯函数,RAG 路径零回归 - 新增 FewShotSample 表模型(样本池),注册到 db/base.py 和 models/__init__.py 供 few-shot 飞轮沉淀已确认风险观测
This commit is contained in:
138
server/src/app/services/embedding_provider.py
Normal file
138
server/src/app/services/embedding_provider.py
Normal file
@@ -0,0 +1,138 @@
|
||||
"""公共 Embedding 提供者。
|
||||
|
||||
把 ``knowledge_rag_runtime`` 里 embedding 调用逻辑抽出来,供 RAG 和
|
||||
few-shot 检索复用。本模块只依赖现有模块级纯函数和 ``RuntimeModelConfig``,
|
||||
不改动 ``_LightRagRuntime`` 的行为,RAG 路径保持零回归风险。
|
||||
|
||||
典型用法::
|
||||
|
||||
provider = EmbeddingProvider.from_settings(session)
|
||||
vectors = provider.embed(["差旅超标", "票单不一致"])
|
||||
dim = provider.dimension()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.core.logging import get_logger
|
||||
from app.services.knowledge_rag_runtime import (
|
||||
DEFAULT_EMBEDDING_TIMEOUT_SECONDS,
|
||||
KnowledgeRagError,
|
||||
RuntimeModelConfig,
|
||||
_build_headers,
|
||||
_ensure_path,
|
||||
_extract_embedding_vectors,
|
||||
_normalize_endpoint,
|
||||
_send_json_request,
|
||||
)
|
||||
|
||||
logger = get_logger("app.services.embedding_provider")
|
||||
|
||||
|
||||
def _runtime_model_config_from_dict(config: dict[str, str]) -> RuntimeModelConfig:
|
||||
"""把 SettingsService.get_runtime_model_config 返回的 dict 转成 dataclass。"""
|
||||
|
||||
return RuntimeModelConfig(
|
||||
slot=str(config.get("slot") or "embedding"),
|
||||
provider=str(config.get("provider") or ""),
|
||||
model=str(config.get("model") or ""),
|
||||
endpoint=str(config.get("endpoint") or ""),
|
||||
api_key=str(config.get("apiKey") or ""),
|
||||
capability=str(config.get("capability") or ""),
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingProvider:
|
||||
"""对 embedding 模型的轻量封装。
|
||||
|
||||
设计要点:
|
||||
- 持有一个 ``RuntimeModelConfig``,构造即固定,不依赖 LightRAG。
|
||||
- 复用 ``knowledge_rag_runtime`` 的 HTTP 调用纯函数,行为与 RAG 完全一致。
|
||||
- 维度采用惰性探测(首次 embed 后缓存),避免空构造就打远端。
|
||||
"""
|
||||
|
||||
def __init__(self, config: RuntimeModelConfig) -> None:
|
||||
self.config = config
|
||||
self._dimension: int | None = None
|
||||
|
||||
@classmethod
|
||||
def from_settings(cls, session: Any) -> "EmbeddingProvider":
|
||||
"""从 SettingsService 取 embedding 配置构造 provider。"""
|
||||
|
||||
from app.services.settings import SettingsService
|
||||
|
||||
raw = SettingsService(session).get_runtime_model_config("embedding")
|
||||
return cls(_runtime_model_config_from_dict(raw))
|
||||
|
||||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
"""对一组文本做 embedding,返回与输入等长的向量列表。"""
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
return _request_embeddings_public(self.config, texts)
|
||||
|
||||
def dimension(self) -> int:
|
||||
"""探测 embedding 维度,结果缓存。失败抛 KnowledgeRagError。"""
|
||||
|
||||
if self._dimension is None:
|
||||
vectors = self.embed(["dimension probe"])
|
||||
if not vectors or not isinstance(vectors[0], list):
|
||||
raise KnowledgeRagError("无法从 embedding 模型返回结果中解析向量维度。")
|
||||
self._dimension = len(vectors[0])
|
||||
if self._dimension <= 0:
|
||||
raise KnowledgeRagError("embedding 模型返回了无效的向量维度。")
|
||||
return self._dimension
|
||||
|
||||
|
||||
def _request_embeddings_public(
|
||||
config: RuntimeModelConfig,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
"""按 provider 分支构造 embedding 请求。
|
||||
|
||||
与 ``_LightRagRuntime._request_embeddings`` 实现保持一致,
|
||||
保证 few-shot 检索与 RAG 走同一套调用语义。
|
||||
"""
|
||||
|
||||
from app.services.model_connectivity import AZURE_API_VERSION
|
||||
|
||||
if config.provider == "Azure OpenAI":
|
||||
from app.services.knowledge_rag_runtime import _build_azure_deployment_base
|
||||
|
||||
url = f"{_build_azure_deployment_base(config.endpoint, config.model)}/embeddings?api-version={AZURE_API_VERSION}"
|
||||
payload: dict[str, Any] = {"input": texts}
|
||||
status_code, body = _send_json_request(
|
||||
"POST",
|
||||
url,
|
||||
headers=_build_headers(config.api_key, use_bearer=False, use_api_key=True),
|
||||
payload=payload,
|
||||
timeout_seconds=DEFAULT_EMBEDDING_TIMEOUT_SECONDS,
|
||||
)
|
||||
elif config.provider == "Ollama":
|
||||
url = _ensure_path(_normalize_endpoint(config.endpoint), "api/embed")
|
||||
payload = {"model": config.model, "input": texts}
|
||||
status_code, body = _send_json_request(
|
||||
"POST",
|
||||
url,
|
||||
headers={"Content-Type": "application/json", "Accept": "application/json"},
|
||||
payload=payload,
|
||||
timeout_seconds=DEFAULT_EMBEDDING_TIMEOUT_SECONDS,
|
||||
)
|
||||
else:
|
||||
url = _ensure_path(_normalize_endpoint(config.endpoint), "embeddings")
|
||||
payload = {"model": config.model, "input": texts}
|
||||
status_code, body = _send_json_request(
|
||||
"POST",
|
||||
url,
|
||||
headers=_build_headers(config.api_key, use_bearer=True),
|
||||
payload=payload,
|
||||
timeout_seconds=DEFAULT_EMBEDDING_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
from http import HTTPStatus
|
||||
|
||||
if status_code >= HTTPStatus.BAD_REQUEST:
|
||||
raise KnowledgeRagError(f"embedding 模型返回异常状态码 {status_code}。")
|
||||
|
||||
return _extract_embedding_vectors(body, provider=config.provider)
|
||||
Reference in New Issue
Block a user