优化OpenAiLLM类中的API调用,使用with语句确保资源正确释放,并增强错误处理机制,添加重试次数提示。

This commit is contained in:
2025-06-20 14:28:57 +08:00
parent cfda8443ae
commit 644b745ab2
+8 -9
View File
@@ -18,6 +18,7 @@ import os
import logging
from rag2_0.tool.APIKeyManager import APIKeyManager
class SiliconFlowEmbeddings(Embeddings):
"""SiliconFlow嵌入模型封装"""
def __init__(self, api_key: str, model: str = "bge-m3"):
@@ -141,7 +142,6 @@ class OpenAiLLM:
def invoke(self, user_prompt="你是谁?", need_retry=True):
# 初始化 OpenAI 客户端
max_retries = 3
retry_count = 0
if "timeout" not in self._kwargs:
@@ -152,7 +152,8 @@ class OpenAiLLM:
while retry_count < max_retries:
try:
api_key = APIKeyManager.get_api_key()
client = OpenAI(api_key=api_key, base_url=self._url)
# 使用with语句创建客户端,确保资源会被正确释放
with OpenAI(api_key=api_key, base_url=self._url) as client:
# 创建 Completion 请求. 超时120s
completion = client.chat.completions.create(
model=self._model,
@@ -164,13 +165,14 @@ class OpenAiLLM:
except Exception as e:
retry_count += 1
if retry_count == max_retries:
raise e
raise RuntimeError(f"调用LLM失败,已重试{max_retries}次: {str(e)}")
else:
time.sleep(5*retry_count) # 重试前等待1秒
time.sleep(5*retry_count) # 重试前等待5秒*重试次数
else:
# 创建 Completion 请求. 超时120s
api_key = APIKeyManager.get_api_key()
client = OpenAI(api_key=api_key, base_url=self._url)
# 使用with语句创建客户端,确保资源会被正确释放
with OpenAI(api_key=api_key, base_url=self._url) as client:
completion = client.chat.completions.create(
model=self._model,
messages=[{'role': 'user', 'content': user_prompt}],
@@ -184,10 +186,7 @@ if __name__ == "__main__":
# 测试用例1:简单问题
query = "如何通过【电力经济评价软件】的【打开】功能加载工程文件?"
documents = ["\n# (电力建设计价通软件) (概预算工程)工程备份管理\n## 操作步骤\n**方法一:** \n\n1、查找工程:输入工程文件名称的关键字,点击“查找”按钮,可以快速定位需查找的工程;\n\n![2](https://172.20.0.145/files/ea1d6edc-090c-4c35-be13-e506f0eeb176/image-preview)\n\n2、根据时间点找备份工程:选中对应工程文件,在右侧选中“备份时间”的备份记录,点击“还原工程”或者“另存为工程”;\n\n **还原工程:** 将工程还原保存在原路径下;\n\n **另存为工程:** 另存为一个新工程,可选择保存路径,保存后,可点击文件——打开,浏览到另存的新工程打开。\n\n注:不确定备份是否是需要时,优先建议另存为工程。\n\n![3](https://172.20.0.145/files/63fb4e9d-06ce-44e6-adbc-45ebc92c3e6f/image-preview)\n\n **方法二:** \n\n1、点击桌面软件快捷图标,右键属性—打开文件位置,直接定位软件安装根目录。 \n\n![打开文件所在位置](https://172.20.0.145/files/28540c85-2b85-4717-9b87-31eba886c5f0/image-preview)\n\n2、在软件安装根目录,点击“数据备份”文件夹,进入到文件夹内,根据修改日期找到对应工程,右键复制粘贴至桌面。\n\n![打开数据备份文件夹](https://172.20.0.145/files/76a4263c-ea75-4083-9ed8-7f2e28412e51/image-preview)\n\n![按照时间排序,找到.bak文件复制到桌面](https://172.20.0.145/files/eaa59154-6fbc-494e-817c-e5f7610e6294/image-preview)\n\n3、定位桌面复制粘贴出来的数据工程,右键\"重命名\",将bak修改成相应的文件后缀(概预算工程及施工图预算工程后缀为zwzj,招标工程及投标工程后缀为zwqd),然后点击“确定”,再通过软件的“文件”——“打开”按钮去浏览工程打开。\n",
"\n# (配网计价通D3)插件管理/全国版和专版切换\n## 使用场景\n1.打开软件提示“当前工程文件为全国版文件,请使用全国版软件打开!”,该如何打开这个工程呢?\n\n![1](https://172.20.0.145/files/3751d1c2-da12-4076-bd9f-3e1c9eced1ab/image-preview)\n\n2.打开软件提示“当前工程文件为辽宁版文件,请确认是否要在全国版软件中打开?”,这是什么意思?点击“确定”又可以打开工程?\n\n![2](https://172.20.0.145/files/68c18647-c4fa-4496-974a-f5b07352063e/image-preview)\n## 知识原理\n\n## 费用去向\n\n",
"\n(电力建设计价通软件) 云造价--停用\n# 工程文件管理\n\n## 【主页】中点击“云端工程管理”,进入博微服务大厅;\n![6](https://172.20.0.145/files/69d1d763-bb6b-459f-b7ba-264c788225e9/image-preview)![7](https://172.20.0.145/files/eb93d8ea-d3c8-45f3-86ab-f23c0d244d96/image-preview)\n## 工程文件管理界面中显示云端备份的工程列表,可支持\n![8](https://172.20.0.145/files/91200171-182b-4de7-a895-f5118bfe9d98/image-preview)\n## 高级设置:可对历史版本数量进行设置,默认数量为10,可设置(5-15);\n![9](https://172.20.0.145/files/31b63307-3313-4402-9652-9b46e9549aec/image-preview)\n## 历史版本:勾选单个工程,点击“历史版本”可查看该工程保存的不同时间节点的历史工程;\n![10](https://172.20.0.145/files/1ae83da5-9e86-4a44-9af1-013199a933c6/image-preview)\n## 在线查阅:可查看工程数据,仅为只读模式不支持任何编辑;\n![11](https://172.20.0.145/files/61ef2bc8-0d73-47bf-b66f-aebcc39d7a46/image-preview)\n## 下载;选择需要的工程点击“下载”,可下载软件版本工程;\n![12](https://172.20.0.145/files/3391d7b2-c289-485a-a459-754ce6c8e294/image-preview)\n",
"\n(配网D3软件)打开工程\n\n# (配网D3软件)打开工程\n\n## 功能入口\n各界面点击“文件”按钮——“打开”按钮 \n\n![打开工程-功能入口](https://172.20.0.145/files/7c6dbff1-932c-4cf5-ab16-b17a6309af0e/image-preview)\n## 操作步骤\n**打开工程:** \n\n点击“打开”按钮,浏览到工程存放位置,选中工程文件,点击“打开”即可。"]
documents = []
results = reranker.rerank(query, documents)
print(f"测试用例1 - 查询:{query}")
for idx, item in enumerate(results):