调整提示词、简化代码

rag2_0/tool/ModelTool.py

@@ -193,7 +193,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:invoke:error:{str(e)}") from e
 
@@ -225,36 +227,16 @@ class OpenAiLLM:
 
 if __name__ == "__main__":
     # 测试重排模型
-    reranker = SiliconFlowReRankerModel()
-    
-    # 测试用例1：简单问题
-    query = "如何通过【电力经济评价软件】的【打开】功能加载工程文件?"
-    documents = []
-    results = reranker.rerank(query, documents)
-    print(f"测试用例1 - 查询:{query}")
-    for idx, item in enumerate(results):
-        print(f"{idx+1}. 文档: {item['document']}, 分数: {item['score']}")
-    print("-" * 50)
-    
-    # 异步测试示例
-    async def test_async():
-        # 测试异步嵌入
-        api_key = APIKeyManager.get_api_key()
-        embeddings = XinferenceEmbeddings(api_key=api_key)
-        query_embedding = await embeddings.embed_query_async("测试查询")
-        print(f"异步嵌入向量维度: {len(query_embedding)}")
-        
-        # 测试异步重排序
-        results = await SiliconFlowReRankerModel.rerank_async(query, documents)
-        print(f"异步重排序结果数量: {len(results)}")
-        
-        # 测试异步LLM调用
-        llm = OpenAiLLM()
-        response = await llm.ainvoke("你好，请简单介绍一下自己")
-        print(f"异步LLM响应: {response.content}")
-    
-    # 如果需要运行异步测试，取消下面的注释
-    # import asyncio
-    # asyncio.run(test_async())
-
+    base_url = os.getenv("OPENAI_API_BASE")
+    model_name = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
+    # 初始化LLM
+    llm_params = {
+        "temperature": 0.4,  # 降低随机性，使结果更确定
+        "top_p": 0.7,
+        "model": model_name,
+        "base_url": base_url
+    }
     
+    _llm = OpenAiLLM(**llm_params) 
+    promt="""你好，请简单介绍一下自己"""
+    print(_llm.invoke(promt))