From 523cc63b51b5f2f242c5d074825f5c5d35dd0c11 Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Sun, 14 Jul 2024 23:59:52 +0800 Subject: [PATCH] add database server api partial support --- src/api.py | 87 ++----------- src/cli.py | 2 +- src/core/__init__.py | 4 +- src/core/database.py | 30 +++++ src/core/filereader.py | 25 ++++ src/core/graphbase.py | 13 ++ .../{preretrieval.py => knowledgebase.py} | 122 +++++++++--------- src/core/retriever.py | 9 +- src/core/startup.py | 9 ++ src/models/__init__.py | 5 + src/utils/__init__.py | 7 +- src/views/__init__.py | 14 ++ src/views/common_view.py | 63 +++++++++ src/views/database_view.py | 71 ++++++++++ 14 files changed, 313 insertions(+), 148 deletions(-) create mode 100644 src/core/database.py create mode 100644 src/core/filereader.py create mode 100644 src/core/graphbase.py rename src/core/{preretrieval.py => knowledgebase.py} (53%) create mode 100644 src/core/startup.py create mode 100644 src/views/__init__.py create mode 100644 src/views/common_view.py create mode 100644 src/views/database_view.py diff --git a/src/api.py b/src/api.py index b12d389b..be3f9046 100644 --- a/src/api.py +++ b/src/api.py @@ -1,84 +1,15 @@ -import os -import json -from flask import Flask, jsonify, Response, request -from flask_cors import CORS from dotenv import load_dotenv -from core import HistoryManager -from core import Retriever -from config import Config -from models import select_model +load_dotenv() + +import os +from views import create_app from utils.logging_config import setup_logger +logger = setup_logger("Server") -load_dotenv() -logger = setup_logger("server") - - -config = Config("config/base.yaml") -model = select_model(config) -retriever = Retriever(config) - - -apps = Flask(__name__)# 这段代码是为了解决跨域问题，Flask默认不支持跨域 -CORS(apps, resources=r'/*')# CORS的用法是 - -@apps.route('/', methods=["GET"]) -def route_index(): - return jsonify({"message": "You Got It!"}) - - -@apps.errorhandler(404) -def page_not_found(e): - return jsonify({"message": "DEBUG: " + str(e)}), 404 - - -@apps.errorhandler(403) -def page_not_found(e): - return jsonify({"message": str(e)}), 403 -@apps.route('/', methods=['GET']) -def chat_get(): - return "Chat Get!" - -@apps.route('/chat', methods=['POST']) -def chat(): - request_data = json.loads(request.data) - query = request_data['query'] - logger.debug(f"Web query: {query}") - - new_query, refs = retriever(query) - - history_manager = HistoryManager(request_data['history']) - messages = history_manager.get_history_with_msg(new_query) - history_manager.add_user(query) - logger.debug(f"Web history: {history_manager}") - - def generate_response(): - content = "" - for delta in model.predict(messages, stream=True): - content += delta.content - response_chunk = json.dumps({ - "history": history_manager.update_ai(content), - "response": content, - "refs": refs # TODO: 优化 refs，不需要每次都返回 - }, ensure_ascii=False).encode('utf8') + b'\n' - yield response_chunk - - return Response(generate_response(), content_type='application/json', status=200) - -@apps.route('/call', methods=['POST']) -def call(): - request_data = json.loads(request.data) - query = request_data['query'] - logger.debug(f"Web query: {query}") - response = model.predict(query, stream=False) - - return jsonify({ - "response": response.content, - }) - - +app = create_app() if __name__ == '__main__': - print("Starting model...") - apps.secret_key = os.urandom(24) - apps.run(host='0.0.0.0', port=8000, debug=True, threaded=True) \ No newline at end of file + logger.info("Starting server") + app.secret_key = os.urandom(24) + app.run(host='0.0.0.0', port=5000, debug=False, threaded=True) diff --git a/src/cli.py b/src/cli.py index 1a0fd1bf..816ec1e9 100644 --- a/src/cli.py +++ b/src/cli.py @@ -1,7 +1,7 @@ import os from dotenv import load_dotenv from core import HistoryManager -from core import PreRetrieval, Retriever +from core import Retriever from config import Config from models import select_model diff --git a/src/core/__init__.py b/src/core/__init__.py index 0338e6fe..c6ba8867 100644 --- a/src/core/__init__.py +++ b/src/core/__init__.py @@ -1,3 +1,3 @@ from .history import * -from .preretrieval import * -from .retriever import * \ No newline at end of file +from .retriever import * +from .database import * \ No newline at end of file diff --git a/src/core/database.py b/src/core/database.py new file mode 100644 index 00000000..a2518de1 --- /dev/null +++ b/src/core/database.py @@ -0,0 +1,30 @@ +from core.knowledgebase import KnowledgeBase + + +class DataBaseManager: + + def __init__(self, config=None) -> None: + self.config = config + self.knowledge_base = KnowledgeBase(config) + + def get_databases(self): + kb = self.knowledge_base.get_collections() + return kb + + def create_database(self, collection_name): + self.knowledge_base.add_collection(collection_name) + return self.get_databases() + + def add_file(self, file, collection_name=None): + self.knowledge_base.add_file(file, collection_name) + return self.get_databases() + + def add_text(self, text, collection_name=None): + self.knowledge_base.add_text(text, collection_name) + return self.get_databases() + + def get_database_info(self, database_name): + return self.knowledge_base.get_collection_info(database_name) + + def get_document_info(self, database_name, document_id): + return self.knowledge_base.search_by_id(database_name, document_id) diff --git a/src/core/filereader.py b/src/core/filereader.py new file mode 100644 index 00000000..ae9034dc --- /dev/null +++ b/src/core/filereader.py @@ -0,0 +1,25 @@ +import os + +from pathlib import Path +from llama_index.readers.file import PDFReader + + +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 + diff --git a/src/core/graphbase.py b/src/core/graphbase.py new file mode 100644 index 00000000..4a384f37 --- /dev/null +++ b/src/core/graphbase.py @@ -0,0 +1,13 @@ +import os + + +class GraphBase: + + def __init__(self, config=None) -> None: + self.config = config + + self._init_config(config) + + def _init_config(self, config): + pass + diff --git a/src/core/preretrieval.py b/src/core/knowledgebase.py similarity index 53% rename from src/core/preretrieval.py rename to src/core/knowledgebase.py index f10240c7..e8a8c914 100644 --- a/src/core/preretrieval.py +++ b/src/core/knowledgebase.py @@ -1,89 +1,81 @@ -# Read Chunking Embedding and save it to Vector Database import os -import utils - -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 -from plugins import OneKE, pdf2txt +from utils import setup_logger, hashstr +from core.filereader import pdfreader, plainreader + +logger = setup_logger("KnowledgeBase") -logger = setup_logger("PreRetrieval") +class KnowledgeBase: -def pdfreader(file_path): - """读取PDF文件并返回text文本""" - assert os.path.exists(file_path), "File not found" - assert file_path.endswith(".pdf"), "File format not supported" - - if utils.is_text_pdf(file_path): - doc = PDFReader().load_data(file=Path(file_path)) - text = "\n\n".join([d.get_content() for d in doc]) - else: - text = pdf2txt(file_path) - - 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 PreRetrieval: - - def __init__(self, config): + def __init__(self, config=None) -> None: self.config = config self._init_config(config) self.embed_model = EmbeddingModel(config) self.client = MilvusClient(config.milvus_local_path) - self.oneke = OneKE(config) 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) - # convert text to graph - # self.oneke.processing_text_to_kg(text, none) - chunks = self.chunking(text) + def get_collections(self): + collections_name = self.client.list_collections() + collections = [] + for collection_name in collections_name: + collection = self.get_collection_info(collection_name) + collections.append(collection) - self.add_documents(chunks, collection_name) + return collections - def add_documents(self, docs, collection_name): - """添加已经分块之后的文本""" - vectors = self.embed_model.encode(docs) + def get_collection_info(self, collection_name): + collection = self.client.describe_collection(collection_name) + collection.update(self.client.get_collection_stats(collection_name)) + # collection["id"] = hashstr(collection_name) + return collection - 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) + def add_collection(self, collection_name): + if self.client.has_collection(collection_name=collection_name): + logger.warning(f"Collection {collection_name} already exists, drop it") + 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 ) + def add_file(self, file, collection_name): + """添加文件到数据库""" + + # 检查 collection 是否存在 + if not self.client.has_collection(collection_name=collection_name): + logger.warning(f"Collection {collection_name} not found, create it") + self.add_collection(collection_name) + + text = self.read_text(file) + chunks = self.chunking(text) + + self.add_documents(chunks, collection_name, filename=file) + + def add_text(self, text, collection_name): + """添加文本到数据库""" + chunks = self.chunking(text) + self.add_documents(chunks, collection_name) + + def add_documents(self, docs, collection_name, filename=None): + """添加已经分块之后的文本""" + vectors = self.embed_model.encode(docs) + + data = [ + {"id": i, "vector": vectors[i], "text": docs[i], "filename": filename} + for i in range(len(vectors)) + ] + 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 + def search(self, query, collection_name, limit=None): limit = limit or self.default_query_limit query_vectors = self.embed_model.encode_queries([query]) @@ -97,6 +89,18 @@ class PreRetrieval: return res[0] # 因为 query 只有一个 + def examples(self, collection_name, limit=20): + res = self.client.query( + collection_name=collection_name, + limit=10, + output_fields=["id", "text"], + ) + return res + + def search_by_id(self, collection_name, id, output_fields=["id", "text"]): + res = self.client.get(collection_name, id, output_fields=output_fields) + return res + def read_text(self, file): support_format = [".pdf", ".txt", "*.md"] assert os.path.exists(file), "File not found" @@ -118,4 +122,4 @@ class PreRetrieval: 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 + return chunks diff --git a/src/core/retriever.py b/src/core/retriever.py index 93e56bed..03fa6741 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -1,19 +1,14 @@ -from core import PreRetrieval - class Retriever: def __init__(self, config): self.config = config - self.pre_retrieval = PreRetrieval(config) def retrieval(self, query): refs = {} - # TODO: 查询分类、查询重写、查询分解、查询伪文档生成（HyDE） - - if self.config.enable_knowledge_base: - refs["knowledge_base"] = self.pre_retrieval.search(query) + # TODO: 查询分类、查询重写、查询分解、查询伪文档生成（HyDE）) + # NOTE：2024-07-14 暂时禁用知识检索 return refs diff --git a/src/core/startup.py b/src/core/startup.py new file mode 100644 index 00000000..369b793b --- /dev/null +++ b/src/core/startup.py @@ -0,0 +1,9 @@ +from core import Retriever, DataBaseManager +from models import select_model +from config import Config + + +config = Config("config/base.yaml") +model = select_model(config) +dbm = DataBaseManager(config) +retriever = Retriever(config) \ No newline at end of file diff --git a/src/models/__init__.py b/src/models/__init__.py index e166b145..bae3a78e 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -1,8 +1,13 @@ +from utils.logging_config import logger + + def select_model(config): model_provider = config.model_provider model_name = config.model_name + logger.info(f"Selecting model from {model_provider} with name {model_name}") + if model_provider == "deepseek": from models.chat_model import DeepSeek return DeepSeek(model_name) diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 27224e83..3c2115b0 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -9,4 +9,9 @@ def is_text_pdf(pdf_path): text = page.get_text() if text.strip(): # 检查是否有文本内容 return True - return False \ No newline at end of file + return False + +def hashstr(input_string, length=16): + import hashlib + hash = hashlib.md5(str(input_string).encode()).hexdigest() + return hash[:length] \ No newline at end of file diff --git a/src/views/__init__.py b/src/views/__init__.py new file mode 100644 index 00000000..784eba5c --- /dev/null +++ b/src/views/__init__.py @@ -0,0 +1,14 @@ +from flask import Flask +from flask_cors import CORS +from views.common_view import common +from views.database_view import db + + +def create_app(): + app = Flask(__name__) + CORS(app, resources=r'/*') + + app.register_blueprint(common) + app.register_blueprint(db) + + return app diff --git a/src/views/common_view.py b/src/views/common_view.py new file mode 100644 index 00000000..22670cda --- /dev/null +++ b/src/views/common_view.py @@ -0,0 +1,63 @@ +import json +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 + + +common = Blueprint('common', __name__) +logger = setup_logger("server-common") + +@common.route('/', methods=["GET"]) +def route_index(): + return jsonify({"message": "You Got It!"}) + +@common.errorhandler(404) +def page_not_found(e): + return jsonify({"message": "DEBUG: " + str(e)}), 404 + + +@common.errorhandler(403) +def page_not_found(e): + return jsonify({"message": str(e)}), 403 +@common.route('/', methods=['GET']) +def chat_get(): + return "Chat Get!" + +@common.route('/chat', methods=['POST']) +def chat(): + request_data = json.loads(request.data) + query = request_data['query'] + logger.debug(f"Web query: {query}") + + new_query, refs = retriever(query) + + history_manager = HistoryManager(request_data['history']) + messages = history_manager.get_history_with_msg(new_query) + history_manager.add_user(query) + logger.debug(f"Web history: {history_manager}") + + def generate_response(): + content = "" + for delta in model.predict(messages, stream=True): + content += delta.content + response_chunk = json.dumps({ + "history": history_manager.update_ai(content), + "response": content, + "refs": refs # TODO: 优化 refs，不需要每次都返回 + }, ensure_ascii=False).encode('utf8') + b'\n' + yield response_chunk + + return Response(generate_response(), content_type='application/json', status=200) + +@common.route('/call', methods=['POST']) +def call(): + request_data = json.loads(request.data) + query = request_data['query'] + logger.debug(f"Web query: {query}") + response = model.predict(query, stream=False) + + return jsonify({ + "response": response.content, + }) \ No newline at end of file diff --git a/src/views/database_view.py b/src/views/database_view.py new file mode 100644 index 00000000..214e01e6 --- /dev/null +++ b/src/views/database_view.py @@ -0,0 +1,71 @@ +import os +import json +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 + +db = Blueprint('database', __name__, url_prefix="/database") + +logger = setup_logger("server-database") + +@db.route('/', methods=['GET']) +def get_databases(): + database = dbm.get_databases() + return jsonify(database) + +@db.route('/', methods=['POST']) +def create_database(): + database_name = request.args.get('database_name') + logger.debug(f"Create database {database_name}") + database = dbm.create_database(database_name) + return jsonify(database) + + +@db.route('/add_by_text', methods=['POST']) +def create_document_by_text(): + form_data = request.form + name = form_data['name'] + text = form_data['text'] + logger.debug(f"Add document in {name} by text: {text}") + database = dbm.add_text(text, name) + return jsonify(database) + +@db.route('/add_by_file', methods=['POST']) +def create_document_by_file(): + """file Type: multipart/form-data""" + name = request.form['database_name'] + file = request.files['file'] + logger.debug(f"Adding document {name} by file: {file.filename}") + if file: + file_path = os.path.join("data/uploads", file.filename) + file.save(file_path) + logger.info(f"Save file {file_path}") + else: + logger.error("No file found") + return jsonify({"message": "No file found"}), 400 + + database = dbm.add_file(file_path, name) + return jsonify(database) + + +@db.route('/info', methods=['GET']) +def get_database_info(): + name = request.args.get('database_name') + logger.debug(f"Get database {name} info") + database = dbm.get_database_info(name) + return jsonify(database) + +@db.route('/info', methods=['DELETE']) +def delete_database(): + return jsonify({"message": "unimplemented"}), 501 + +@db.route('/document', methods=['GET']) +def get_document_info(): + name = request.args.get('database_name') + id = request.args.get('id') + logger.debug(f"Get document {id} in {name}") + document = dbm.get_document_info(name, id) + return jsonify(document) +