From 3057c29e471c0bf23f07a2cff7022d81aa40c0aa Mon Sep 17 00:00:00 2001 From: showmecodett Date: Sun, 21 Jan 2024 20:56:26 +0800 Subject: [PATCH 1/2] perf typing --- server/db/repository/knowledge_file_repository.py | 2 +- server/knowledge_base/kb_service/base.py | 4 ++-- server/knowledge_base/kb_service/faiss_kb_service.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/server/db/repository/knowledge_file_repository.py b/server/db/repository/knowledge_file_repository.py index 0a6b782d..4388e7ad 100644 --- a/server/db/repository/knowledge_file_repository.py +++ b/server/db/repository/knowledge_file_repository.py @@ -83,7 +83,7 @@ def add_file_to_db(session, kb_file: KnowledgeFile, docs_count: int = 0, custom_docs: bool = False, - doc_infos: List[str] = [], # 形式：[{"id": str, "metadata": dict}, ...] + doc_infos: List[Dict] = [], # 形式：[{"id": str, "metadata": dict}, ...] ): kb = session.query(KnowledgeBaseModel).filter_by(kb_name=kb_file.kb_name).first() if kb: diff --git a/server/knowledge_base/kb_service/base.py b/server/knowledge_base/kb_service/base.py index 86b99054..44c0d64e 100644 --- a/server/knowledge_base/kb_service/base.py +++ b/server/knowledge_base/kb_service/base.py @@ -24,7 +24,7 @@ from server.knowledge_base.utils import ( list_kbs_from_folder, list_files_from_folder, ) -from typing import List, Union, Dict, Optional +from typing import List, Union, Dict, Optional, Tuple from server.embeddings_api import embed_texts, aembed_texts, embed_documents from server.knowledge_base.model.kb_document_model import DocumentWithVSId @@ -261,7 +261,7 @@ class KBService(ABC): query: str, top_k: int, score_threshold: float, - ) -> List[Document]: + ) -> List[Tuple[Document, float]]: """ 搜索知识库子类实自己逻辑 """ diff --git a/server/knowledge_base/kb_service/faiss_kb_service.py b/server/knowledge_base/kb_service/faiss_kb_service.py index f073b4ed..501b2e1f 100644 --- a/server/knowledge_base/kb_service/faiss_kb_service.py +++ b/server/knowledge_base/kb_service/faiss_kb_service.py @@ -61,7 +61,7 @@ class FaissKBService(KBService): query: str, top_k: int, score_threshold: float = SCORE_THRESHOLD, - ) -> List[Document]: + ) -> List[Tuple[Document, float]]: embed_func = EmbeddingsFunAdapter(self.embed_model) embeddings = embed_func.embed_query(query) with self.load_vector_store().acquire() as vs: From 71deed93dfeab8bd67c20e229dcec29cdec8ae87 Mon Sep 17 00:00:00 2001 From: showmecodett Date: Sun, 21 Jan 2024 21:00:12 +0800 Subject: [PATCH 2/2] perf typing --- server/knowledge_base/kb_service/faiss_kb_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/knowledge_base/kb_service/faiss_kb_service.py b/server/knowledge_base/kb_service/faiss_kb_service.py index 501b2e1f..c05f6136 100644 --- a/server/knowledge_base/kb_service/faiss_kb_service.py +++ b/server/knowledge_base/kb_service/faiss_kb_service.py @@ -7,7 +7,7 @@ from server.knowledge_base.kb_cache.faiss_cache import kb_faiss_pool, ThreadSafe from server.knowledge_base.utils import KnowledgeFile, get_kb_path, get_vs_path from server.utils import torch_gc from langchain.docstore.document import Document -from typing import List, Dict, Optional +from typing import List, Dict, Optional, Tuple class FaissKBService(KBService):