diff --git a/.gitignore b/.gitignore index 88bf7442..990ebc34 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,4 @@ logs ### IDE .vscode +*.nogit.* \ No newline at end of file diff --git a/README.md b/README.md index b6dfe59f..5f47a457 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ ## Project: Athena +![home](web/public/home.png) + ### 准备 1. 提供 API 服务商的 API_KEY，并放置在 `src/.env` 文件中，参考 `src/.env.template`。默认使用的是智谱AI。 diff --git a/src/.env.template b/src/.env.template index 816e4113..c172cd7b 100644 --- a/src/.env.template +++ b/src/.env.template @@ -1,3 +1,6 @@ DIFYAPI=optional DEEPSEEKAPI=optional -ZHIPUAPI=optional \ No newline at end of file +ZHIPUAPI=optional +QIANFAN_ACCESS_KEY=optional +QIANFAN_SECRET_KEY=optional +HF_ENDPOINT='https://hf-mirror.com' \ No newline at end of file diff --git a/src/api.py b/src/api.py index d28d99b0..e559a626 100644 --- a/src/api.py +++ b/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}") diff --git a/src/cli.py b/src/cli.py index 0edf9352..11663bb4 100644 --- a/src/cli.py +++ b/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) diff --git a/src/config/__init__.py b/src/config/__init__.py index 352cfd84..bd06628c 100644 --- a/src/config/__init__.py +++ b/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): diff --git a/src/config/base.yaml b/src/config/base.yaml index ba90d2ce..233065be 100644 --- a/src/config/base.yaml +++ b/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 \ No newline at end of file +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 \ No newline at end of file diff --git a/src/core/__init__.py b/src/core/__init__.py index e69de29b..414a12cc 100644 --- a/src/core/__init__.py +++ b/src/core/__init__.py @@ -0,0 +1,2 @@ +from .history import * +from .preretrival import * \ No newline at end of file diff --git a/src/core/preretrival.py b/src/core/preretrival.py new file mode 100644 index 00000000..51e73bab --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/core/retrieval.py b/src/core/retrieval.py new file mode 100644 index 00000000..e69de29b diff --git a/src/models/__init__.py b/src/models/__init__.py index d2bfca61..e166b145 100644 --- a/src/models/__init__.py +++ b/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`") diff --git a/src/models/chat_model.py b/src/models/chat_model.py index 56d8f8bf..8df94f7b 100644 --- a/src/models/chat_model.py +++ b/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"]) \ No newline at end of file + return QianfanResponse(response["body"]["result"]) \ No newline at end of file diff --git a/src/models/embedding.py b/src/models/embedding.py new file mode 100644 index 00000000..f4d7924a --- /dev/null +++ b/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") \ No newline at end of file diff --git a/src/requirements.txt b/src/requirements.txt index b7205d79..32f414e5 100644 --- a/src/requirements.txt +++ b/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 diff --git a/web/public/home.png b/web/public/home.png index 74d4aef7..9c87dbf5 100644 Binary files a/web/public/home.png and b/web/public/home.png differ diff --git a/web/src/views/ChatView.vue b/web/src/views/ChatView.vue index 2ce3fc5f..95362ead 100644 --- a/web/src/views/ChatView.vue +++ b/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" >

{{ message.text }}

@@ -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); + } +}