增加对Rerank功能支持

backend/app/engine/__init__.py

@@ -11,6 +11,7 @@ from sqlalchemy import create_engine, Engine
 from app.engine.loaders.db import makeDescriptionByEngine
 from app.engine.tools import ToolFactory
 from app.engine.index import get_index
+from app.settings import get_node_postprocessors
 
 sql_database = None
 sql_obj_index = None
@@ -53,12 +54,15 @@ def get_chat_engine(filters=None, params=None):
                                                             )
 
         # 创建向量检索查询工具
+        postprocess = get_node_postprocessors()
         query_engine = index.as_query_engine(
-            similarity_top_k=top_k, filters=filters
+            similarity_top_k=top_k, filters=filters,
+            node_postprocessors=postprocess,
         )
         query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine, name="zj_query_tool",
-                                                          description="由博微公司编制的关于电力造价知识、电力造价编制软件知识和造价工程文件结构的知识库。适用于查询电力领域、电力造价领域、博微、博微电力、博微造价等业务等内容。如果本知识库没有直接答案但有解决思路的可以返回解决办法后建议使用“zjdata_query_tool”工具。",
+                                                          description="由博微公司编制的关于电力造价知识、电力造价编制软件知识和造价工程文件结构的知识库。适用于查询电力领域、电力造价领域、博微、博微电力、博微造价等业务等内容。如果本知识库没有直接答案但有解决思路的可以返回解决办法后建议使用“zjdata_query_tool”工具。如果你不知道答案，就说你不知道，不要编造答案。",
                                                           )
+
         tools.append(summary_query_tool)
         tools.append(query_engine_tool)
         #tools.append(sql_query_tool)

backend/app/settings.py

@@ -6,9 +6,17 @@ from llama_index.core.settings import Settings
 from llama_index.llms.xinference import Xinference
 from llama_index.llms.xinference.base import DEFAULT_XINFERENCE_TEMP
 
-from app.xinference.base import XinferenceEmbedding
+from app.xinference.base import XinferenceEmbedding, XinferenceRerank
 
 
+def get_node_postprocessors():
+    rerank_model = os.getenv("RERANK_MODEL")
+    rerank_url = os.getenv("RERANK_BASE_URL")
+    postprocess = None
+    if rerank_model is None:
+        postprocess = [XinferenceRerank(rerank_model, rerank_url)]
+    return postprocess
+
 def init_settings():
     model_provider = os.getenv("MODEL_PROVIDER")
     match model_provider: