rag/embed.py

"""文本向量化 —— 把"怎么算向量"收敛到一处。

两种后端,用环境变量 `EMBED_BACKEND` 切换(默认 ollama),与 agents/llm.py 同构:
- ollama:本机 Ollama 原生 `/api/embeddings`,逐条调用。
- openai:OpenAI 兼容 `/v1/embeddings`,批量;部署指向托管 embedding 服务。

bge-m3:多语种、中文检索 SOTA(1024维),对齐 PLAN 的 BGE 选型。
想退回 nomic-embed-text 只需 export FINSIGHT_EMBED_MODEL=nomic-embed-text。
"""
import os

import requests

EMBED_BACKEND = os.getenv("EMBED_BACKEND", "ollama")           # ollama | openai
EMBED_MODEL = os.getenv("FINSIGHT_EMBED_MODEL", "bge-m3")

OLLAMA_EMBED_URL = "http://localhost:11434/api/embeddings"
# embedding 端点默认复用 LLM 的 OPENAI_BASE_URL,可用 OPENAI_EMBED_BASE_URL 单独指定
OPENAI_BASE_URL = os.getenv("OPENAI_EMBED_BASE_URL",
                            os.getenv("OPENAI_BASE_URL", "http://localhost:11434/v1"))
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "ollama")


def _ollama_embed(texts: list[str], timeout: int) -> list[list[float]]:
    out = []
    for t in texts:
        r = requests.post(OLLAMA_EMBED_URL,
                          json={"model": EMBED_MODEL, "prompt": t}, timeout=timeout)
        r.raise_for_status()
        out.append(r.json()["embedding"])
    return out


def _openai_embed(texts: list[str], timeout: int) -> list[list[float]]:
    r = requests.post(f"{OPENAI_BASE_URL}/embeddings",
                      json={"model": EMBED_MODEL, "input": texts},
                      headers={"Authorization": f"Bearer {OPENAI_API_KEY}"}, timeout=timeout)
    r.raise_for_status()
    return [d["embedding"] for d in r.json()["data"]]


def embed(texts: list[str], timeout: int = 60) -> list[list[float]]:
    """list[str] -> list[向量]。"""
    if EMBED_BACKEND == "openai":
        return _openai_embed(texts, timeout)
    return _ollama_embed(texts, timeout)