feat(api): 接入智谱 embedding-3(1024 维)并迁移 memory_chunks 向量列

This commit is contained in:
Kevin
2026-03-30 13:54:35 +08:00
parent aac484463d
commit 42ae2a5e91
11 changed files with 157 additions and 26 deletions

View File

@@ -1,21 +0,0 @@
"""OpenAI embedding adapter — implements EmbeddingProvider port."""
from openai import AsyncOpenAI
class OpenAIEmbeddingProvider:
def __init__(self, api_key: str, model: str = "text-embedding-3-small"):
self._client = AsyncOpenAI(api_key=api_key) if api_key else None
self._model = model
async def embed_text(self, text: str) -> list[float]:
if not self._client:
return []
resp = await self._client.embeddings.create(input=[text], model=self._model)
return resp.data[0].embedding
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
if not self._client or not texts:
return []
resp = await self._client.embeddings.create(input=texts, model=self._model)
return [item.embedding for item in sorted(resp.data, key=lambda d: d.index)]

View File

@@ -0,0 +1,56 @@
"""智谱 BigModel 国内 embedding API — 实现 EmbeddingProviderzai-sdk / ZhipuAiClient"""
from __future__ import annotations
import asyncio
from zai import ZhipuAiClient
from app.core.embedding import MEMORY_EMBEDDING_DIMENSION
# 单次请求最多 64 条文本(智谱 Embedding-3 文档)
_EMBED_BATCH_SIZE = 64
class ZhipuEmbeddingProvider:
def __init__(
self,
*,
api_key: str,
base_url: str | None = None,
model: str = "embedding-3",
) -> None:
self._model = model
if not api_key:
self._client = None
elif base_url:
self._client = ZhipuAiClient(
api_key=api_key,
base_url=base_url.rstrip("/"),
)
else:
self._client = ZhipuAiClient(api_key=api_key)
def _create_vectors_sync(self, texts: list[str]) -> list[list[float]]:
assert self._client is not None
resp = self._client.embeddings.create(
input=texts,
model=self._model,
dimensions=MEMORY_EMBEDDING_DIMENSION,
)
ordered = sorted(resp.data, key=lambda d: d.index or 0)
return [list(item.embedding) for item in ordered]
async def embed_text(self, text: str) -> list[float]:
vectors = await self.embed_texts([text])
return vectors[0] if vectors else []
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
if not self._client or not texts:
return []
out: list[list[float]] = []
for i in range(0, len(texts), _EMBED_BATCH_SIZE):
batch = texts[i : i + _EMBED_BATCH_SIZE]
part = await asyncio.to_thread(self._create_vectors_sync, batch)
out.extend(part)
return out