youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-04 01:18:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

r2

Browse Source
This commit is contained in:
jyong
2025-05-29 09:53:42 +08:00
parent 7f59ffe7af
commit 797d044714
5 changed files with 50 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
api/services/dataset_service.py
View File

@ -53,6 +53,7 @@ from services.entities.knowledge_entities.knowledge_entities import (
)
from services.entities.knowledge_entities.rag_pipeline_entities import (
KnowledgeBaseUpdateConfiguration,
KnowledgeConfiguration,
RagPipelineDatasetCreateEntity,
)
from services.errors.account import InvalidActionError, NoPermissionError
@ -495,11 +496,11 @@ class DatasetService:
@staticmethod
def update_rag_pipeline_dataset_settings(session: Session,
dataset: Dataset,
knowledge_base_setting: KnowledgeBaseUpdateConfiguration,
knowledge_configuration: KnowledgeConfiguration,
has_published: bool = False):
if not has_published:
dataset.chunk_structure = knowledge_base_setting.chunk_structure
index_method = knowledge_base_setting.index_method
dataset.chunk_structure = knowledge_configuration.chunk_structure
index_method = knowledge_configuration.index_method
dataset.indexing_technique = index_method.indexing_technique
if index_method == "high_quality":
model_manager = ModelManager()
@ -519,26 +520,26 @@ class DatasetService:
dataset.keyword_number = index_method.economy_setting.keyword_number
else:
raise ValueError("Invalid index method")
dataset.retrieval_model = knowledge_base_setting.retrieval_setting.model_dump()
dataset.retrieval_model = knowledge_configuration.retrieval_setting.model_dump()
session.add(dataset)
else:
if dataset.chunk_structure and dataset.chunk_structure != knowledge_base_setting.chunk_structure:
if dataset.chunk_structure and dataset.chunk_structure != knowledge_configuration.chunk_structure:
raise ValueError("Chunk structure is not allowed to be updated.")
action = None
if dataset.indexing_technique != knowledge_base_setting.index_method.indexing_technique:
if dataset.indexing_technique != knowledge_configuration.index_method.indexing_technique:
# if update indexing_technique
if knowledge_base_setting.index_method.indexing_technique == "economy":
if knowledge_configuration.index_method.indexing_technique == "economy":
raise ValueError("Knowledge base indexing technique is not allowed to be updated to economy.")
elif knowledge_base_setting.index_method.indexing_technique == "high_quality":
elif knowledge_configuration.index_method.indexing_technique == "high_quality":
action = "add"
# get embedding model setting
try:
model_manager = ModelManager()
embedding_model = model_manager.get_model_instance(
tenant_id=current_user.current_tenant_id,
provider=knowledge_base_setting.index_method.embedding_setting.embedding_provider_name,
provider=knowledge_configuration.index_method.embedding_setting.embedding_provider_name,
model_type=ModelType.TEXT_EMBEDDING,
model=knowledge_base_setting.index_method.embedding_setting.embedding_model_name,
model=knowledge_configuration.index_method.embedding_setting.embedding_model_name,
)
dataset.embedding_model = embedding_model.model
dataset.embedding_model_provider = embedding_model.provider
@ -607,9 +608,9 @@ class DatasetService:
except ProviderTokenNotInitError as ex:
raise ValueError(ex.description)
elif dataset.indexing_technique == "economy":
if dataset.keyword_number != knowledge_base_setting.index_method.economy_setting.keyword_number:
dataset.keyword_number = knowledge_base_setting.index_method.economy_setting.keyword_number
dataset.retrieval_model = knowledge_base_setting.retrieval_setting.model_dump()
if dataset.keyword_number != knowledge_configuration.index_method.economy_setting.keyword_number:
dataset.keyword_number = knowledge_configuration.index_method.economy_setting.keyword_number
dataset.retrieval_model = knowledge_configuration.retrieval_setting.model_dump()
session.add(dataset)
session.commit()
if action:

30
api/services/rag_pipeline/rag_pipeline.py
View File

@ -47,7 +47,7 @@ from models.workflow import (
WorkflowType,
)
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeBaseUpdateConfiguration, PipelineTemplateInfoEntity
from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeBaseUpdateConfiguration, KnowledgeConfiguration, PipelineTemplateInfoEntity
from services.errors.app import WorkflowHashNotEqualError
from services.rag_pipeline.pipeline_template.pipeline_template_factory import PipelineTemplateRetrievalFactory
@ -262,7 +262,6 @@ class RagPipelineService:
session: Session,
pipeline: Pipeline,
account: Account,
knowledge_base_setting: KnowledgeBaseUpdateConfiguration,
) -> Workflow:
draft_workflow_stmt = select(Workflow).where(
Workflow.tenant_id == pipeline.tenant_id,
@ -291,16 +290,23 @@ class RagPipelineService:
# commit db session changes
session.add(workflow)
# update dataset
dataset = pipeline.dataset
if not dataset:
raise ValueError("Dataset not found")
DatasetService.update_rag_pipeline_dataset_settings(
session=session,
dataset=dataset,
knowledge_base_setting=knowledge_base_setting,
has_published=pipeline.is_published
)
graph = workflow.graph_dict
nodes = graph.get("nodes", [])
for node in nodes:
if node.get("data", {}).get("type") == "knowledge_index":
knowledge_configuration = node.get("data", {}).get("knowledge_configuration", {})
knowledge_configuration = KnowledgeConfiguration(**knowledge_configuration)
# update dataset
dataset = pipeline.dataset
if not dataset:
raise ValueError("Dataset not found")
DatasetService.update_rag_pipeline_dataset_settings(
session=session,
dataset=dataset,
knowledge_configuration=knowledge_configuration,
has_published=pipeline.is_published
)
# return new workflow
return workflow