support for knowledge base RAG and fix UI bug

.gitignore

@@ -22,3 +22,4 @@ logs
 
 ### IDE
 .vscode
+*.nogit.*

README.md

@@ -1,5 +1,7 @@
 ## Project: Athena
 
+![home](web/public/home.png)
+
 ### 准备
 
 1. 提供 API 服务商的 API_KEY，并放置在 `src/.env` 文件中，参考 `src/.env.template`。默认使用的是智谱AI。

src/.env.template

@@ -1,3 +1,6 @@
 DIFYAPI=optional
 DEEPSEEKAPI=optional
-ZHIPUAPI=optional
+ZHIPUAPI=optional
+QIANFAN_ACCESS_KEY=optional
+QIANFAN_SECRET_KEY=optional
+HF_ENDPOINT='https://hf-mirror.com'

src/api.py

@@ -3,7 +3,8 @@ import json
 from flask import Flask, jsonify, Response, request
 from flask_cors import CORS
 from dotenv import load_dotenv
-from core.history import HistoryManager
+from core import HistoryManager
+from core import PreRetrival
 from config import Config
 from models import select_model
 from utils.logging_config import setup_logger
@@ -12,14 +13,15 @@ from utils.logging_config import setup_logger
 load_dotenv()
 logger = setup_logger("server")
 
-apps = Flask(__name__)# 这段代码是为了解决跨域问题，Flask默认不支持跨域
-CORS(apps, resources=r'/*')# CORS的用法是
-
 
 config = Config("config/base.yaml")
 model = select_model(config)
+pre_retrival = PreRetrival(config)
 
 
+apps = Flask(__name__)# 这段代码是为了解决跨域问题，Flask默认不支持跨域
+CORS(apps, resources=r'/*')# CORS的用法是
+
 @apps.route('/', methods=["GET"])
 def route_index():
     return jsonify({"message": "You Got It!"})
@@ -44,6 +46,18 @@ def chat():
     query = request_data['query']
     logger.debug(f"Web query: {query}")
 
+    external = ""
+    if config.enable_knowledge_base:
+        kb_res = pre_retrival.search(query)
+        if kb_res:
+            kb_res = "\n".join([f"{r['id']}: {r['entity']['text']}" for r in kb_res[0]])
+            kb_res = f"知识库信息: {kb_res}"
+            external += kb_res
+
+    if len(external) > 0:
+        query = f"以下是参考资料：\n\n\n {external} 请根据前面的知识回答：{query}"
+
+
     history_manager = HistoryManager(request_data['history'])
     messages = history_manager.add_user(query)
     logger.debug(f"Web history: {history_manager}")

src/cli.py

@@ -1,6 +1,7 @@
 import os
 from dotenv import load_dotenv
-from core.history import HistoryManager
+from core import HistoryManager
+from core import PreRetrival
 from config import Config
 from models import select_model
 
@@ -10,6 +11,8 @@ load_dotenv()
 if __name__ == "__main__":
     config = Config("config/base.yaml")
     model = select_model(config)
+    pre_retrival = PreRetrival(config)
+    # pre_retrival.add_file("/home/zwj/workspace/ProjectAthena/src/data/file/鉴定工作报告、技术报告-0708.pdf")
 
     print(f"[{config.model_provider}:{config.get('model_name', 'default')}] Type 'exit' to quit")
 
@@ -19,6 +22,18 @@ if __name__ == "__main__":
         if message == "exit":
             break
 
+        external = ""
+
+        if config.enable_knowledge_base:
+            kb_res = pre_retrival.search(message)
+            if kb_res:
+                kb_res = "\n".join([f"{r['id']}: {r['entity']['text']}" for r in kb_res[0]])
+                kb_res = f"知识库信息: {kb_res}"
+                external += kb_res
+
+        if len(external) > 0:
+            message = f"以下是参考资料：\n\n\n {external} 请根据前面的知识回答：{message}"
+
         messages = history_manager.add_user(message)
         response = model.predict(messages, stream=config.stream)
 

src/config/__init__.py

@@ -38,6 +38,15 @@ class Config(SimpleConfig):
         self.stream = False
 
         self.load()
+        self.handle_self()
+
+    def handle_self(self):
+        ### handle local model
+        model_root_dir = os.getenv("MODEL_ROOT_DIR", "pretrained_models")
+        for model, model_rel_path in self.model_local_paths.items():
+            if not model_rel_path.startswith("/"):
+                self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path)
+
 
     def load(self):
         if self.filename is not None and os.path.exists(self.filename):

src/config/base.yaml

@@ -2,7 +2,22 @@ name: base
 
 ## model
 ### model_provider, option in deepseek, zhipu
-model_provider: wenxin
+model_provider: qianfan
+embed_model: bge-large-zh  # option in ["bge-large-zh"]
 
 ## startup
-stream: True
+stream: True
+
+## knowledge
+enable_query_rewrite: True
+enable_knowledge_base: True
+enable_knowledge_graph: True
+enable_search_engine: True
+
+## vector_store
+milvus_local_path: data/vector_base/milvus.db  # 初步
+
+
+## model dir
+model_local_paths:
+  bge-large-zh: bge-large-zh-v1.5

src/core/__init__.py

@@ -0,0 +1,2 @@
+from .history import *
+from .preretrival import *

src/core/preretrival.py

@@ -0,0 +1,114 @@
+# Read Chunking Embedding and save it to Vector Database
+import os
+
+from pathlib import Path
+from llama_index.readers.file import PDFReader
+
+from models.embedding import EmbeddingModel
+from utils.logging_config import setup_logger
+
+from pymilvus import MilvusClient
+
+
+logger = setup_logger("PreRetrival")
+
+def pdfreader(file_path):
+    """读取PDF文件并返回text文本"""
+    assert os.path.exists(file_path), "File not found"
+    assert file_path.endswith(".pdf"), "File format not supported"
+
+    doc = PDFReader().load_data(file=Path(file_path))
+
+    # 简单的拼接起来之后返回纯文本
+    text = "\n\n".join([d.get_content() for d in doc])
+    return text
+
+def plainreader(file_path):
+    """读取普通文本文件并返回text文本"""
+    assert os.path.exists(file_path), "File not found"
+
+    with open(file_path, "r") as f:
+        text = f.read()
+    return text
+
+
+class PreRetrival:
+
+    def __init__(self, config):
+        self.config = config
+        self._init_config(config)
+
+        self.embed_model = EmbeddingModel(config)
+        self.client = MilvusClient(config.milvus_local_path)
+
+    def _init_config(self, config):
+        self.vector_dim = 1024  # 暂时不知道这个和 embedding model 的 embedding 大小有什么关系
+        self.default_query_limit = 2
+        self.default_collection_name = "default"
+
+    def add_file(self, file, collection_name=None):
+        """添加文件到数据库"""
+        collection_name = collection_name or self.default_collection_name
+        text = self.read_text(file)
+        chunks = self.chunking(text)
+
+        self.add_documents(chunks, collection_name)
+
+    def add_documents(self, docs, collection_name):
+        """添加已经分块之后的文本"""
+        vectors = self.embed_model.encode(docs)
+
+        data = [
+            {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
+            for i in range(len(vectors))
+        ]
+
+        # for testing, we drop the collection if it already exists
+        # if self.client.has_collection(collection_name=collection_name):
+        #     self.client.drop_collection(collection_name=collection_name)
+
+        self.client.create_collection(
+            collection_name=collection_name,
+            dimension=self.vector_dim,  # The vectors we will use in this demo has 768 dimensions
+        )
+
+        res = self.client.insert(collection_name=collection_name, data=data)
+        return res
+
+    def search(self, query, collection_name=None, limit=None):
+        collection_name = collection_name or self.default_collection_name
+        limit = limit or self.default_query_limit
+
+        query_vectors = self.embed_model.encode_queries([query])
+
+        res = self.client.search(
+            collection_name=collection_name,  # target collection
+            data=query_vectors,  # query vectors
+            limit=limit,  # number of returned entities
+            output_fields=["text", "subject"],  # specifies fields to be returned
+        )
+
+        return res
+
+    def read_text(self, file):
+        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"):
+                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!")
+
+    def chunking(self, text, chunk_size=1024):
+        """将文本切分成固定大小的块"""
+        chunks = []
+        for i in range(0, len(text), chunk_size):
+            chunks.append(text[i:i + chunk_size])
+        return chunks

src/models/__init__.py

@@ -1,7 +1,3 @@
-from models.chat_model import DeepSeek, Zhipu
-
-
-
 def select_model(config):
 
     model_provider = config.model_provider
@@ -15,9 +11,9 @@ def select_model(config):
         from models.chat_model import Zhipu
         return Zhipu(model_name)
 
-    elif model_provider == "wenxin":
-        from models.chat_model import Wenxin
-        return Wenxin(model_name)
+    elif model_provider == "qianfan":
+        from models.chat_model import Qianfan
+        return Qianfan(model_name)
 
     elif model_provider is None:
         raise ValueError("Model provider not specified, please modify `model_provider` in `src/config/base.yaml`")

src/models/chat_model.py

@@ -60,12 +60,12 @@ class Zhipu(OpenAIBase):
 import qianfan
 
 
-class WenxinResponse:
+class QianfanResponse:
     def __init__(self, content):
         self.content = content
 
 
-class Wenxin:
+class Qianfan:
 
     def __init__(self, model_name="ernie-lite-8k") -> None:
         self.model_name = model_name
@@ -93,7 +93,7 @@ class Wenxin:
             stream=True,
         )
         for chunk in response:
-            yield WenxinResponse(chunk["body"]["result"])
+            yield QianfanResponse(chunk["body"]["result"])
 
     def _get_response(self, messages):
         response = self.client.do(
@@ -101,4 +101,4 @@ class Wenxin:
             messages=messages,
             stream=False,
         )
-        return WenxinResponse(response["body"]["result"])
+        return QianfanResponse(response["body"]["result"])

src/models/embedding.py

@@ -0,0 +1,32 @@
+from FlagEmbedding import FlagModel
+
+from utils.logging_config import setup_logger
+
+
+logger = setup_logger("EmbeddingModel")
+
+SUPPORT_LIST = {
+    "bge-large-zh": "BAAI/bge-large-zh-v1.5",
+}
+
+QUERY_INSTRUCTION = {
+    "bge-large-zh": "为这个句子生成表示以用于检索相关文章：",
+}
+
+class EmbeddingModel(FlagModel):
+    def __init__(self, config, **kwargs):
+
+        assert config.embed_model in SUPPORT_LIST.keys(), f"Unsupported embed model: {config.embed_model}, only support {SUPPORT_LIST.keys()}"
+
+        if config.embed_model in config.model_local_paths.keys():
+            model_name_or_path = config.model_local_paths[config.embed_model]
+        else:
+            model_name_or_path = SUPPORT_LIST[config.embed_model]
+
+        logger.info(f"Loading embedding model {config.embed_model} from {model_name_or_path}")
+
+        super().__init__(model_name_or_path,
+                query_instruction_for_retrieval=QUERY_INSTRUCTION[config.embed_model],
+                use_fp16=False, **kwargs)
+
+        logger.info(f"Embedding model {config.embed_model} loaded")

src/requirements.txt

@@ -1,5 +1,9 @@
+FlagEmbedding==1.2.10
 Flask==3.0.3
 Flask_Cors==4.0.1
+llama_index==0.10.53
 openai==1.35.10
+pymilvus==2.4.4
 python-dotenv==1.0.1
 PyYAML==6.0.1
+qianfan==0.4.0.1

web/src/views/ChatView.vue

@@ -6,7 +6,7 @@
           v-for="message in state.messages"
           :key="message.id"
           class="message-box"
-          :class="message.type"
+          :class="message.role"
         >
           <img v-if="message.filetype === 'image'" :src="message.url" class="message-image" alt="">
           <p v-else style="white-space: pre-line" class="message-text">{{ message.text }}</p>
@@ -97,10 +97,15 @@ const scrollToBottom = () => {
   }, 10) // 10ms 后滚动到底部
 }
 
-const appendMessage = (message, type) => {
+/**
+ *
+ * @param {*} message 消息内容
+ * @param {*} role 消息
+ */
+const appendMessage = (message, role) => {
   state.messages.push({
     id: state.messages.length + 1,
-    type,
+    role,
     text: message
   })
   scrollToBottom()
@@ -196,9 +201,7 @@ const getFormattedDocs = (docs) => {
 }
 
 const sendDeafultMessage = () => {
-  setTimeout(() => {
-    appendMessage('你好？我是 Project: Athena，有什么可以帮你？😊', 'received')
-  }, 1000);
+  appendMessage('你好？我是 Project: Athena，有什么可以帮你？😊', 'received')
 }
 
 const clearChat = () => {
@@ -272,6 +275,7 @@ div.chat, div.info {
   font-weight: 400;
   box-sizing: border-box;
   box-shadow: 0px 0.3px 0.9px rgba(0, 0, 0, 0.12), 0px 1.6px 3.6px rgba(0, 0, 0, 0.16);
+  animation: slideInUp 0.1s ease-in;
 }
 
 .message-box.sent {
@@ -287,6 +291,7 @@ div.chat, div.info {
   color: #111111;
   background-color: #ffffff;
   text-align: left;
+  // animation-delay: 0.2s; /* 延迟 100ms 开始动画 */
 }
 
 p.message-text {
@@ -392,4 +397,15 @@ p.note {
   padding: 1rem;
   color: #ccc;
 }
+
+@keyframes slideInUp {
+  from {
+    opacity: 0;
+    transform: translateY(100%);
+  }
+  to {
+    opacity: 1;
+    transform: translateY(0);
+  }
+}
 </style>