add database server api partial support

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)
+    logger.info("Starting server")
+    app.secret_key = os.urandom(24)
+    app.run(host='0.0.0.0', port=5000, debug=False, threaded=True)

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
 

src/core/__init__.py

@@ -1,3 +1,3 @@
 from .history import *
-from .preretrieval import *
-from .retriever import *
+from .retriever import *
+from .database import *

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)

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
+

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
+

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
+        return chunks

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
 

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)

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)

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
+    return False
+
+def hashstr(input_string, length=16):
+    import hashlib
+    hash = hashlib.md5(str(input_string).encode()).hexdigest()
+    return hash[:length]

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

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,
+    })

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)
+