refactor: 移除硬编码配置,使用环境变量管理API密钥和路径

重构XinferenceEmbeddings和XinferenceReRankerModel类,移除硬编码的API密钥和路径
统一使用环境变量管理配置,提高安全性和灵活性
修改数据库路径为相对路径,增强可移植性
This commit is contained in:
2025-09-24 10:21:09 +08:00
parent 0541db1e34
commit d8d32ced9a
4 changed files with 30 additions and 23 deletions
+14 -9
View File
@@ -23,13 +23,13 @@ from urllib.parse import urljoin
class XinferenceEmbeddings(Embeddings):
"""SiliconFlow嵌入模型封装"""
def __init__(self, api_key: str, model: str = os.getenv("EMBEDDING_MODEL_NAME", "bge-m3")):
self.api_key = api_key
def __init__(self, model: str = os.getenv("EMBEDDING_MODEL_NAME", "bge-m3")):
self.model = model
base_url = os.getenv("XINFERENCE_URL", "http://10.1.16.39:9995")
base_url = os.getenv("EMBEDDING_BASE_URL", "http://10.1.16.39:9995")
self.url = urljoin(base_url.rstrip('/') + '/', 'v1/embeddings')
api_key = os.getenv("EMBEDDING_API_KEY", "")
self.headers = {
"Authorization": f"Bearer {self.api_key}",
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
@@ -89,12 +89,13 @@ class XinferenceReRankerModel:
List[dict]: 重排序后的文档列表,每个元素包含document内容、相关性分数和原始索引
"""
base_url = os.getenv("XINFERENCE_URL", "http://10.1.16.39:9995")
base_url = os.getenv("RERANKER_BASE_URL", "http://10.1.16.39:9995")
model_name = os.getenv("RERANKER_MODEL_NAME", "bge-reranker-v2-m3")
api_key = os.getenv("RERANKER_API_KEY", "")
rerank_url = urljoin(base_url.rstrip('/') + '/', 'v1/rerank')
params = {"documents": documents, "query": query, "top_n": top_k, "return_documents": True, "model": model_name}
headers = {
"Authorization": "Bearer <token>", # 这里需要替换为实际的token
"Authorization": f"Bearer {api_key}", # 这里需要替换为实际的token
"Content-Type": "application/json"
}
@@ -123,12 +124,14 @@ class XinferenceReRankerModel:
Returns:
List[dict]: 重排序后的文档列表,每个元素包含document内容、相关性分数和原始索引
"""
base_url = os.getenv("XINFERENCE_URL", "http://10.1.16.39:9995")
base_url = os.getenv("RERANKER_BASE_URL", "http://10.1.16.39:9995")
rerank_url = urljoin(base_url.rstrip('/') + '/', 'v1/rerank')
model_name = os.getenv("RERANKER_MODEL_NAME", "bge-reranker-v2-m3")
api_key = os.getenv("RERANKER_API_KEY", "")
params = {"documents": documents, "query": query, "top_n": top_k, "return_documents": True, "model": model_name}
headers = {
"Authorization": "Bearer <token>", # 这里需要替换为实际的token
"Authorization": f"Bearer {api_key}", # 这里需要替换为实际的token
"Content-Type": "application/json"
}
@@ -214,7 +217,9 @@ class OpenAiLLM:
messages=[{'role': 'user', 'content': user_prompt}],
**kwargs
)
return completion.choices[0].message
message = completion.choices[0].message
message.usage = completion.usage
return message
except Exception as e:
raise RuntimeError(f"OpenAiLLM:ainvoke:error:{str(e)}") from e