Update embedding_adapters.py

This commit is contained in:
CNlaojing
2025-03-13 05:54:39 +08:00
committed by GitHub
parent de5c08fb0e
commit f192efadd9
+42 -2
View File
@@ -1,10 +1,10 @@
# embedding_adapters.py # embedding_adapters.py
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import logging import logging
import requests
import traceback import traceback
from typing import List from typing import List
from langchain_openai import OpenAIEmbeddings, AzureOpenAIEmbeddings import requests
from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
def ensure_openai_base_url_has_v1(url: str) -> str: def ensure_openai_base_url_has_v1(url: str) -> str:
""" """
@@ -187,6 +187,44 @@ class GeminiEmbeddingAdapter(BaseEmbeddingAdapter):
logging.error(f"Gemini embed_content parse error: {e}\n{traceback.format_exc()}") logging.error(f"Gemini embed_content parse error: {e}\n{traceback.format_exc()}")
return [] return []
class SiliconFlowEmbeddingAdapter(BaseEmbeddingAdapter):
"""
基于 SiliconFlow 的 embedding 适配器
"""
def __init__(self, api_key: str, base_url: str, model_name: str):
# 自动为 base_url 添加 scheme(如果缺失)
if not base_url.startswith("http://") and not base_url.startswith("https://"):
base_url = "https://" + base_url
self.url = base_url if base_url else "https://api.siliconflow.cn/v1/embeddings"
self.payload = {
"model": model_name,
"input": "Silicon flow embedding online: fast, affordable, and high-quality embedding services. come try it out!",
"encoding_format": "float"
}
self.headers = {
"Authorization": "Bearer {api_key}".format(api_key=api_key),
"Content-Type": "application/json"
}
def embed_documents(self, texts: List[str]) -> List[List[float]]:
embeddings = []
for text in texts:
self.payload["input"] = text
response = requests.post(self.url, json=self.payload, headers=self.headers)
result = response.json()
# 从返回数据中提取第一个 embedding
emb = result.get("data", [{}])[0].get("embedding", [])
embeddings.append(emb)
return embeddings
def embed_query(self, query: str) -> List[float]:
self.payload["input"] = query
# print('SiliconFlowEmbeddingAdapter发送',self.payload)
response = requests.post(self.url, json=self.payload, headers=self.headers)
result = response.json()
return result.get("data", [{}])[0].get("embedding", [])
def create_embedding_adapter( def create_embedding_adapter(
interface_format: str, interface_format: str,
api_key: str, api_key: str,
@@ -207,5 +245,7 @@ def create_embedding_adapter(
return MLStudioEmbeddingAdapter(api_key, base_url, model_name) return MLStudioEmbeddingAdapter(api_key, base_url, model_name)
elif fmt == "gemini": elif fmt == "gemini":
return GeminiEmbeddingAdapter(api_key, model_name, base_url) return GeminiEmbeddingAdapter(api_key, model_name, base_url)
elif fmt == "siliconflow":
return SiliconFlowEmbeddingAdapter(api_key, base_url, model_name)
else: else:
raise ValueError(f"Unknown embedding interface_format: {interface_format}") raise ValueError(f"Unknown embedding interface_format: {interface_format}")