From 7e1472a95b55cd038754731916c9f67484aabe98 Mon Sep 17 00:00:00 2001 From: imClumsyPanda Date: Sat, 5 Aug 2023 00:10:10 +0800 Subject: [PATCH] add KnowledgeBase class --- server/knowledge_base/knowledge_base.py | 18 ++++++++++++++++++ server/knowledge_base/utils.py | 2 -- 2 files changed, 18 insertions(+), 2 deletions(-) create mode 100644 server/knowledge_base/knowledge_base.py diff --git a/server/knowledge_base/knowledge_base.py b/server/knowledge_base/knowledge_base.py new file mode 100644 index 0000000..65f9bad --- /dev/null +++ b/server/knowledge_base/knowledge_base.py @@ -0,0 +1,18 @@ +from server.knowledge_base.utils import (get_vs_path, get_kb_path, get_doc_path) + +SUPPORTED_VS_TYPES = ["faiss", "milvus"] + + +class KnowledgeBase: + def __init__(self, + knowledge_base_name: str, + vector_store_type: str, + ): + self.kb_name = knowledge_base_name + if vector_store_type not in SUPPORTED_VS_TYPES: + raise ValueError(f"暂未支持向量库类型 {vector_store_type}") + self.vs_type = vector_store_type + self.kb_path = get_kb_path(self.kb_name) + self.doc_path = get_doc_path(self.kb_name) + if self.vs_type in ["faiss"]: + self.vs_path = get_vs_path(self.kb_name) \ No newline at end of file diff --git a/server/knowledge_base/utils.py b/server/knowledge_base/utils.py index 896a8ad..df35234 100644 --- a/server/knowledge_base/utils.py +++ b/server/knowledge_base/utils.py @@ -1,6 +1,4 @@ import os -from typing import List -from server.utils import torch_gc from configs.model_config import KB_ROOT_PATH from langchain.vectorstores import FAISS from langchain.embeddings.huggingface import HuggingFaceEmbeddings