"""文本向量化 —— 把"怎么算向量"收敛到一处。 两种后端,用环境变量 `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)