embedding.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. """
  2. Embedding 封装。默认使用 OpenAI 兼容 embedding 接口。
  3. """
  4. from __future__ import annotations
  5. import hashlib
  6. import logging
  7. from typing import List
  8. from core.config import get_config
  9. logger = logging.getLogger(__name__)
  10. def _hash_embedding(text: str, dim: int = 64) -> List[float]:
  11. """
  12. 本地兜底 embedding(仅在外部 embedding 失败时使用)。
  13. 目的不是高质量召回,而是保证流程可运行。
  14. """
  15. digest = hashlib.sha256(text.encode("utf-8")).digest()
  16. vals: List[float] = []
  17. for i in range(dim):
  18. b = digest[i % len(digest)]
  19. vals.append((b / 255.0) * 2.0 - 1.0)
  20. return vals
  21. def embed_texts(texts: List[str]) -> List[List[float]]:
  22. cfg = get_config()
  23. model = cfg.rag.embedding_model
  24. try:
  25. from openai import OpenAI
  26. client = OpenAI(
  27. api_key=cfg.rag.embedding_api_key or cfg.llm.api_key,
  28. base_url=cfg.rag.embedding_base_url or cfg.llm.base_url,
  29. )
  30. resp = client.embeddings.create(model=model, input=texts)
  31. return [d.embedding for d in resp.data]
  32. except Exception as e:
  33. logger.warning("Embedding API 调用失败,回退哈希向量: %s", e)
  34. return [_hash_embedding(t, dim=cfg.rag.fallback_embedding_dim) for t in texts]