mirror of
https://github.com/RYDE-WORK/Langchain-Chatchat.git
synced 2026-02-03 05:03:12 +08:00
3acbf4d5d1
* close #1172: 给webui_page/utils添加一些log信息，方便定位错误 * 修复：重建知识库时页面未实时显示进度 * skip model_worker running when using online model api such as chatgpt * 修改知识库管理相关内容: 1.KnowledgeFileModel增加3个字段：file_mtime(文件修改时间),file_size(文件大小),custom_docs(是否使用自定义docs)。为后面比对上传文件做准备。 2.给所有String字段加上长度，防止mysql建表错误(pr#1177) 3.统一[faiss/milvus/pgvector]_kb_service.add_doc接口，使其支持自定义docs 4.为faiss_kb_service增加一些方法，便于调用 5.为KnowledgeFile增加一些方法，便于获取文件信息，缓存file2text的结果。 * 修复/chat/fastchat无法流式输出的问题 * 新增功能： 1、KnowledgeFileModel增加"docs_count"字段，代表该文件加载到向量库中的Document数量，并在WEBUI中进行展示。 2、重建知识库`python init_database.py --recreate-vs`支持多线程。 其它： 统一代码中知识库相关函数用词：file代表一个文件名称或路径，doc代表langchain加载后的Document。部分与API接口有关或含义重叠的函数暂未修改。 --------- Co-authored-by: liunux4odoo <liunux@qq.com>, hongkong9771
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
from typing import List
|
|
|
|
import numpy as np
|
|
from faiss import normalize_L2
|
|
from langchain.embeddings.base import Embeddings
|
|
from langchain.schema import Document
|
|
from langchain.vectorstores import Milvus
|
|
from sklearn.preprocessing import normalize
|
|
|
|
from configs.model_config import SCORE_THRESHOLD, kbs_config
|
|
|
|
from server.knowledge_base.kb_service.base import KBService, SupportedVSType, EmbeddingsFunAdapter, \
|
|
score_threshold_process
|
|
from server.knowledge_base.utils import KnowledgeFile
|
|
|
|
|
|
class MilvusKBService(KBService):
|
|
milvus: Milvus
|
|
|
|
@staticmethod
|
|
def get_collection(milvus_name):
|
|
from pymilvus import Collection
|
|
return Collection(milvus_name)
|
|
|
|
@staticmethod
|
|
def search(milvus_name, content, limit=3):
|
|
search_params = {
|
|
"metric_type": "L2",
|
|
"params": {"nprobe": 10},
|
|
}
|
|
c = MilvusKBService.get_collection(milvus_name)
|
|
return c.search(content, "embeddings", search_params, limit=limit, output_fields=["content"])
|
|
|
|
def do_create_kb(self):
|
|
pass
|
|
|
|
def vs_type(self) -> str:
|
|
return SupportedVSType.MILVUS
|
|
|
|
def _load_milvus(self, embeddings: Embeddings = None):
|
|
if embeddings is None:
|
|
embeddings = self._load_embeddings()
|
|
self.milvus = Milvus(embedding_function=EmbeddingsFunAdapter(embeddings),
|
|
collection_name=self.kb_name, connection_args=kbs_config.get("milvus"))
|
|
|
|
def do_init(self):
|
|
self._load_milvus()
|
|
|
|
def do_drop_kb(self):
|
|
if self.milvus.col:
|
|
self.milvus.col.drop()
|
|
|
|
def do_search(self, query: str, top_k: int, score_threshold: float, embeddings: Embeddings):
|
|
self._load_milvus(embeddings=EmbeddingsFunAdapter(embeddings))
|
|
return score_threshold_process(score_threshold, top_k, self.milvus.similarity_search_with_score(query, top_k))
|
|
|
|
def do_add_doc(self, docs: List[Document], embeddings: Embeddings, **kwargs):
|
|
self.milvus.add_documents(docs)
|
|
|
|
def do_delete_doc(self, kb_file: KnowledgeFile, **kwargs):
|
|
filepath = kb_file.filepath.replace('\\', '\\\\')
|
|
delete_list = [item.get("pk") for item in
|
|
self.milvus.col.query(expr=f'source == "{filepath}"', output_fields=["pk"])]
|
|
self.milvus.col.delete(expr=f'pk in {delete_list}')
|
|
|
|
def do_clear_vs(self):
|
|
if self.milvus.col:
|
|
self.milvus.col.drop()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# 测试建表使用
|
|
from server.db.base import Base, engine
|
|
|
|
Base.metadata.create_all(bind=engine)
|
|
milvusService = MilvusKBService("test")
|
|
milvusService.add_doc(KnowledgeFile("README.md", "test"))
|
|
milvusService.delete_doc(KnowledgeFile("README.md", "test"))
|
|
milvusService.do_drop_kb()
|
|
print(milvusService.search_docs("如何启动api服务"))
|