mirror of
https://github.com/RYDE-WORK/Langchain-Chatchat.git
synced 2026-01-31 03:03:22 +08:00
9ff7bef2c2
* 新功能： - 知识库管理界面支持查看、编辑、删除向量库文档。暂不支持增加（aggrid添加新行比较麻烦，需要另外实现） - 去除知识库管理界面中重建知识库和删除知识库按钮，建议用户到终端命令操作 修复： - 所有与知识库名称、文件名称有关的数据库操作函数都改成大小写不敏感，所有路径统一为 posix 风格，避免因路径文本不一致导致数据重复和操作失效 （close #2232） 开发者： - 添加 update_docs_by_id 函数与 API 接口。当前仅支持 FAISS，暂时未用到，未将来对知识库做更细致的修改做准备 - 统一 DocumentWithScore 与 DocumentWithVsId - FAISS 返回的 Document.metadata 中包含 ID， 方便后续查找比对 - /knowledge_base/search_docs 接口支持 file_name, metadata 参数，可以据此检索文档 * fix bug
114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
from typing import List, Dict, Optional
|
|
|
|
from langchain.schema import Document
|
|
from langchain.vectorstores.milvus import Milvus
|
|
|
|
from configs import 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)
|
|
|
|
# def save_vector_store(self):
|
|
# if self.milvus.col:
|
|
# self.milvus.col.flush()
|
|
|
|
def get_doc_by_ids(self, ids: List[str]) -> List[Document]:
|
|
result = []
|
|
if self.milvus.col:
|
|
data_list = self.milvus.col.query(expr=f'pk in {ids}', output_fields=["*"])
|
|
for data in data_list:
|
|
text = data.pop("text")
|
|
result.append(Document(page_content=text, metadata=data))
|
|
return result
|
|
|
|
def del_doc_by_ids(self, ids: List[str]) -> bool:
|
|
self.milvus.col.delete(expr=f'pk in {ids}')
|
|
|
|
@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):
|
|
self.milvus = Milvus(embedding_function=EmbeddingsFunAdapter(self.embed_model),
|
|
collection_name=self.kb_name,
|
|
connection_args=kbs_config.get("milvus"),
|
|
index_params=kbs_config.ge("milvus_kwargs")["index_params"],
|
|
search_params=kbs_config.get("milvus_kwargs")["search_params"]
|
|
)
|
|
|
|
def do_init(self):
|
|
self._load_milvus()
|
|
|
|
def do_drop_kb(self):
|
|
if self.milvus.col:
|
|
self.milvus.col.release()
|
|
self.milvus.col.drop()
|
|
|
|
def do_search(self, query: str, top_k: int, score_threshold: float):
|
|
self._load_milvus()
|
|
embed_func = EmbeddingsFunAdapter(self.embed_model)
|
|
embeddings = embed_func.embed_query(query)
|
|
docs = self.milvus.similarity_search_with_score_by_vector(embeddings, top_k)
|
|
return score_threshold_process(score_threshold, top_k, docs)
|
|
|
|
def do_add_doc(self, docs: List[Document], **kwargs) -> List[Dict]:
|
|
# TODO: workaround for bug #10492 in langchain
|
|
for doc in docs:
|
|
for k, v in doc.metadata.items():
|
|
doc.metadata[k] = str(v)
|
|
for field in self.milvus.fields:
|
|
doc.metadata.setdefault(field, "")
|
|
doc.metadata.pop(self.milvus._text_field, None)
|
|
doc.metadata.pop(self.milvus._vector_field, None)
|
|
|
|
ids = self.milvus.add_documents(docs)
|
|
doc_infos = [{"id": id, "metadata": doc.metadata} for id, doc in zip(ids, docs)]
|
|
return doc_infos
|
|
|
|
def do_delete_doc(self, kb_file: KnowledgeFile, **kwargs):
|
|
if self.milvus.col:
|
|
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.do_drop_kb()
|
|
self.do_init()
|
|
|
|
|
|
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"))
|
|
|
|
print(milvusService.get_doc_by_ids(["444022434274215486"]))
|
|
# milvusService.delete_doc(KnowledgeFile("README.md", "test"))
|
|
# milvusService.do_drop_kb()
|
|
# print(milvusService.search_docs("如何启动api服务"))
|