Merge pull request #4 from xerrors/dev

Dev

src/config/__init__.py

@@ -34,9 +34,23 @@ class Config(SimpleConfig):
         self.filename = filename
         logger.info(f"Loading config from {filename}")
 
-        ### startup
+        ### >>> 默认配置
+        # 可以在 config/base.yaml 中覆盖
         self.mode = "cli"
-        self.stream = False
+        self.stream = True
+
+        # 功能选项
+        self.enable_query_rewrite = True
+        self.enable_knowledge_base = True
+        self.enable_knowledge_graph = True
+        self.enable_search_engine = True
+
+        # 模型配置
+        ## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
+        ## 如果需要自定义路径，则在 config/base.yaml 中配置 model_local_paths
+        self.embed_model = "bge-large-zh-v1.5"
+        self.reranker = "bge-reranker-v2-m3"
+        ### <<< 默认配置结束
 
         self.load()
         self.handle_self()
@@ -51,6 +65,7 @@ class Config(SimpleConfig):
 
 
     def load(self):
+        """根据传入的文件覆盖掉默认配置"""
         if self.filename is not None and os.path.exists(self.filename):
             if self.filename.endswith(".json"):
                 with open(self.filename, 'r') as f:

src/config/base.yaml

@@ -1,25 +1,12 @@
+# 默认配置请参考 config/__init__.py
 name: base
 
 ## model
 ### model_provider, option in deepseek, zhipu
-model_provider: qianfan
+model_provider: zhipu
-
-# 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
-# 如果需要自定义路径，则在下面配置 model_local_paths
-embed_model: bge-large-zh  # option in ["bge-large-zh"]
-
-## startup
-stream: True
-
-## knowledge
-enable_query_rewrite: True
-enable_knowledge_base: True
-enable_knowledge_graph: True
-enable_search_engine: True
-
 
 ## model dir 可以写相对路径和绝对路径
 ### 相对路径是相对于环境变量中 MODEL_ROOT_DIR 的路径
 model_local_paths:
-  bge-large-zh: bge-large-zh-v1.5
+  bge-large-zh-v1.5: bge-large-zh-v1.5
   oneke: oneke

src/core/__init__.py

@@ -1,4 +1,2 @@
 from .history import *
-from .retriever import *
+from .database import *
-from .database import *
-from .graphbase import *

src/core/database.py

@@ -1,7 +1,8 @@
 import os
 import json
 import time
-from utils import hashstr, setup_logger
+from utils import hashstr, setup_logger, is_text_pdf
+from plugins import pdf2txt
 from core.knowledgebase import KnowledgeBase
 from core.filereader import pdfreader, plainreader
 from core.graphbase import GraphDatabase
@@ -146,18 +147,25 @@ class DataBaseManager:
         support_format = [".pdf", ".txt", "*.md"]
         assert os.path.exists(file), "File not found"
         logger.info(f"Try to read file {file}")
-        if os.path.isfile(file):
+
-            if file.endswith(".pdf"):
+        if not os.path.isfile(file):
-                return pdfreader(file)
-            elif file.endswith(".txt") or file.endswith(".md"):
-                return plainreader(file)
-            else:
-                logger.error(f"File format not supported, only support {support_format}")
-                raise Exception(f"File format not supported, only support {support_format}")
-        else:
             logger.error(f"Directory not supported now!")
             raise NotImplementedError("Directory not supported now!")
 
+        if file.endswith(".pdf"):
+            if is_text_pdf(file):
+                return pdfreader(file)
+            else:
+                return pdf2txt(file, return_text=True)
+
+        elif file.endswith(".txt") or file.endswith(".md"):
+            return plainreader(file)
+
+        else:
+            logger.error(f"File format not supported, only support {support_format}")
+            raise Exception(f"File format not supported, only support {support_format}")
+
+
     def delete_file(self, db_id, file_id):
         db = self.get_kb_by_id(db_id)
         file_idx_to_delete = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0]
@@ -187,6 +195,16 @@ class DataBaseManager:
             chunks.append(text[i:i + chunk_size])
         return chunks
 
+    def delete_database(self, db_id):
+        db = self.get_kb_by_id(db_id)
+        if db is None:
+            return {"message": "database not found"}, 404
+
+        self.knowledge_base.client.drop_collection(db.metaname)
+        self.data["databases"] = [d for d in self.data["databases"] if d.db_id != db_id]
+        self._save_databases()
+        return {"message": "删除成功"}
+
     def get_kb_by_id(self, db_id):
         for db in self.data["databases"]:
             if db.db_id == db_id:

src/core/graphbase.py

@@ -57,7 +57,7 @@ class GraphDatabase:
         """切换到指定数据库"""
         self.driver = GD.driver(f"{os.environ.get('NEO4J_URI')}/{kgdb_name}", auth=(os.environ.get('NEO4J_USERNAME'), os.environ.get('NEO4J_PASSWORD')))
 
-    def txt_add_entity(self, triples, kgdb_name):
+    def txt_add_entity(self, triples, kgdb_name='neo4j'):
         """添加实体三元组"""
         self.use_database(kgdb_name)
         def create(tx, triples):
@@ -75,7 +75,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             session.execute_write(create, triples)
 
-    def file_add_entity(self, file_path, output_path,kgdb_name):
+    def file_add_entity(self, file_path, output_path, kgdb_name='neo4j'):
         self.use_database(kgdb_name)  # 切换到指定数据库
         text_path = pdf2txt(file_path)
         oneke = OneKE()
@@ -112,7 +112,7 @@ class GraphDatabase:
         """
         tx.run(query)
 
-    def query_all_nodes_and_relationships(self, kgdb_name, hops = 2):
+    def query_all_nodes_and_relationships(self, kgdb_name='neo4j', hops = 2):
         """查询图数据库中所有三元组信息"""
         self.use_database(kgdb_name)
         def query(tx, hops):
@@ -125,7 +125,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, hops)
 
-    def query_specific_entity(self, entity_name, kgdb_name, hops = 2):
+    def query_specific_entity(self, entity_name, kgdb_name='neo4j', hops = 2):
         """查询指定实体三元组信息"""
         self.use_database(kgdb_name)
         def query(tx, entity_name, hops):
@@ -138,7 +138,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, entity_name, hops)
 
-    def query_by_relationship_type(self, relationship_type, kgdb_name, hops = 2):
+    def query_by_relationship_type(self, relationship_type, kgdb_name='neo4j', hops = 2):
         """查询指定关系三元组信息"""
         self.use_database(kgdb_name)
         def query(tx, relationship_type, hops):
@@ -151,7 +151,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, relationship_type, hops)
 
-    def query_entity_like(self, keyword, kgdb_name, hops = 2):
+    def query_entity_like(self, keyword, kgdb_name='neo4j', hops = 2):
         """模糊查询"""
         self.use_database(kgdb_name)
         def query(tx, keyword, hops):
@@ -166,7 +166,7 @@ class GraphDatabase:
         with self.driver.session() as session:
             return session.execute_read(query, keyword, hops)
 
-    def query_node_info(self, node_name, kgdb_name, hops = 2):
+    def query_node_info(self, node_name, kgdb_name='neo4j', hops = 2):
         """查询指定节点的详细信息返回信息"""
         self.use_database(kgdb_name)  # 切换到指定数据库
         def query(tx, node_name, hops):
@@ -179,6 +179,17 @@ class GraphDatabase:
 
         with self.driver.session() as session:
             return session.execute_read(query, node_name, hops)
+        
+    # def format_query_results(self, results):
+    #     formatted_results = []
+    #     for row in results:
+    #         n, rs, m = row
+    #         entity_a = n['name']
+    #         entity_b = m['name']
+    #         for rel in rs:
+    #             relationship = rel.type
+    #             formatted_results.append(f"实体 {entity_a} 和 实体 {entity_b} 的关系是 {relationship}")
+    #     return formatted_results
 
 
 

src/core/knowledgebase.py

@@ -66,8 +66,7 @@ class KnowledgeBase:
         res = self.client.insert(collection_name=collection_name, data=data)
         return res
 
-    def search(self, query, collection_name, limit=None):
+    def search(self, query, collection_name, limit=3):
-        limit = limit or self.default_query_limit
 
         query_vectors = self.embed_model.encode_queries([query])
 

src/core/retriever.py

@@ -1,30 +1,40 @@
+from core.startup import dbm, model
+from models.embedding import ReRanker
+
 class Retriever:
 
     def __init__(self, config):
         self.config = config
+        self.reranker = ReRanker(config)
 
-    def retrieval(self, query):
+    def retrieval(self, query, history, meta):
 
         refs = {}
 
         # TODO: 查询分类、查询重写、查询分解、查询伪文档生成（HyDE）)
-        # NOTE：2024-07-14 暂时禁用知识检索
+        refs["meta"] = meta
+        refs["rewrite_query"] = self.rewrite_query(query, history)
+        refs["knowledge_base"] = self.query_knowledgebase(query, history, meta)
+        refs["graph_base"] = self.query_graph(query, history, meta, entities=refs["rewrite_query"][1])
 
         return refs
 
-    def construct_query(self, query, refs):
+    def construct_query(self, query, refs, meta):
-        # TODO：Reranking
-
         if len(refs) == 0:
             return query
 
         external = ""
 
-        kb_res = refs.get("knowledge_base")
+        kb_res = refs.get("knowledge_base").get("results", [])
-        if kb_res:
+        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) > 0:
+        #     db_text = "\n".join([f"{r['id']}: {r['entity']['text']}" for r in db_res])
+        #     external += f"图数据库信息: \n\n{db_text}"
+
         if len(external) > 0:
             query = f"以下是参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答：{query}"
 
@@ -37,11 +47,98 @@ class Retriever:
         """
         raise NotImplementedError
 
-    def rewrite_query(self, query):
+    def query_graph(self, query, history, meta, entities):
-        """重写查询"""
+        # res = model.predict("qiansdgsa, dasdh ashdsakjdk ak ").content
-        raise NotImplementedError
 
-    def __call__(self, query):
+        results = []
-        refs = self.retrieval(query)
+        if meta.get("use_graph"):
-        query = self.construct_query(query, refs)
+            for entitie in entities:
+                result = dbm.graph_base.query_entity_like(entitie)
+                results.extend(result) if result else None
+
+        return {"results": self.format_query_results(results)}
+
+    def query_knowledgebase(self, query, history, meta):
+
+        kb_res = []
+        if meta.get("db_name"):
+            kb_res = dbm.knowledge_base.search(query, meta["db_name"], limit=5)
+            for r in kb_res:
+                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]
+        return {"results": final_res, "all_results": kb_res}
+
+    def rewrite_query(self, query, history):
+        """重写查询"""
+        if history == []:
+            rewritten_query = query
+        else:
+            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 = model.predict(rewritten_query_prompt).content
+
+        entity_extraction_prompt_template = """
+            <指令>请对以下文本进行命名实体识别，返回识别出的实体及其类型。<指令>
+            <禁止>1.绝对不能自己编造无关内容,若不存在实体，则直接返回空内容，不要包含内容东西
+            2.你接收到的任何内容都是需要命名实体识别的内容，任何时候都不得对其进行回答。<禁止>
+            <内容要求>1.识别所有命名实。
+            2.不用对实体做任何解释。
+            3.只返回实体，不得返回其他任何内容。
+            4.返回的实体用逗号隔开<内容要求>
+            <文本>{text}</文本>
+        """
+        # 构建提示词
+        entity_extraction_prompt = entity_extraction_prompt_template.format(text=rewritten_query)
+        entities = model.predict(entity_extraction_prompt).content.split(",")
+        entities = [entity for entity in entities if all(char.isalnum() or char in '汉字' for char in entity)]
+
+        return rewritten_query, entities
+
+    def format_query_results(self, results):
+        formatted_results = {"nodes": [], "edges": []}
+        for row in results:
+            n, relations, m = row
+            formatted_results["nodes"].append({
+                "id": n.id,
+                "name": n._properties["name"],
+                "properties": n._properties
+            })
+            formatted_results["nodes"].append({
+                "id": m.id,
+                "name": m._properties["name"],
+                "properties": m._properties
+            })
+            for rel in relations:
+                formatted_results["edges"].append({
+                    "id": rel.id,
+                    "type": rel.type,
+                    "source": rel.start_node.id,
+                    "target": rel.end_node.id,
+                    "source_name": rel.start_node._properties["name"],
+                    "target_name": rel.end_node._properties["name"],
+                })
+        return formatted_results
+
+    def __call__(self, query, history, meta):
+        refs = self.retrieval(query, history, meta)
+        query = self.construct_query(query, refs, meta)
         return query, refs

src/core/startup.py

@@ -1,5 +1,4 @@
-from core import Retriever, DataBaseManager
+from core import DataBaseManager
-from core.graphbase import GraphDatabase
 from models import select_model
 from config import Config
 
@@ -7,6 +6,5 @@ from config import Config
 config = Config("config/base.yaml")
 model = select_model(config)
 dbm = DataBaseManager(config)
-retriever = Retriever(config)
 
 # 启动本地图数据库

src/models/chat_model.py

@@ -67,15 +67,13 @@ class QianfanResponse:
 
 class Qianfan:
 
-    def __init__(self, model_name="ernie-lite-8k") -> None:
+    def __init__(self, model_name="ernie_speed") -> None:
         self.model_name = model_name
         access_key = os.getenv("QIANFAN_ACCESS_KEY")
         secret_key = os.getenv("QIANFAN_SECRET_KEY")
         self.client = qianfan.ChatCompletion(ak=access_key, sk=secret_key)
 
     def predict(self, message, stream=False):
-
-        logger.debug(message)
         if isinstance(message, str):
             messages=[{"role": "user", "content": message}]
         else:

src/models/embedding.py

@@ -1,4 +1,4 @@
-from FlagEmbedding import FlagModel
+from FlagEmbedding import FlagModel, FlagReranker
 
 from utils.logging_config import setup_logger
 
@@ -6,11 +6,15 @@ from utils.logging_config import setup_logger
 logger = setup_logger("EmbeddingModel")
 
 SUPPORT_LIST = {
-    "bge-large-zh": "BAAI/bge-large-zh-v1.5",
+    "bge-large-zh-v1.5": "BAAI/bge-large-zh-v1.5",
+}
+
+RERANKER_LIST = {
+    "bge-reranker-v2-m3": "BAAI/bge-reranker-v2-m3",
 }
 
 QUERY_INSTRUCTION = {
-    "bge-large-zh": "为这个句子生成表示以用于检索相关文章：",
+    "bge-large-zh-v1.5": "为这个句子生成表示以用于检索相关文章：",
 }
 
 class EmbeddingModel(FlagModel):
@@ -25,4 +29,16 @@ class EmbeddingModel(FlagModel):
                 query_instruction_for_retrieval=QUERY_INSTRUCTION[config.embed_model],
                 use_fp16=False, **kwargs)
 
-        logger.info(f"Embedding model {config.embed_model} loaded")
+        logger.info(f"Embedding model {config.embed_model} loaded")
+
+
+class ReRanker(FlagReranker):
+    def __init__(self, config, **kwargs):
+
+        assert config.reranker in RERANKER_LIST.keys(), f"Unsupported ReRanker: {config.reranker}, only support {RERANKER_LIST.keys()}"
+
+        model_name_or_path = config.model_local_paths.get(config.reranker, RERANKER_LIST[config.reranker])
+        logger.info(f"Loading ReRanker model {config.re_ranker} from {model_name_or_path}")
+
+        super().__init__(model_name_or_path, use_fp16=True, **kwargs)
+        logger.info(f"ReRanker model {config.re_ranker} loaded")

src/plugins/pdf2txt.py

@@ -7,7 +7,7 @@ from copy import deepcopy
 from tqdm import tqdm
 
 
-def pdf2txt(pdf_path):
+def pdf2txt(pdf_path, return_text=False):
     output_dir = os.path.join('tmp', 'pdf2txt', os.path.basename(pdf_path).split('.')[0])
     os.makedirs(output_dir, exist_ok=True)
 
@@ -63,6 +63,9 @@ def pdf2txt(pdf_path):
     with open(respath, 'w', encoding='utf-8') as f:
         f.write(whole_text)
 
+    if return_text:
+        return whole_text
+
     return respath
 
 if __name__ == "__main__":

src/utils/logging_config.py

@@ -6,7 +6,7 @@ from datetime import datetime
 # DATETIME = datetime.now().strftime('%Y-%m-%d-%H%M%S')
 DATETIME = "debug" # 为了方便，调试的时候输出到 debug.log 文件
 
-def setup_logger(name, log_file=None, level=logging.DEBUG, console=True):
+def setup_logger(name, log_file=None, level=logging.DEBUG, console=False):
 
     if log_file is None:
         log_file = f'output/log/project-{DATETIME}.log'

src/views/common_view.py

@@ -3,11 +3,13 @@ from flask import Blueprint, jsonify, request, Response
 
 from core import HistoryManager
 from utils.logging_config import setup_logger
-from core.startup import config, model, retriever
+from core.startup import config, model
+from core.retriever import Retriever
 
 
 common = Blueprint('common', __name__)
 logger = setup_logger("server-common")
+retriever = Retriever(config)
 
 @common.route('/', methods=["GET"])
 def route_index():
@@ -29,11 +31,12 @@ def chat_get():
 def chat():
     request_data = json.loads(request.data)
     query = request_data['query']
+    meta = request_data.get('meta')
     logger.debug(f"Web query: {query}")
-
-    new_query, refs = retriever(query)
-
     history_manager = HistoryManager(request_data['history'])
+
+    new_query, refs = retriever(query, history_manager.messages, meta)
+
     messages = history_manager.get_history_with_msg(new_query)
     history_manager.add_user(query)
     logger.debug(f"Web history: {history_manager}")
@@ -55,9 +58,13 @@ def chat():
 def call():
     request_data = json.loads(request.data)
     query = request_data['query']
-    logger.debug(f"Web query: {query}")
     response = model.predict(query)
+    logger.debug(f"Call query: {query} Response: {response.content}")
 
     return jsonify({
         "response": response.content,
-    })
+    })
+
+@common.route('/config', methods=['get'])
+def get_config():
+    return jsonify(config)

src/views/database_view.py

@@ -5,7 +5,7 @@ from flask import Blueprint, jsonify, request, Response
 
 from core import HistoryManager
 from utils.logging_config import setup_logger
-from core.startup import config, model, retriever, dbm
+from core.startup import config, model, dbm
 
 db = Blueprint('database', __name__, url_prefix="/database")
 
@@ -28,6 +28,15 @@ def create_database():
     database = dbm.create_database(database_name, description, db_type)
     return jsonify(database)
 
+# TODO: 删除数据库
+@db.route('/', methods=['DELETE'])
+def delete_database():
+    data = json.loads(request.data)
+    db_id = data.get('db_id')
+    logger.debug(f"Delete database {db_id}")
+    dbm.delete_database(db_id)
+    return jsonify({"message": "删除成功"})
+
 
 @db.route('/add_by_file', methods=['POST'])
 def create_document_by_file():
@@ -53,9 +62,6 @@ def get_database_info():
 
     return jsonify(database)
 
-@db.route('/info', methods=['DELETE'])
-def delete_database():
-    return jsonify({"message": "unimplemented"}), 501
 
 @db.route('/document', methods=['DELETE'])
 def delete_document():

web/package.json

@@ -12,6 +12,8 @@
   },
   "dependencies": {
     "@ant-design/icons-vue": "^6.1.0",
+    "@antv/g6": "^5.0.9",
+    "@vueuse/core": "^10.11.0",
     "ant-design-vue": "^4.2.3",
     "axios": "^1.3.4",
     "d3": "^7.8.3",

web/src/components/ChatComponent.vue

@@ -19,7 +19,56 @@
         </div>
       </div>
       <div class="header__right">
-        <div class="nav-btn text" @click="myAlert('未开发')">张文杰</div>
+        <!-- <div class="nav-btn text metas">
+          <CompassFilled v-if="meta.use_web" />
+          <GoldenFilled v-if="meta.use_graph"/>
+        </div> -->
+        <a-dropdown v-if="state.selectedKB !== null">
+          <a class="ant-dropdown-link nav-btn text" @click.prevent>
+            <component :is="state.selectedKB === null ? BookOutlined : BookFilled" />&nbsp;
+            {{ state.selectedKB === null ? '未选择' : state.databases[state.selectedKB]?.name }}
+          </a>
+          <template #overlay>
+            <a-menu>
+              <a-menu-item v-for="(db, index) in state.databases" :key="index" @click="state.selectedKB=index">
+                <a href="javascript:;">{{ db.name }}</a>
+              </a-menu-item>
+              <a-menu-item  @click="state.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>
+        <div v-if="state.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;
+                  {{ state.selectedKB === null ? '未选择' : state.databases[state.selectedKB]?.name }}
+                </a>
+                <template #overlay>
+                  <a-menu>
+                    <a-menu-item v-for="(db, index) in state.databases" :key="index" @click="state.selectedKB=index">
+                      <a href="javascript:;">{{ db.name }}</a>
+                    </a-menu-item>
+                    <a-menu-item  @click="state.selectedKB = null">
+                      <a href="javascript:;">不使用</a>
+                    </a-menu-item>
+                  </a-menu>
+                </template>
+              </a-dropdown>
+            </div>
+          </div>
+          <div class="graphbase flex-center" @click="meta.use_graph = !meta.use_graph">
+            图数据库 <div @click.stop><a-switch v-model:checked="meta.use_graph" /></div>
+          </div>
+          <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>
       </div>
     </div>
     <div v-if="conv.messages.length == 0" class="chat-examples">
@@ -43,16 +92,26 @@
         :class="message.role"
       >
         <p v-if="message.role=='sent'" style="white-space: pre-line" class="message-text">{{ message.text }}</p>
-        <p v-else v-html="renderMarkdown(message.text)" class="message-md" ></p>
+        <p v-else
+          v-html="renderMarkdown(message.text)"
+          class="message-md"
+          @click="consoleMsg(message)"></p>
       </div>
     </div>
     <div class="input-box">
-      <input
+      <a-textarea
+        class="user-input"
+        v-model:value="conv.inputText"
+        @keydown="handleKeyDown"
+        placeholder="输入问题……"
+        :auto-size="{ minRows: 1, maxRows: 10 }"
+      />
+      <!-- <input
         class="user-input"
         v-model="conv.inputText"
         @keydown.enter="sendMessage"
         placeholder="输入问题……"
-      />
+      /> -->
       <a-button size="large" @click="sendMessage" :disabled="(!conv.inputText && !isStreaming)">
         <template #icon> <SendOutlined v-if="!isStreaming" /> <LoadingOutlined v-else/> </template>
       </a-button>
@@ -62,8 +121,18 @@
 </template>
 
 <script setup>
-import { reactive, ref, onMounted, toRefs } from 'vue'
+import { reactive, ref, onMounted, toRefs, nextTick, computed } from 'vue'
-import { SendOutlined, MenuOutlined, FormOutlined, LoadingOutlined } from '@ant-design/icons-vue'
+import { onClickOutside } from '@vueuse/core'
+import {
+  SendOutlined,
+  MenuOutlined,
+  FormOutlined,
+  LoadingOutlined,
+  BookOutlined,
+  BookFilled,
+  CompassFilled,
+  GoldenFilled,
+} from '@ant-design/icons-vue'
 import { marked } from 'marked';
 
 const props = defineProps({
@@ -76,14 +145,22 @@ const emit = defineEmits(['renameTitle'])
 const { conv, state } = toRefs(props)
 const chatBox = ref(null)
 const isStreaming = ref(false)
+const panel = ref(null)
 const examples = ref([
   '写一个冒泡排序',
-  '介绍一下 MECT',
+  '肉碱是什么？',
   '介绍一下江南大学',
   'A大于B，B小于C，A和C哪个大？',
   '今天天气怎么样？'
 ])
 
+const meta = reactive({
+  db_name: computed(() => state.value.databases[state.value.selectedKB]?.metaname),
+  use_graph: false,
+  use_search: false,
+  graph_name: "neo4j",
+})
+
 marked.setOptions({
   gfm: true,
   breaks: true,
@@ -91,6 +168,29 @@ marked.setOptions({
   // 更多选项可以在 marked 文档中找到：https://marked.js.org/
 });
 
+onClickOutside(panel, () => setTimeout(() => state.value.showPanel = false, 30))
+
+const handleKeyDown = (e) => {
+  if (e.key === 'Enter' && !e.shiftKey) {
+    e.preventDefault()
+    sendMessage()
+    console.log('Enter')
+  } else if (e.key === 'Enter' && e.shiftKey) {
+    console.log('Shift + Enter')
+    // Insert a newline character at the current cursor position
+    const textarea = e.target;
+    const start = textarea.selectionStart;
+    const end = textarea.selectionEnd;
+    conv.value.inputText.value =
+      conv.value.inputText.value.substring(0, start) +
+      '\n' +
+      conv.value.inputText.value.substring(end);
+    nextTick(() => {
+      textarea.setSelectionRange(start + 1, start + 1);
+    });
+  }
+}
+
 const renameTitle = () => {
   const prompt = '请用一个很短的句子关于下面的对话内容的主题起一个名字，不要带标点符号：'
   const firstUserMessage = conv.value.messages[0].text
@@ -101,13 +201,8 @@ const renameTitle = () => {
   })
 }
 
-const myAlert = (message) => {
+const myAlert = (message) => alert(message)
-  alert(message)
+const renderMarkdown = (text) => marked(text)
-}
-
-const renderMarkdown = (text) => {
-  return marked(text)
-}
 
 const scrollToBottom = () => {
   setTimeout(() => {
@@ -115,6 +210,8 @@ const scrollToBottom = () => {
   }, 10)
 }
 
+const consoleMsg = (message) => console.log(message)
+
 const generateRandomHash = (length) => {
     let chars = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789';
     let hash = '';
@@ -192,6 +289,7 @@ const sendMessage = () => {
       body: JSON.stringify({
         query: user_input,
         history: conv.value.history,
+        meta: meta
       }),
       headers: {
         'Content-Type': 'application/json'
@@ -226,7 +324,7 @@ const sendMessage = () => {
             conv.value.history = data.history
             buffer = ''
           } catch (e) {
-            console.log(e)
+            // console.log(e)
           }
           return readChunk()
         })
@@ -274,10 +372,19 @@ onMounted(() => {
   padding: 1rem;
 }
 
-.chat div.header .header__left {
+.chat div.header {
-  display: flex;
+  user-select: none;
-  align-items: center;
+
-  gap: 1rem;
+  .header__left, .header__right {
+    display: flex;
+    align-items: center;
+    gap: 1rem;
+  }
+}
+
+.ant-dropdown-link {
+  color: var(--c-text-light-1);
+  cursor: pointer;
 }
 
 .nav-btn {
@@ -302,6 +409,43 @@ onMounted(() => {
   }
 }
 
+.metas {
+  display: flex;
+  gap: 8px;
+}
+
+.my-panal {
+  position: absolute;
+  top: 100%;
+  right: 0;
+  margin-top: 5px;
+  background-color: white;
+  border: 1px solid #ccc;
+  box-shadow: 0px 0px 10px 1px rgba(0, 0, 0, 0.05);
+  border-radius: 12px;
+  padding: 12px;
+  z-index: 100;
+  width: 250px;
+
+  .flex-center {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    gap: 10px;
+  }
+
+  .graphbase {
+    padding: 8px 16px;
+    border-radius: 12px;
+    cursor: pointer;
+    transition: background-color 0.3s;
+
+    &:hover {
+      background-color: #EAEAEA;
+    }
+  }
+}
+
 
 div.chat-examples {
   padding: 0 50px;
@@ -407,14 +551,14 @@ img.message-image {
   max-width: calc(1100px - 2rem);
   margin: 0 auto;
   display: flex;
-  align-items: center;
+  align-items: flex-end;
   background-color: #F4F4F4;
   border-radius: 2rem;
-  height: 3.5rem;
+  height: auto;
   padding: 0.5rem;
 }
 
-input.user-input {
+.user-input {
   flex: 1;
   height: 40px;
   padding: 0.5rem 1rem;
@@ -425,9 +569,15 @@ input.user-input {
   color: #111111;
   font-size: 16px;
   font-variation-settings: 'wght' 400, 'opsz' 10.5;
+  outline: none;
 
   &:focus {
     outline: none;
+    box-shadow: none;
+  }
+
+  &:active {
+    outline: none;
   }
 }
 

web/src/router/index.js

@@ -45,6 +45,12 @@ const router = createRouter({
           path: ':database_id',
           name: 'databaseInfo',
           component: () => import('../views/DataBaseInfoView.vue'),
+        },
+        {
+          path: 'graph',
+          name: 'graph',
+          component: () => import('../views/GraphView.vue'),
+          meta: { keepAlive: true }
         }
       ]
     },

web/src/views/ChatView.vue

@@ -39,6 +39,9 @@ const convs = reactive(JSON.parse(localStorage.getItem('chat-convs')) || [
 
 const state = reactive({
   isSidebarOpen: true,
+  selectedKB: null,
+  showPanel: false,
+  databases: [],
 })
 
 const curConvId = ref(0)
@@ -86,6 +89,18 @@ 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,
@@ -97,6 +112,7 @@ 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++) {
@@ -167,6 +183,7 @@ onMounted(() => {
   padding: 16px;
   cursor: pointer;
   width: 100%;
+  user-select: none;
 
   &__title {
     white-space: nowrap; /* 禁止换行 */

web/src/views/DataBaseInfoView.vue

@@ -1,15 +1,22 @@
 <template>
   <div style="display: flex;">
   <div class="sider">
-    <a-button type="link" @click="backToDatabase"><LeftOutlined />返回知识库</a-button>
+    <div class="sider-top">
-    <div class="top">
+      <div class="header-actions">
-      <div class="icon"><ReadFilled /></div>
+        <a-button type="text" @click="backToDatabase"><LeftOutlined /></a-button>
-      <div class="info">
+        <a-button type="text" danger class="del-db" @click="deleteDatabse"><DeleteOutlined /></a-button>
-        <h3>{{ database.name }}</h3>
-        <p><span>{{ database.metaname }}</span> · <span>{{ database.metadata?.row_count }}</span></p>
       </div>
+      <div class="top">
+        <div class="icon"><ReadFilled /></div>
+        <div class="info">
+          <h3>{{ database.name }}</h3>
+          <p><span>{{ database.metaname }}</span> · <span>{{ database.metadata?.row_count }}</span></p>
+        </div>
+      </div>
+      <p class="description">{{ database.description }}</p>
+    </div>
+    <div class="sider-bottom">
     </div>
-    <p class="description">{{ database.description }}</p>
   </div>
   <div class="db-info-container">
     <h2>向知识库中添加文件</h2>
@@ -59,7 +66,11 @@
           <ClockCircleFilled style="color: #FFCD43;"/>
         </template>
         <template v-else-if="column.key === 'action'">
-          <a-button class="del-btn" type="link" @click="deleteFile(text)" :disabled="state.lock">删除</a-button>
+          <a-button class="del-btn" type="link"
+            @click="deleteFile(text)"
+            :disabled="state.lock || record.status != 'done' "
+            >删除
+          </a-button>
         </template>
         <span v-else-if="column.key === 'created_at'">{{ formatRelativeTime(Math.round(text*1000)) }}</span>
         <span v-else>{{ text }}</span>
@@ -84,7 +95,7 @@
 
 <script setup>
 import { onMounted, reactive, ref, watch } from 'vue';
-import { message } from 'ant-design-vue';
+import { message, Modal } from 'ant-design-vue';
 import { useRoute, useRouter } from 'vue-router';
 import {
   ReadFilled,
@@ -93,6 +104,7 @@ import {
   HourglassFilled,
   CloseCircleFilled,
   ClockCircleFilled,
+  DeleteOutlined,
 } from '@ant-design/icons-vue'
 
 
@@ -109,6 +121,7 @@ const state = reactive({
   refrashing: false,
   lock: false,
   drawer: false,
+  refreshInterval: null,
 });
 
 const handleFileUpload = (event) => {
@@ -139,6 +152,41 @@ const handleRefresh = () => {
   })
 }
 
+const deleteDatabse = () => {
+
+  Modal.confirm({
+    title: '删除数据库',
+    content: '确定要删除该数据库吗？',
+    okText: '确认',
+    cancelText: '取消',
+    onOk: () => {
+      state.lock = true
+      fetch('/api/database/', {
+        method: "DELETE",
+        body: JSON.stringify({
+          db_id: databaseId.value
+        }),
+      })
+      .then(response => response.json())
+      .then(data => {
+        console.log(data)
+        message.success(data.message)
+        router.push('/database')
+      })
+      .catch(error => {
+        console.error(error)
+        message.error(error.message)
+      })
+      .finally(() => {
+        state.lock = false
+      })
+    },
+    onCancel: () => {
+      console.log('Cancel');
+    },
+  });
+}
+
 const openFileDetail = (record) => {
   state.lock = true
   fetch(`/api/database/document?db_id=${databaseId.value}&file_id=${record.file_id}`, {
@@ -230,7 +278,7 @@ const addDocumentByFile = () => {
 
   state.loading = true
   state.lock = true
-  const refreshInterval = setInterval(() => {
+  state.refreshInterval = setInterval(() => {
     getDatabaseInfo();
   }, 1000);
   fetch('/api/database/add_by_file', {
@@ -252,7 +300,7 @@ const addDocumentByFile = () => {
     })
     .finally(() => {
       getDatabaseInfo()
-      clearInterval(refreshInterval)
+      clearInterval(state.refreshInterval)
       state.loading = false
     })
 }
@@ -293,7 +341,8 @@ const columns = [
 watch(() => route.params.database_id, (newId) => {
     databaseId.value = newId;
     console.log(newId)
-    getDatabaseInfo();
+    clearInterval(state.refreshInterval)
+    getDatabaseInfo()
   }
 );
 
@@ -308,13 +357,34 @@ onMounted(() => {
 
 <style lang="less" scoped>
 .sider {
+  display: flex;
+  flex-direction: column;
+  justify-content: space-between;
   width: 300px;
   height: 100%;
-  padding: 20px;
+  padding: 0;
   border-right: 1px solid #E0EAFF;
 
-  button {
+  .sider-top {
-    margin-bottom: 20px;
+
+    & > * {
+      padding: 0 20px;
+    }
+    .header-actions {
+      display: flex;
+      justify-content: space-between;
+      margin-bottom: 20px;
+      padding-top: 10px;
+      padding-bottom: 10px;
+      background-color: #FAFAFA;
+      border-bottom: 1px solid #E0EAFF;
+
+      button {
+        height: auto;
+        font-size: 16px;
+        color: var(--c-text-light-1);
+      }
+    }
   }
 }
 
@@ -394,9 +464,16 @@ onMounted(() => {
     }
   }
 
-  button.del-btn:hover {
+  button.del-btn {
     cursor: pointer;
-    color: var(--error-color);
+
+    &:hover {
+      color: var(--error-color);
+    }
+    &:disabled {
+      cursor: not-allowed;
+    }
   }
+
 }
 </style>

web/src/views/DataBaseView.vue

@@ -45,7 +45,7 @@
     <h2>图数据库</h2>
     <p>基于 neo4j 构建的图数据库。</p>
     <div :class="{'graphloading': graphloading}">
-      <div class="dbcard graphbase" >
+      <div class="dbcard graphbase" @click="navigateToGraph">
         <div class="top">
           <div class="icon"><AppstoreFilled /></div>
           <div class="info">
@@ -64,11 +64,12 @@
 </template>
 
 <script setup>
-import { ref, onMounted, reactive } from 'vue'
+import { ref, onMounted, reactive, watch } from 'vue'
-import { useRouter } from 'vue-router';
+import { useRouter, useRoute } from 'vue-router';
 import { message, Button } from 'ant-design-vue'
 import { ReadFilled, PlusOutlined, AppstoreFilled } from '@ant-design/icons-vue'
 
+const route = useRoute()
 const router = useRouter()
 const databases = ref([])
 const graph = ref(null)
@@ -80,7 +81,7 @@ const newDatabase = reactive({
   loading: false,
 })
 
-const load_databases = () => {
+const loadDatabases = () => {
   loadGraph()
   fetch('/api/database/', {
     method: "GET",
@@ -112,7 +113,7 @@ const createDatabase = () => {
   .then(response => response.json())
   .then(data => {
     console.log(data)
-    load_databases()
+    loadDatabases()
     newDatabase.open = false
     newDatabase.name = ''
     newDatabase.description = ''
@@ -126,6 +127,10 @@ const navigateToDatabase = (databaseId) => {
   router.push({ path: `/database/${databaseId}` });
 };
 
+const navigateToGraph = () => {
+  router.push({ path: `/database/graph` });
+};
+
 const loadGraph = () => {
   graphloading.value = true
   fetch('/api/database/graph', {
@@ -144,8 +149,14 @@ const loadGraph = () => {
     })
 }
 
+watch(() => route.path, (newPath, oldPath) => {
+  if (newPath === '/database') {
+    loadDatabases();
+  }
+});
+
 onMounted(() => {
-  load_databases()
+  loadDatabases()
 })
 
 </script>

web/src/views/GraphView.vue

@@ -0,0 +1,48 @@
+<template>
+  <div class="graph-container">
+    <div class="main" id="container"></div>
+  </div>
+</template>
+
+<script setup>
+import { Graph } from "@antv/g6";
+import { onMounted } from 'vue';
+
+const getCurWidth = () => document.getElementById("container").offsetWidth
+const getCurHeight = () => document.getElementById("container").offsetHeight
+
+
+onMounted(() => {
+  const graph = new Graph({
+    container: document.getElementById("container"),
+    width: getCurWidth(),
+    height: getCurHeight(),
+    data: {
+      nodes: [
+        {
+          id: "node-1",
+          style: { x: 50, y: 100 },
+        },
+        {
+          id: "node-2",
+          style: { x: 150, y: 100 },
+        },
+      ],
+      edges: [{ id: "edge-1", source: "node-1", target: "node-2" }],
+    },
+    behaviors: ["drag-canvas", "zoom-canvas", "drag-element"],
+  });
+
+  graph.render();
+});
+</script>
+
+<style scoped>
+.graph-container {}
+
+#container {
+  width: 100%;
+  height: 100%;
+}
+
+</style>

web/vite.config.js

@@ -20,5 +20,8 @@ export default defineConfig({
         rewrite: (path) => path.replace(/^\/api/, '')
       }
     },
+    watch: {
+      ignored: ['**/node_modules/**', '**/dist/**'],
+    },
   }
 })