This commit is contained in:
zR 2024-01-22 11:48:16 +08:00
commit 8483c0bc96
3 changed files with 5 additions and 5 deletions

View File

@ -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:

View File

@ -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]]:
"""
搜索知识库子类实自己逻辑
"""

View File

@ -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):
@ -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: