15308555518/ MaxKB
1
0
Fork 0
mirror of https://github.com/1Panel-dev/MaxKB.git synced 2025-12-26 10:12:51 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

feat: Add document_id_list parameter to query methods in vector classes

This commit is contained in:
CaptainB 2025-10-21 21:18:24 +08:00
parent 09c5c9c042
commit 219fe97c87
4 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
apps/application/chat_pipeline/step/search_dataset_step/impl/base_search_dataset_step.py
View File

@ -67,7 +67,7 @@ class BaseSearchDatasetStep(ISearchDatasetStep):
embedding_model = ModelManage.get_model(model_id, lambda _id: get_model(model))
embedding_value = embedding_model.embed_query(exec_problem_text)
vector = VectorStore.get_embedding_vector()
embedding_list = vector.query(exec_problem_text, embedding_value, knowledge_id_list, exclude_document_id_list,
embedding_list = vector.query(exec_problem_text, embedding_value, knowledge_id_list, None, exclude_document_id_list,
exclude_paragraph_id_list, True, top_n, similarity, SearchMode(search_mode))
if embedding_list is None:
return []

6
apps/application/flow/step_node/search_knowledge_node/impl/base_search_knowledge_node.py
View File

@ -79,12 +79,14 @@ class BaseSearchKnowledgeNode(ISearchKnowledgeStepNode):
self.context['question'] = question
self.context['show_knowledge'] = show_knowledge
document_id_list = None
if search_scope_type == 'referencing': # 引用上一步知识库/文档
if search_scope_source == 'knowledge': # 知识库
knowledge_id_list = self.get_reference_content(search_scope_reference)
else: # 文档
document_id_list = self.get_reference_content(search_scope_reference)
knowledge_id_list = QuerySet(Document).filter(
id__in=self.get_reference_content(search_scope_reference)
id__in=document_id_list
).values_list(
'knowledge_id', flat=True
).distinct()
@ -105,7 +107,7 @@ class BaseSearchKnowledgeNode(ISearchKnowledgeStepNode):
QuerySet(Document).filter(
knowledge_id__in=knowledge_id_list,
is_active=False)]
embedding_list = vector.query(question, embedding_value, knowledge_id_list, exclude_document_id_list,
embedding_list = vector.query(question, embedding_value, knowledge_id_list, document_id_list, exclude_document_id_list,
exclude_paragraph_id_list, True, knowledge_setting.get('top_n'),
knowledge_setting.get('similarity'),
SearchMode(knowledge_setting.get('search_mode')))

1
apps/knowledge/vector/base_vector.py
View File

@ -126,6 +126,7 @@ class BaseVectorStore(ABC):
@abstractmethod
def query(self, query_text: str, query_embedding: List[float], knowledge_id_list: list[str],
document_id_list: list[str] | None,
exclude_document_id_list: list[str],
exclude_paragraph_list: list[str], is_active: bool, top_n: int, similarity: float,
search_mode: SearchMode):

3
apps/knowledge/vector/pg_vector.py
View File

@ -97,6 +97,7 @@ class PGVector(BaseVectorStore):
return search_handle.handle(query_set, query_text, embedding_query, top_number, similarity, search_mode)
def query(self, query_text: str, query_embedding: List[float], knowledge_id_list: list[str],
document_id_list: list[str],
exclude_document_id_list: list[str],
exclude_paragraph_list: list[str], is_active: bool, top_n: int, similarity: float,
search_mode: SearchMode):
@ -104,6 +105,8 @@ class PGVector(BaseVectorStore):
if knowledge_id_list is None or len(knowledge_id_list) == 0:
return []
query_set = QuerySet(Embedding).filter(knowledge_id__in=knowledge_id_list, is_active=is_active)
if document_id_list is not None and len(document_id_list) > 0:
query_set = query_set.filter(document_id__in=document_id_list)
if exclude_document_id_list is not None and len(exclude_document_id_list) > 0:
query_set = query_set.exclude(document_id__in=exclude_document_id_list)
if exclude_paragraph_list is not None and len(exclude_paragraph_list) > 0: