From 2942730c9a049cc1d89de8925cb1ffe5b76dd9c5 Mon Sep 17 00:00:00 2001
From: paituo <330435863@qq.com>
Date: Mon, 19 Aug 2024 08:59:08 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=AF=B9Rerank=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=E6=94=AF=E6=8C=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/.env.xinference        |  3 +++
 backend/app/engine/__init__.py |  8 ++++++--
 backend/app/settings.py        | 10 +++++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/backend/.env.xinference b/backend/.env.xinference
index 71bdc2b..a8e9396 100644
--- a/backend/.env.xinference
+++ b/backend/.env.xinference
@@ -19,6 +19,9 @@ EMBEDDING_MODEL=bge-m3
 EMBEDDING_BASE_URL=http://10.1.16.39:9995
 # Dimension of the embedding model to use.
 EMBEDDING_DIM=1024
+# Rerank model
+RERANK_MODEL=bge-reranker-v2-m3
+RERANK_BASE_URL=http://10.1.16.39:9995
 ##----------   OpenAI    ----------------
 ## The provider for the AI models to use.
 #MODEL_PROVIDER=openai
diff --git a/backend/app/engine/__init__.py b/backend/app/engine/__init__.py
index def5e51..444f925 100644
--- a/backend/app/engine/__init__.py
+++ b/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)
diff --git a/backend/app/settings.py b/backend/app/settings.py
index 82e83e8..ceec7ba 100644
--- a/backend/app/settings.py
+++ b/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: