database query test

src/config/__init__.py

@@ -42,7 +42,7 @@ class Config(SimpleConfig):
         self.add_item("stream", default=True, des="是否开启流式输出")
         self.add_item("save_dir", default="saves", des="保存目录")
         # 功能选项
-        self.add_item("enable_query_rewrite", default=True, des="是否开启查询重写")
+        self.add_item("enable_reranker", default=False, des="是否开启重排序")
         self.add_item("enable_knowledge_base", default=True, des="是否开启知识库")
         self.add_item("enable_knowledge_graph", default=True, des="是否开启知识图谱")
         self.add_item("enable_search_engine", default=True, des="是否开启搜索引擎")
@@ -58,6 +58,7 @@ class Config(SimpleConfig):
         ### <<< 默认配置结束
 
         self.filename = filename or os.path.join(self.save_dir, "config", "config.yaml")
+        os.makedirs(os.path.dirname(self.filename), exist_ok=True)
 
         self.load()
         self.handle_self()
@@ -79,7 +80,6 @@ class Config(SimpleConfig):
                 if not model_rel_path.startswith("/"):
                     self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path)
 
-
     def load(self):
         """根据传入的文件覆盖掉默认配置"""
         logger.info(f"Loading config from {self.filename}")
@@ -108,6 +108,7 @@ class Config(SimpleConfig):
         if self.filename is None:
             logger.warning("Config file is not specified, save to default config/base.yaml")
             self.filename = os.path.join(self.save_dir, "config", "config.yaml")
+            os.makedirs(os.path.dirname(self.filename), exist_ok=True)
 
         if self.filename.endswith(".json"):
             with open(self.filename, 'w+') as f:

src/core/database.py

@@ -21,8 +21,12 @@ class DataBaseLite:
         self.metadata = kwargs.get("metaname", {})
         self.files = kwargs.get("files", [])
         self.embed_model = kwargs.get("embed_model", None)
-        self.id2file = {f["file_id"]: f for f in self.files}
 
+    def id2file(self, file_id):
+        for f in self.files:
+            if f["file_id"] == file_id:
+                return f
+        return None
 
     def update(self, metadata):
         self.metadata = metadata

src/core/graphbase.py

@@ -277,7 +277,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, keyword)
 
-    def query_by_vector(self, entity_name,  num_of_res=2, threshold=0.9,kgdb_name='neo4j', hops=2):
+    def query_by_vector(self, entity_name, num_of_res=2, threshold=0.9, kgdb_name='neo4j', hops=2):
         self.use_database(kgdb_name)
         result = self.query_by_vector_tep(entity_name)
         querys = []

src/core/retriever.py

@@ -7,18 +7,20 @@ class Retriever:
 
     def __init__(self, config, dbm, model):
         self.config = config
-        self.reranker = Reranker(config)
         self.dbm = dbm
         self.model = model
 
+        if self.config.enable_reranker:
+            self.reranker = Reranker(config)
+
     def retrieval(self, query, history, meta):
 
         refs = {}
-
         refs["meta"] = meta
-        refs["rewrite_query"] = self.rewrite_query(query, history, meta)
+        refs["rewritten_query"] = self.rewrite_query(query, history, refs)
-        refs["knowledge_base"] = self.query_knowledgebase(query, history, meta)
+        refs["entities"] = self.reco_entities(query, history, refs)
-        refs["graph_base"] = self.query_graph(query, history, meta, entities=refs["rewrite_query"][1])
+        refs["knowledge_base"] = self.query_knowledgebase(query, history, refs)
+        refs["graph_base"] = self.query_graph(query, history, refs)
 
         return refs
 
@@ -28,16 +30,19 @@ class Retriever:
 
         external = ""
 
+        # 解析知识库的结果
         kb_res = refs.get("knowledge_base").get("results", [])
         if len(kb_res) > 0:
             kb_text = "\n".join([f"{r['id']}: {r['entity']['text']}" for r in kb_res])
             external += f"知识库信息: \n\n{kb_text}"
 
+        # 解析图数据库的结果
         db_res = refs.get("graph_base").get("results", [])
         if len(db_res["nodes"]) > 0:
             db_text = '\n'.join([f"{edge['source_name']}和{edge['target_name']}的关系是{edge['type']}" for edge in db_res['edges']])
             external += f"图数据库信息: \n\n{db_text}"
 
+        # 构造查询
         if len(external) > 0:
             query = f"以下是参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答：{query}"
 
@@ -50,101 +55,90 @@ class Retriever:
         """
         raise NotImplementedError
 
-    def query_graph(self, query, history, meta, entities):
+    def query_graph(self, query, history, refs):
         # res = model.predict("qiansdgsa, dasdh ashdsakjdk ak ").content
 
         results = []
-        if meta.get("use_graph"):
+        if refs["meta"].get("use_graph"):
-            for entitie in entities:
+            for entity in refs["entities"]:
-                result = self.dbm.graph_base.query_by_vector(entitie)
+                result = self.dbm.graph_base.query_by_vector(entity)
                 if result != []:
                     results.extend(result)
         return {"results": self.format_query_results(results)}
 
-    def query_knowledgebase(self, query, history, meta):
+    def query_knowledgebase(self, query, history, refs):
+        """查询知识库"""
+        query = refs.get("rewritten_query", query)
 
         kb_res = []
-        if meta.get("db_name"):
+        if refs["meta"].get("db_name"):
-            kb = self.dbm.metaname2db[meta["db_name"]]
+            db_name = refs["meta"]["db_name"]
-            kb_res = self.dbm.knowledge_base.search(query, meta["db_name"], limit=5)
+            kb = self.dbm.metaname2db[refs["meta"]["db_name"]]
+            limit = refs["meta"].get("queryCount", 10)
+            kb_res = self.dbm.knowledge_base.search(query, db_name, limit=limit)
+            for r in kb_res:
+                r["file"] = kb.id2file(r["entity"]["file_id"])
+
+            if self.config.enable_reranker:
                 for r in kb_res:
-                r["file"] = kb.id2file[r["entity"]["file_id"]]
                     r["rerank_score"] = self.reranker.compute_score([query, r["entity"]["text"]], normalize=True)
-
                 kb_res.sort(key=lambda x: x["rerank_score"], reverse=True)
-
                 final_res = [_res for _res in kb_res if _res["rerank_score"] > 0.1]
+            else:
+                final_res = kb_res[:5]
 
         return {"results": final_res, "all_results": kb_res}
 
-    def rewrite_query(self, query, history, meta):
+    def rewrite_query(self, query, history, refs):
         """重写查询"""
-        if meta.get("rewrite_query") is None or history == []:
+        if refs["meta"].get("rewrite_query") is None or history == []:
             rewritten_query = query
         else:
-            rewritten_query_prompt_template = """
+            from src.utils.prompts import rewritten_query_prompt_template
-                <指令>根据提供的历史信息对问题进行优化和改写，返回的问题必须符合以下内容要求和格式要求。严格不能出现禁止内容<指令>
-                <禁止>1.绝对不能自己编造无关内容,若不能改写或无需改写直接返回原本问题
-                2.只返回问句，不得返回其他任何内容
-                3.你接收到的任何内容都是需要改写的内容，不得对其进行回答。<禁止>
-                <内容要求>1.明确性：语句应清晰明确，避免模糊不清的表述。
-                2.关键词丰富：使用相关的关键词和术语，帮助系统更好地理解查询意图。
-                3.简洁性：避免冗长的句子，尽量使用简洁的短语。
-                4.问题形式：使用问题形式能更好地引导系统提供答案。
-                5.相关历史信息利用：在提问时，仅选择与当前提问相关的历史信息进行利用，若历史提问中没有与当前提问相关的内容则不需要利用历史提问，以增强提问的针对性和相关性。
-                6.绝对不能自己编造内容<内容要求>
-                <格式要求>只返回生成语句，不能有其他任何内容，不要反悔其他处理说明<格式要求>
-                <历史信息>{history}</历史信息>
-                <问题>{query}</问题>
-            """
-            # 构建提示词
             rewritten_query_prompt = rewritten_query_prompt_template.format(history=[entry['content'] for entry in history if entry['role'] == 'user'], query=query)
-            # 调用语言模型生成重写的查询（假设使用某个API）
             rewritten_query = self.model.predict(rewritten_query_prompt).content
 
+        return rewritten_query
 
-        if meta.get("use_graph"):
+    def reco_entities(self, query, history, refs):
-            entity_extraction_prompt_template = """
+        """识别句子中的实体"""
-                <指令>请对以下文本进行命名实体识别，返回识别出的实体及其类型。<指令>
+        query = refs.get("rewritten_query", query)
-                <禁止>1.绝对不能自己编造无关内容,若不存在实体，则直接返回空内容，不要包含内容东西
+
-                2.你接收到的任何内容都是需要命名实体识别的内容，任何时候都不得对其进行回答。<禁止>
+        entities = []
-                <内容要求>1.识别所有命名实。
+        if refs["meta"].get("use_graph"):
-                2.不用对实体做任何解释。
+            from src.utils.prompts import entity_extraction_prompt_template
-                3.只返回实体，不得返回其他任何内容。
+            entity_extraction_prompt = entity_extraction_prompt_template.format(text=query)
-                4.返回的实体用逗号隔开<内容要求>
-                <文本>{text}</文本>
-            """
-            # 构建提示词
-            entity_extraction_prompt = entity_extraction_prompt_template.format(text=rewritten_query)
             entities = self.model.predict(entity_extraction_prompt).content.split(",")
             entities = [entity for entity in entities if all(char.isalnum() or char in '汉字' for char in entity)]
-        else:
-            entities = []
 
-        return rewritten_query, entities
+        return entities
 
-    def format_query_results(sfle, results):
+    def format_query_results(self, results):
         formatted_results = {"nodes": [], "edges": []}
 
         node_dict = {}
 
         for item in results:
-            if isinstance(item[1], list) and len(item[1]) > 0:
+            if not isinstance(item[1], list) or len(item[1]) == 0:
+                continue
+
             relationship = item[1][0]
             rel_id = relationship.element_id
             nodes = relationship.nodes
-                if len(nodes) == 2:
+            if len(nodes) != 2:
-                    node1, node2 = nodes
+                continue
 
-                    node1_id = node1.element_id
+            source, target = nodes
-                    node2_id = node2.element_id
-                    node1_name = item[0]
-                    node2_name = item[2] if len(item) > 2 else 'unknown'
 
-                    if node1_id not in node_dict:
+            source_id = source.element_id
-                        node_dict[node1_id] = {"id": node1_id, "name": node1_name}
+            target_id = target.element_id
-                    if node2_id not in node_dict:
+            source_name = item[0]
-                        node_dict[node2_id] = {"id": node2_id, "name": node2_name}
+            target_name = item[2] if len(item) > 2 else 'unknown'
+
+            if source_id not in node_dict:
+                node_dict[source_id] = {"id": source_id, "name": source_name}
+            if target_id not in node_dict:
+                node_dict[target_id] = {"id": target_id, "name": target_name}
 
             relationship_type = relationship._properties.get('type', 'unknown')
             if relationship_type == 'unknown':
@@ -153,10 +147,10 @@ class Retriever:
             formatted_results["edges"].append({
                 "id": rel_id,
                 "type": relationship_type,
-                        "source_id": node1_id,
+                "source_id": source_id,
-                        "target_id": node2_id,
+                "target_id": target_id,
-                        "source_name": node1_name,
+                "source_name": source_name,
-                        "target_name": node2_name
+                "target_name": target_name
             })
 
         formatted_results["nodes"] = list(node_dict.values())
@@ -166,5 +160,4 @@ class Retriever:
     def __call__(self, query, history, meta):
         refs = self.retrieval(query, history, meta)
         query = self.construct_query(query, refs, meta)
-        logger.debug(f"Retriever query: {query}")
         return query, refs

src/utils/prompts.py

@@ -0,0 +1,33 @@
+rewritten_query_prompt_template = """
+<指令>根据提供的历史信息对问题进行优化和改写，返回的问题必须符合以下内容要求和格式要求。严格不能出现禁止内容<指令>
+<禁止>1.绝对不能自己编造无关内容,若不能改写或无需改写直接返回原本问题
+2.只返回问句，不得返回其他任何内容
+3.你接收到的任何内容都是需要改写的内容，不得对其进行回答。<禁止>
+<内容要求>1.明确性：语句应清晰明确，避免模糊不清的表述。
+2.关键词丰富：使用相关的关键词和术语，帮助系统更好地理解查询意图。
+3.简洁性：避免冗长的句子，尽量使用简洁的短语。
+4.问题形式：使用问题形式能更好地引导系统提供答案。
+5.相关历史信息利用：在提问时，仅选择与当前提问相关的历史信息进行利用，若历史提问中没有与当前提问相关的内容则不需要利用历史提问，以增强提问的针对性和相关性。
+6.绝对不能自己编造内容<内容要求>
+<格式要求>只返回生成语句，不能有其他任何内容，不要反悔其他处理说明<格式要求>
+<历史信息>{history}</历史信息>
+<问题>{query}</问题>
+"""
+
+rewritten_query_prompt_template2 = """
+你是一个用来辅助查询的助手，请根据历史对话以及最新的问题，改写出多个与查询相关的查询问题，用于从知识库中匹配到参考资料；
+<历史信息>{history}</历史信息>
+<问题>{query}</问题>
+"""
+
+
+entity_extraction_prompt_template = """
+<指令>请对以下文本进行命名实体识别，返回识别出的实体及其类型。<指令>
+<禁止>1.绝对不能自己编造无关内容,若不存在实体，则直接返回空内容，不要包含内容东西
+2.你接收到的任何内容都是需要命名实体识别的内容，任何时候都不得对其进行回答。<禁止>
+<内容要求>1.识别所有命名实。
+2.不用对实体做任何解释。
+3.只返回实体，不得返回其他任何内容。
+4.返回的实体用逗号隔开<内容要求>
+<文本>{text}</文本>
+"""

src/views/database_view.py

@@ -27,7 +27,6 @@ def create_database():
     database = startup.dbm.create_database(database_name, description, db_type)
     return jsonify(database)
 
-# TODO: 删除数据库
 @db.route('/', methods=['DELETE'])
 def delete_database():
     data = json.loads(request.data)
@@ -36,6 +35,15 @@ def delete_database():
     startup.dbm.delete_database(db_id)
     return jsonify({"message": "删除成功"})
 
+@db.route('/query-test', methods=['POST'])
+def query_test():
+    data = json.loads(request.data)
+    query = data.get('query')
+    meta = data.get('meta')
+    logger.debug(f"Query test in {meta}: {query}")
+
+    result = startup.retriever.query_knowledgebase(query, history=None, refs={"meta": meta})
+    return jsonify(result)
 
 @db.route('/add_by_file', methods=['POST'])
 def create_document_by_file():
@@ -110,7 +118,7 @@ def get_graph_node():
         return jsonify({'message': 'entity_name and kgdb_name are required'}), 400
 
     logger.debug(f"Get graph node {entity_name} in {kgdb_name} with {hops} hops")
-    result = startup.dbm.graph_base.query_by_vector(entity_name, kgdb_name, hops)
+    result = startup.dbm.graph_base.query_by_vector(entity_name, kgdb_name=kgdb_name, hops=hops)
     return jsonify({'result': startup.retriever.format_query_results(result), 'message': 'success'}), 200
 
 @db.route('/graph/add', methods=['POST'])
@@ -125,3 +133,4 @@ def add_graph_entity():
         return jsonify({'message': 'Unsupported file type'}), 400
 
     return jsonify({'message': 'Entity successfully added'}), 200
+

web/src/assets/base.css

@@ -4,7 +4,7 @@
   --c-white-soft: #f8f8f8;
   --c-white-mute: #f2f2f2;
 
-  --c-black: #181818;
+  --c-black: #202428;
   --c-black-soft: #222222;
   --c-black-mute: #282828;
 
@@ -31,7 +31,7 @@
   --color-border-hover: var(--c-black-light-1);
 
   --color-heading: var(--c-text-light-1);
-  --color-text: var(--c-text-light-1);
+  --color-text: var(--c-black);
 
   --section-gap: 160px;
 
@@ -41,6 +41,8 @@
   --main-light-2: #D3EAED;
   --main-light-3: #EDF4F5;
   --main-light-4: #F2F6F7;
+  --main-light-5: #F7FAFB;
+  --main-light-6: #FAFDFD;
   --min-width: 400px;
   --min-header-width: 80px;
   --min-sider-width: 100px;
@@ -74,8 +76,6 @@ body {
   display: flow-root;
   min-height: 100vh;
   color: var(--color-text);
-  background: var(--color-background);
-  transition: color 0.5s, background-color 0.5s;
   line-height: 1.6;
   font-family: 'HarmonyOS Sans SC', Inter, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu,
     Cantarell, 'Fira Sans', 'Droid Sans', 'Helvetica Neue', sans-serif;

web/src/components/ChatComponent.vue

@@ -22,40 +22,42 @@
           <CompassFilled v-if="meta.use_web" />
           <GoldenFilled v-if="meta.use_graph"/>
         </div> -->
-        <a-dropdown v-if="state.selectedKB !== null">
+        <a-dropdown v-if="meta.selectedKB !== null">
           <a class="ant-dropdown-link nav-btn" @click.prevent>
-            <component :is="state.selectedKB === null ? BookOutlined : BookFilled" />
+            <!-- <component :is="meta.selectedKB === null ? BookOutlined : BookFilled" /> -->
-            <span class="text">{{ state.selectedKB === null ? '不使用' : state.databases[state.selectedKB]?.name }}</span>
+             <BookOutlined />
+            <span class="text">{{ meta.selectedKB === null ? '不使用' : opts.databases[meta.selectedKB]?.name }}</span>
           </a>
           <template #overlay>
             <a-menu>
-              <a-menu-item v-for="(db, index) in state.databases" :key="index" @click="state.selectedKB=index">
+              <a-menu-item v-for="(db, index) in opts.databases" :key="index" @click="meta.selectedKB=index">
                 <a href="javascript:;" >{{ db.name }}</a>
               </a-menu-item>
-              <a-menu-item  @click="state.selectedKB = null">
+              <a-menu-item  @click="meta.selectedKB = null">
                 <a href="javascript:;">不使用</a>
               </a-menu-item>
             </a-menu>
           </template>
         </a-dropdown>
-        <div class="nav-btn text" @click="state.showPanel = !state.showPanel">
+        <div class="nav-btn text" @click="opts.showPanel = !opts.showPanel">
-          <component :is="state.showPanel ? FolderOpenOutlined : FolderOutlined" /> <span class="text">选项</span>
+          <component :is="opts.showPanel ? FolderOpenOutlined : FolderOutlined" /> <span class="text">选项</span>
         </div>
-        <div v-if="state.showPanel" class="my-panal" ref="panel">
+        <div v-if="opts.showPanel" class="my-panal" ref="panel">
           <div class="graphbase flex-center">
             知识库
             <div @click.stop>
               <a-dropdown>
                 <a class="ant-dropdown-link " @click.prevent>
-                  <component :is="state.selectedKB === null ? BookOutlined : BookFilled" />&nbsp;
+                  <!-- <component :is="meta.selectedKB === null ? BookOutlined : BookFilled" />&nbsp; -->
-                  {{ state.selectedKB === null ? '不使用' : state.databases[state.selectedKB]?.name }}
+                  <BookOutlined />&nbsp;
+                  <span class="text">{{ meta.selectedKB === null ? '不使用' : opts.databases[meta.selectedKB]?.name }}</span>
                 </a>
                 <template #overlay>
                   <a-menu>
-                    <a-menu-item v-for="(db, index) in state.databases" :key="index" @click="state.selectedKB=index">
+                    <a-menu-item v-for="(db, index) in opts.databases" :key="index" @click="meta.selectedKB=index">
                       <a href="javascript:;">{{ db.name }}</a>
                     </a-menu-item>
-                    <a-menu-item  @click="state.selectedKB = null">
+                    <a-menu-item  @click="meta.selectedKB = null">
                       <a href="javascript:;">不使用</a>
                     </a-menu-item>
                   </a-menu>
@@ -69,12 +71,15 @@
           <div class="graphbase flex-center" @click="meta.use_web = !meta.use_web">
             搜索引擎（Bing） <div @click.stop><a-switch v-model:checked="meta.use_web" /></div>
           </div>
+          <div class="graphbase flex-center" @click="meta.rewrite_query = !meta.rewrite_query">
+            重写查询 <div @click.stop><a-switch v-model:checked="meta.rewrite_query" /></div>
+          </div>
         </div>
       </div>
     </div>
     <div v-if="conv.messages.length == 0" class="chat-examples">
       <h1>你好，我是 Athena 😊</h1>
-      <div class="opt">
+      <div class="opts">
         <div
           class="opt__button"
           v-for="(exp, key) in examples"
@@ -151,8 +156,7 @@
 </template>
 
 <script setup>
-import { reactive, ref, onMounted, toRefs, nextTick, computed } from 'vue'
+import { reactive, ref, onMounted, toRefs, nextTick, computed, watch } from 'vue'
-import { onClickOutside } from '@vueuse/core'
 import {
   SendOutlined,
   MenuOutlined,
@@ -168,6 +172,7 @@ import {
   FolderOutlined,
   FolderOpenOutlined,
 } from '@ant-design/icons-vue'
+import { onClickOutside } from '@vueuse/core'
 import { marked } from 'marked';
 import { useConfigStore } from '@/stores/config'
 
@@ -192,24 +197,29 @@ const examples = ref([
 ])
 
 const opts = reactive({
-  openDetail: false
+  showPanel: false,
+  openDetail: false,
+  databases: [],
 })
 
-const meta = reactive({
+const meta = reactive(JSON.parse(localStorage.getItem('meta')) || {
-  db_name: computed(() => state.value.databases[state.value.selectedKB]?.metaname),
   use_graph: false,
   use_web: false,
   graph_name: "neo4j",
+  rewrite_query: true,
+  selectedKB: null,
 })
 
+// 更多选项可以在 marked 文档中找到：https://marked.js.org/
 marked.setOptions({
   gfm: true,
   breaks: true,
   tables: true,
-  // 更多选项可以在 marked 文档中找到：https://marked.js.org/
 });
 
-onClickOutside(panel, () => setTimeout(() => state.value.showPanel = false, 30))
+const renderMarkdown = (text) => marked(text)
+const consoleMsg = (message) => console.log(message)
+onClickOutside(panel, () => setTimeout(() => opts.showPanel = false, 30))
 
 const handleKeyDown = (e) => {
   if (e.key === 'Enter' && !e.shiftKey) {
@@ -236,20 +246,16 @@ const renameTitle = () => {
   const firstAiMessage = conv.value.messages[1].text
   const context = `${prompt}\n\n问题: ${firstUserMessage}\n\n回复: ${firstAiMessage}，主题是（一句话）：`
   simpleCall(context).then((data) => {
-    emit('renameTitle', data.response.split("：")[0])
+    emit('rename-title', data.response.split("：")[0])
   })
 }
 
-const myAlert = (message) => alert(message)
-const renderMarkdown = (text) => marked(text)
-
 const scrollToBottom = () => {
   setTimeout(() => {
     chatContainer.value.scrollTop = chatContainer.value.scrollHeight - chatContainer.value.clientHeight
   }, 10)
 }
 
-const consoleMsg = (message) => console.log(message)
 
 const generateRandomHash = (length) => {
     let chars = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789';
@@ -260,15 +266,6 @@ const generateRandomHash = (length) => {
     return hash;
 }
 
-const appendMessage = (message, role) => {
-  conv.value.messages.push({
-    id: generateRandomHash(16),
-    role,
-    text: message
-  })
-  scrollToBottom()
-}
-
 const appendUserMessage = (message) => {
   conv.value.messages.push({
     id: generateRandomHash(16),
@@ -298,7 +295,6 @@ const updateMessage = (text, id, refs, status) => {
   } else {
     console.error('Message not found')
   }
-
   scrollToBottom()
 }
 
@@ -324,17 +320,12 @@ const updateStatus = (id, status) => {
   }
 }
 
-
 const simpleCall = (message) => {
   return new Promise((resolve, reject) => {
     fetch('/api/call', {
       method: 'POST',
-      body: JSON.stringify({
+      body: JSON.stringify({ query: message, }),
-        query: message,
+      headers: { 'Content-Type': 'application/json' }
-      }),
-      headers: {
-        'Content-Type': 'application/json'
-      }
     })
     .then((response) => response.json())
     .then((data) => resolve(data))
@@ -342,6 +333,15 @@ const simpleCall = (message) => {
   })
 }
 
+const loadDatabases = () => {
+  fetch('/api/database/', { method: "GET", })
+    .then(response => response.json())
+    .then(data => {
+      console.log(data)
+      opts.databases = data.databases
+    })
+}
+
 const sendMessage = () => {
   const user_input = conv.value.inputText.trim()
   if (user_input) {
@@ -350,6 +350,7 @@ const sendMessage = () => {
     appendAiMessage("检索中……", null)
     const cur_res_id = conv.value.messages[conv.value.messages.length - 1].id
     conv.value.inputText = ''
+    meta.db_name = opts.databases[meta.selectedKB]?.metaname
     fetch('/api/chat', {
       method: 'POST',
       body: JSON.stringify({
@@ -371,9 +372,7 @@ const sendMessage = () => {
             console.log('Finished')
             updateStatus(cur_res_id, "finished")
             isStreaming.value = false
-            if (conv.value.messages.length === 2) {
+            if (conv.value.messages.length === 2) { renameTitle() }
-              renameTitle()
-            }
             return
           }
 
@@ -394,7 +393,7 @@ const sendMessage = () => {
       return readChunk()
     })
   } else {
-    console.log('Please enter a message')
+    console.log('请输入消息')
   }
 }
 
@@ -403,14 +402,30 @@ const autoSend = (message) => {
   sendMessage()
 }
 
-const clearChat = () => {
+// const clearChat = () => {
-  conv.value.messages = []
+//   conv.value.messages = []
-  conv.value.history = []
+//   conv.value.history = []
-}
+// }
 
+// 从本地存储加载数据
 onMounted(() => {
   scrollToBottom()
-})
+  loadDatabases()
+  const storedMeta = localStorage.getItem('meta');
+  if (storedMeta) {
+    const parsedMeta = JSON.parse(storedMeta);
+    Object.assign(meta, parsedMeta);
+  }
+});
+
+// 监听 meta 对象的变化，并保存到本地存储
+watch(
+  () => meta,
+  (newMeta) => {
+    localStorage.setItem('meta', JSON.stringify(newMeta));
+  },
+  { deep: true }
+);
 </script>
 
 <style lang="less" scoped>
@@ -525,7 +540,7 @@ onMounted(() => {
     color: #333;
   }
 
-  .opt {
+  .opts {
     display: flex;
     flex-wrap: wrap;
     justify-content: center;

web/src/views/ChatView.vue

@@ -40,9 +40,6 @@ const convs = reactive(JSON.parse(localStorage.getItem('chat-convs')) || [
 
 const state = reactive({
   isSidebarOpen: true,
-  selectedKB: null,
-  showPanel: false,
-  databases: [],
 })
 
 const curConvId = ref(0)
@@ -90,18 +87,6 @@ const delConv = (index) => {
   }
 }
 
-const loadDatabases = () => {
-  fetch('/api/database/', {
-    method: "GET",
-  })
-    .then(response => response.json())
-    .then(data => {
-      console.log(data)
-      state.databases = data.databases
-    }
-  )
-}
-
 // Watch convs and save to localStorage
 watch(
   () => convs,
@@ -113,7 +98,6 @@ watch(
 
 // Load convs from localStorage on mount
 onMounted(() => {
-  loadDatabases()
   const savedSonvs = JSON.parse(localStorage.getItem('chat-convs'))
   if (savedSonvs) {
     for (let i = 0; i < savedSonvs.length; i++) {

web/src/views/DataBaseInfoView.vue

@@ -10,7 +10,7 @@
         <div class="icon"><ReadFilled /></div>
         <div class="info">
           <h3>{{ database.name }}</h3>
-          <p><span>{{ database.metaname }}</span> · <span>{{ database.metadata?.row_count }}行</span></p>
+          <p><span>{{ database.metadata?.row_count }}行 · {{  database.files?.length || 0 }}文件</span></p>
         </div>
       </div>
       <p class="description">{{ database.description }}</p>
@@ -18,7 +18,31 @@
         <a-tag color="blue" v-if="database.embed_model">Embed: {{ database.embed_model }}</a-tag>
       </div>
       <a-divider/>
-      <h3 style="margin-top: 20px;">向知识库中添加文件</h3>
+      <div class="pagebtns">
+        <a-button @click="state.curPage='add'" :class="{ 'active': state.curPage === 'add' }">
+          <CloudUploadOutlined />添加文件
+        </a-button>
+        <a-button @click="state.curPage='query-test'" :class="{ 'active': state.curPage === 'query-test' }">
+          <SearchOutlined />检索测试
+        </a-button>
+      </div>
+      <div class="query-params" v-if="state.curPage == 'query-test'">
+        <p style="text-align: center; margin: 0;"><strong>参数配置</strong></p>
+        <div class="params-item">
+          <p>检索数量：</p>
+          <a-input-number size="small" v-model:value="meta.queryCount" :min="1" :max="20" />
+        </div>
+        <div class="params-item">
+          <p>过滤低质量：</p>
+          <a-switch v-model:checked="meta.filter"  size="small" />
+        </div>
+      </div>
+    </div>
+    <div class="sider-bottom">
+    </div>
+  </div>
+  <div class="db-info-container" v-if="state.curPage == 'add'">
+    <h3>向知识库中添加文件</h3>
     <div class="upload">
       <a-upload-dragger
         class="upload-dragger"
@@ -41,19 +65,11 @@
       @click="addDocumentByFile"
       :loading="state.loading"
       :disabled="fileList.length === 0"
-        style="margin: 0px 20px 20px 20px;"
+      style="margin: 0px 20px 20px 0;"
     >
       添加到知识库
     </a-button>
     <a-button @click="handleRefresh" :loading="state.refrashing">刷新状态</a-button>
-    </div>
-    <div class="sider-bottom">
-    </div>
-  </div>
-  <div class="db-info-container">
-    <div class="query-test">
-      用于测试检索
-    </div>
     <a-table :columns="columns" :data-source="database.files" row-key="file_id" class="my-table">
       <template #bodyCell="{ column, text, record }">
         <template v-if="column.key === 'file_id'">
@@ -97,6 +113,31 @@
       </p>
     </a-drawer>
   </div>
+  <div class="db-info-container" v-else-if="state.curPage == 'query-test'">
+    <h3>检索测试</h3>
+    <div class="query-action">
+      <a-textarea
+        v-model:value="queryText"
+        placeholder="填写需要查询的句子"
+        :auto-size="{ minRows: 2, maxRows: 10 }"
+      />
+      <!-- :loading="state.searchLoading" -->
+      <a-button @click="onQuery" :disabled="queryText.length == 0">
+        <SearchOutlined v-if="!state.searchLoading"/>检索
+      </a-button>
+    </div>
+    <div class="query-test" v-if="queryResult">
+      <div class="query-card" v-for="(result, idx) in (meta.filter ? queryResult.results : queryResult.all_results)" :key="idx">
+        <p>
+          <strong>#{{ idx + 1 }}&nbsp;&nbsp;&nbsp;</strong>
+          <span>{{ result.file.filename }}&nbsp;&nbsp;&nbsp;</span>
+          <span><strong>距离</strong>：{{ result.distance.toFixed(4) }}&nbsp;&nbsp;&nbsp;</span>
+          <span v-if="result.rerank_score"><strong>重排序</strong>：{{ result.rerank_score.toFixed(4) }}</span>
+        </p>
+        <p class="query-text">{{ result.entity.text }}</p>
+      </div>
+    </div>
+  </div>
   </div>
 </template>
 
@@ -112,6 +153,8 @@ import {
   CloseCircleFilled,
   ClockCircleFilled,
   DeleteOutlined,
+  CloudUploadOutlined,
+  SearchOutlined,
 } from '@ant-design/icons-vue'
 
 
@@ -123,14 +166,55 @@ const database = ref({});
 const fileList = ref([]);
 const selectedFile = ref(null);
 
+// 查询测试
+const queryText = ref('');
+const queryResult = ref(null)
+
 const state = reactive({
   loading: false,
   refrashing: false,
+  searchLoading: false,
   lock: false,
   drawer: false,
   refreshInterval: null,
+  curPage: "add",
 });
 
+const meta = reactive({
+  queryCount: 10,
+  filter: false,
+});
+
+const onQuery = () => {
+  console.log(queryText.value)
+  state.searchLoading = true
+  if (!queryText.value.trim()) {
+    message.error('请输入查询内容')
+    state.searchLoading = false
+    return
+  }
+  meta.db_name = database.value.metaname
+  fetch('/api/database/query-test', {
+    method: "POST",
+    body: JSON.stringify({
+      query: queryText.value.trim(),
+      meta: meta
+    }),
+  })
+  .then(response => response.json())
+  .then(data => {
+    console.log(data)
+    queryResult.value = data
+  })
+  .catch(error => {
+    console.error(error)
+    message.error(error.message)
+  })
+  .finally(() => {
+    state.searchLoading = false
+  })
+}
+
 const handleFileUpload = (event) => {
   console.log(event)
   console.log(fileList.value)
@@ -371,10 +455,11 @@ onMounted(() => {
   display: flex;
   flex-direction: column;
   justify-content: space-between;
-  width: 300px;
+  width: 275px;
   height: 100%;
   padding: 0;
   border-right: 1px solid #E0EAFF;
+  flex: 0 0 275px;
 
   .sider-top {
 
@@ -387,7 +472,7 @@ onMounted(() => {
       margin-bottom: 20px;
       padding-top: 10px;
       padding-bottom: 10px;
-      background-color: #FAFAFA;
+      background-color: var(--main-light-5);
       border-bottom: 1px solid #E0EAFF;
 
       button {
@@ -397,21 +482,101 @@ onMounted(() => {
       }
     }
   }
+
+  .pagebtns {
+    display: flex;
+    flex-direction: column;
+    gap: 16px;
+
+    button {
+      padding: 10px 16px;
+      height: auto;
+      border-radius: 8px;
+      border: none;
+      background: var(--main-light-4);
+    }
+
+    .active {
+      font-weight: bold;
+      color: var(--main-color);
+      background: var(--main-light-2);
+    }
+  }
+
+  .query-params {
+    display: flex;
+    flex-direction: column;
+    padding: 20px;
+    background: var(--main-light-4);
+    border-radius: 8px;
+    margin: 20px;
+    box-sizing: border-box;
+    gap: 12px;
+
+    .params-item {
+      display: flex;
+      align-items: center;
+      justify-content: space-between;
+      gap: 8px;
+
+      p {
+        margin: 0;
+      }
+    }
+  }
 }
 
 .db-info-container {
   padding: 20px;
   flex: 1 1 auto;
+  overflow: scroll;
+
+  .query-action {
+    display: flex;
+    gap: 8px;
+    margin-bottom: 20px;
+
+    textarea {
+      padding: 12px 16px;
+      border: 1px solid var(--main-light-2);
+    }
+
+    button {
+      height: auto;
+      width: 120px;
+      box-shadow: none;
+      border: none;
+      font-weight: bold;
+      background: var(--main-light-2);
+      color: var(--main-color);
+
+      &:disabled {
+        cursor: not-allowed;
+        background: var(--main-light-4);
+        color: var(--c-black-light-3);
+      }
+    }
+  }
 
   .query-test {
-    width: 100%;
-    height: 300px;
     display: flex;
-    justify-content: center;
+    flex-direction: column;
-    align-items: center;
-    background: var(--main-light-3);
     margin-bottom: 20px;
     border-radius: 8px;
+    gap: 16px;
+
+    .query-card {
+      padding: 20px 16px;
+      border-radius: 8px;
+      background: var(--main-light-5);
+      border: 1px solid var(--main-light-3);
+    }
+
+    .query-text {
+      font-size: 14px;
+      margin-bottom: 0;
+      color: var(--c-text-light-1);
+    }
   }
 }
 

web/src/views/DataBaseView.vue

@@ -47,7 +47,7 @@
     </div>
     <h2>图数据库 &nbsp; <a-spin v-if="graphloading" :indicator="indicator" /></h2>
     <p>基于 neo4j 构建的图数据库。</p>
-    <div :class="{'graphloading': graphloading}" v-if="graph">
+    <div :class="{'graphloading': graphloading, 'databases': true}" v-if="graph">
       <div class="dbcard graphbase" @click="navigateToGraph">
         <div class="top">
           <div class="icon"><AppstoreFilled /></div>
@@ -189,6 +189,7 @@ onMounted(() => {
 }
 
 .database, .graphbase {
+  flex: 1 1 380px;
   background-color: white;
   box-shadow: 0px 1px 2px 0px rgba(16,24,40,.06),0px 1px 3px 0px rgba(16,24,40,.1);
   border: 2px solid white;

web/src/views/SettingView.vue

@@ -46,6 +46,7 @@
           <a-select style="width: 160px"
             :value="configStore.config?.reranker"
             @change="handleChange('reranker', $event)"
+            :disabled="!configStore.config.enable_reranker"
           >
             <a-select-option
               v-for="(name, idx) in items?.reranker.choices" :key="idx"
@@ -78,10 +79,14 @@
           />
         </div>
         <div class="card">
-          <span class="label">{{ items?.enable_query_rewrite.des }}</span>
+          <span class="label">{{ items?.enable_reranker.des }}
+            <a-button small v-if="needRestart.enable_reranker" @click="sendRestart">
+              <ReloadOutlined />需要重启
+            </a-button>
+          </span>
           <a-switch
-            :checked="configStore.config.enable_query_rewrite"
+            :checked="configStore.config.enable_reranker"
-            @change="handleChange('enable_query_rewrite', !configStore.config.enable_query_rewrite)"
+            @change="handleChange('enable_reranker', !configStore.config.enable_reranker)"
           />
         </div>
       </div>