siliconflow.py•2.08 kB
import os
from typing import List
import openai
from memos.configs.embedder import SiliconFlowEmbedderConfig
from memos.embedders.base import BaseEmbedder
from memos.log import get_logger
logger = get_logger(__name__)
class SiliconFlowEmbedder(BaseEmbedder):
"""SiliconFlow API embedder implementation."""
def __init__(self, config: SiliconFlowEmbedderConfig):
"""Initialize the SiliconFlow embedder with the given configuration."""
self.config = config
# 初始化OpenAI客户端，使用SiliconFlow API
self.client = openai.OpenAI(
api_key=config.api_key,
base_url=config.api_base
)
# 设置模型名称
self.model_name = config.model_name or "Qwen/Qwen3-Embedding-0.6B"
logger.info(f"Initialized SiliconFlow embedder with model: {self.model_name}")
def embed(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for the given texts using SiliconFlow API."""
try:
# 调用SiliconFlow嵌入API
response = self.client.embeddings.create(
model=self.model_name,
input=texts
)
# 提取嵌入向量
embeddings = [data.embedding for data in response.data]
logger.debug(f"Generated embeddings for {len(texts)} texts")
return embeddings
except Exception as e:
logger.error(f"Failed to generate embeddings: {e}")
raise RuntimeError(f"SiliconFlow embedding failed: {e}")
def get_embedding_dimension(self) -> int:
"""Get the dimension of the embedding vectors."""
if self.config.embedding_dims:
return self.config.embedding_dims
# 对于Qwen3-Embedding-0.6B，维度是1024
if "Qwen3-Embedding" in self.model_name:
return 1024
# 对于BAAI/bge-m3，维度是1024
elif "bge-m3" in self.model_name:
return 1024
# 默认维度
return 1024