From f192efadd903d89b51c2f597cf2bfa6ed4e5c0e4 Mon Sep 17 00:00:00 2001
From: CNlaojing <154053522+CNlaojing@users.noreply.github.com>
Date: Thu, 13 Mar 2025 05:54:39 +0800
Subject: [PATCH] Update embedding_adapters.py

---
 embedding_adapters.py | 44 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/embedding_adapters.py b/embedding_adapters.py
index 6bc6ddb..03eda9e 100644
--- a/embedding_adapters.py
+++ b/embedding_adapters.py
@@ -1,10 +1,10 @@
 # embedding_adapters.py
 # -*- coding: utf-8 -*-
 import logging
-import requests
 import traceback
 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:
     """
@@ -187,6 +187,44 @@ class GeminiEmbeddingAdapter(BaseEmbeddingAdapter):
             logging.error(f"Gemini embed_content parse error: {e}\n{traceback.format_exc()}")
             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(
     interface_format: str,
     api_key: str,
@@ -207,5 +245,7 @@ def create_embedding_adapter(
         return MLStudioEmbeddingAdapter(api_key, base_url, model_name)
     elif fmt == "gemini":
         return GeminiEmbeddingAdapter(api_key, model_name, base_url)
+    elif fmt == "siliconflow":
+         return SiliconFlowEmbeddingAdapter(api_key, base_url, model_name)    
     else:
         raise ValueError(f"Unknown embedding interface_format: {interface_format}")