mtfyun/dify
1
0 Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-11-13 14:51:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): establish dataset application boundaries (#41729)

Browse Source 
Co-authored-by: Jyong <jyong@dify.ai>
This commit is contained in:
Byron.wangandJyong authored and GitHub committed 2026-09-29 08:42:31 +00:00
1 parent c005932db9
commit a67cd87eb0
417 files changed
+30522 -23751
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.github/CODEOWNERS
+1 -4
View File
@@ -57,13 +57,10 @@
# Backend - RAG (Retrieval Augmented Generation)
/api/core/rag/ @JohnJyong
/api/services/knowledge/ @JohnJyong
/api/services/rag_pipeline/ @JohnJyong
/api/services/dataset_service.py @JohnJyong
/api/services/knowledge_service.py @JohnJyong
/api/services/external_knowledge_service.py @JohnJyong
/api/services/hit_testing_service.py @JohnJyong
/api/services/metadata_service.py @JohnJyong
/api/services/vector_service.py @JohnJyong
/api/services/entities/knowledge_entities/ @JohnJyong
/api/services/entities/external_knowledge_entities/ @JohnJyong
/api/controllers/console/datasets/ @JohnJyong
api/.importlinter
+193 -17
View File
@@ -1,12 +1,14 @@
[importlinter]
root_packages =
core
core.rag.extractor.entity
constants
context
configs
controllers
extensions
factories
fields
libs
machinery
models
@@ -16,7 +18,7 @@ root_packages =
include_external_packages = True
[importlinter:contract:backend-layers]
name = Backend modules follow the transport, application, domain, and shared-library layers
name = Backend modules follow the transport, application, execution-kernel, and shared-library layers
type = layers
layers =
controllers
@@ -27,7 +29,7 @@ layers =
# error makes this list shrink whenever a legacy dependency is removed.
unmatched_ignore_imports_alerting = error
ignore_imports =
core.app.app_config.easy_ui_based_app.dataset.manager -> services.dataset_service
core.app.app_config.easy_ui_based_app.dataset.manager -> services.knowledge.dataset_service
core.app.apps.advanced_chat.app_generator -> controllers.console.app.workflow
core.app.apps.advanced_chat.app_generator -> services.conversation_service
core.app.apps.advanced_chat.app_generator -> services.errors.conversation
@@ -48,45 +50,44 @@ ignore_imports =
core.app.apps.message_based_app_generator -> services.errors.conversation
core.app.apps.message_based_app_generator -> services.errors.message
core.app.apps.pipeline.pipeline_generator -> services
core.app.apps.pipeline.pipeline_generator -> services.dataset_service
core.app.apps.pipeline.pipeline_generator -> services.datasource_provider_service
core.app.apps.pipeline.pipeline_generator -> services.knowledge.dataset_service
core.app.apps.pipeline.pipeline_generator -> services.data_source.provider_service
core.app.apps.pipeline.pipeline_generator -> services.feature_service
core.app.apps.pipeline.pipeline_generator -> services.workflow_draft_variable_service
core.app.apps.pipeline.pipeline_runner -> services.dataset_ref_service
core.app.apps.workflow.app_generator -> controllers.console.app.workflow
core.app.apps.workflow.app_generator -> services.snippet_generate_service
core.app.apps.workflow.app_generator -> services.workflow_draft_variable_service
core.app.apps.workflow_app_runner -> tasks.mail_human_input_delivery_task
core.app.features.annotation_reply.annotation_reply -> services.annotation_service
core.app.features.annotation_reply.annotation_reply -> services.dataset_service
core.app.features.annotation_reply.annotation_reply -> services.knowledge.dataset_service
core.app.llm.quota -> services.credit_pool_service
core.app.task_pipeline.message_cycle_manager -> services.account_service
core.app.task_pipeline.message_cycle_manager -> services.annotation_service
core.app.workflow.layers.persistence -> services.workflow.inspector_events
core.datasource.datasource_manager -> services.datasource_provider_service
core.helper.credential_utils -> services.enterprise.plugin_manager_service
core.helper.credential_utils -> services.system_feature_service
core.indexing_runner -> services.vector_space_admission_service
core.mcp.auth_client -> services
core.provider_manager -> services.credential_permission_service
core.provider_manager -> services.credit_pool_service
core.provider_manager -> services.feature_service
core.rag.datasource.retrieval_service -> services.external_knowledge_service
core.rag.datasource.retrieval_service -> services.knowledge.external.service
core.rag.index_processor.index_processor -> services.vector_space_admission_service
core.rag.index_processor.index_processor_base -> services.file_service
core.rag.index_processor.processor.paragraph_index_processor -> services.account_service
core.rag.index_processor.processor.paragraph_index_processor -> services.summary_index_service
core.rag.index_processor.processor.paragraph_index_processor -> services.knowledge.summaries.adapters
core.rag.index_processor.processor.parent_child_index_processor -> services.account_service
core.rag.index_processor.processor.parent_child_index_processor -> services.summary_index_service
core.rag.index_processor.processor.qa_index_processor -> services.summary_index_service
core.rag.retrieval.dataset_retrieval -> services.external_knowledge_service
core.rag.index_processor.processor.parent_child_index_processor -> services.knowledge.summaries.adapters
core.rag.index_processor.processor.qa_index_processor -> services.knowledge.summaries.adapters
core.rag.retrieval.dataset_retrieval -> services.knowledge.external.service
core.rag.retrieval.dataset_retrieval -> services.knowledge.retrieval.attachments
core.rag.retrieval.dataset_retrieval -> services.feature_service
core.rag.summary_index.summary_index -> services.summary_index_service
core.rag.summary_index.summary_index -> services.knowledge.summaries.adapters
core.rag.summary_index.summary_index -> tasks.generate_summary_index_task
core.repositories.sqlalchemy_workflow_node_execution_repository -> services.file_service
core.repositories.sqlalchemy_workflow_node_execution_repository -> services.variable_truncator
core.tools.tool_manager -> services
core.workflow.node_runtime -> services
# Node construction delegates datasource dependency assembly to the composition root.
core.workflow.node_factory -> extensions.application_services.data_sources
core.workflow.nodes.agent.message_transformer -> services
core.workflow.nodes.agent_v2.agent_node -> services.agent.prompt_mentions
core.workflow.nodes.agent_v2.agent_node -> services.agent.workspace_service
@@ -99,7 +100,7 @@ ignore_imports =
core.workflow.nodes.agent_v2.workspace_retirement_layer -> tasks.collect_agent_resources_task
libs.email_i18n -> services.entities.feature_entities
libs.email_i18n -> services.system_feature_service
libs.external_api -> core
libs.external_api -> core.plugin
libs.external_api -> core.errors.error
libs.external_api -> extensions.ext_logging
libs.flask_utils -> models
@@ -168,6 +169,141 @@ forbidden_modules =
sqlalchemy
werkzeug
[importlinter:contract:dataset-application-boundary]
name = Dataset application services and contracts are framework and persistence neutral
type = forbidden
source_modules =
core.rag.extractor.entity.extract_setting
services.knowledge.entities.segments
services.knowledge.entities.datasets
services.knowledge.entities.document_creation
services.data_source.binding_application_service
services.data_source.credential_gateway
services.data_source.notion_import_application_service
services.data_source.entities.credential
services.data_source.entities.notion_import
services.knowledge.entities.indexing_estimate
services.knowledge.dataset_access
services.knowledge.datasets.application
services.knowledge.datasets.retrieval
services.knowledge.document_sync
services.knowledge.documents.application
services.knowledge.external.application
services.knowledge.indexing.errors
services.knowledge.indexing.estimate
services.knowledge.indexing.execution
services.knowledge.resource_scope
services.knowledge.segments.application
services.knowledge.metadata.application
services.knowledge.summaries.application
forbidden_modules =
fields
configs
controllers
extensions
flask
models
repositories
sqlalchemy
tasks
werkzeug
services.knowledge.dataset_service
services.data_source.provider_service
services.knowledge.indexing.adapters
services.knowledge.segments.adapters
services.knowledge.segments.indexing
core.datasource.datasource_manager
core.rag.extractor.notion_extractor
[importlinter:contract:knowledge-repository-boundary]
name = Knowledge persistence adapters only depend on their consumer-owned application contracts
type = forbidden
source_modules =
repositories.knowledge
forbidden_modules =
controllers
extensions
flask
services
tasks
werkzeug
allow_indirect_imports = True
ignore_imports =
repositories.knowledge.dataset_api_key_repository -> services.auth.api_key_contracts
repositories.knowledge.dataset_api_key_repository -> services.knowledge.api_key_service
repositories.knowledge.dataset_repository -> services.knowledge.segments.application
repositories.knowledge.dataset_repository -> services.knowledge.dataset_access
repositories.knowledge.dataset_repository -> services.knowledge.indexing.estimate
repositories.knowledge.dataset_repository -> services.knowledge.resource_scope
repositories.knowledge.document_repository -> services.knowledge.document_sync
repositories.knowledge.document_repository -> services.knowledge.indexing.estimate
repositories.knowledge.document_repository -> services.knowledge.indexing.errors
repositories.knowledge.document_repository -> services.knowledge.indexing.execution
repositories.knowledge.document_repository -> services.knowledge.resource_scope
repositories.knowledge.segment_repository -> services.knowledge.resource_scope
repositories.knowledge.segment_repository -> services.knowledge.indexing.execution
repositories.knowledge.segment_repository -> services.knowledge.segments.application
repositories.knowledge.segment_repository -> services.knowledge.entities.segments
repositories.knowledge.metadata_repository -> services.errors.metadata
repositories.knowledge.metadata_repository -> services.knowledge.metadata.application
repositories.knowledge.metadata_repository -> services.knowledge.resource_scope
repositories.knowledge.summary_repository -> services.knowledge.summaries.application
[importlinter:contract:datasource-credential-repository-boundary]
name = Data source credential persistence does not depend on SQLAlchemy service helpers
type = forbidden
source_modules =
repositories.credentials.query_repository
repositories.data_source.credential_repository
forbidden_modules =
services.data_source.provider_service
allow_indirect_imports = True
[importlinter:contract:datasource-provider-persistence-boundary]
name = Datasource provider service consumes injected persistence contracts
type = forbidden
source_modules =
services.data_source.provider_service
forbidden_modules =
repositories
extensions.application_services
extensions.ext_application_services
allow_indirect_imports = True
[importlinter:contract:credential-permission-query-policy-boundary]
name = Credential query contracts do not expose persistence implementations
type = forbidden
source_modules =
services.credentials.query
forbidden_modules =
repositories
sqlalchemy
models
extensions
flask
allow_indirect_imports = True
[importlinter:contract:credential-query-implementation-boundary]
name = Credential query implementations are wired only by composition or other repositories
type = forbidden
source_modules =
controllers
services
core
forbidden_modules =
repositories.credentials.query_repository
allow_indirect_imports = True
[importlinter:contract:pipeline-document-lookup-boundary]
name = Pipeline core depends on its document store port
type = forbidden
source_modules =
core.app.apps.pipeline.pipeline_generator
core.app.apps.pipeline.pipeline_runner
forbidden_modules =
repositories
allow_indirect_imports = True
[importlinter:contract:init-validation-service-boundary]
name = Initialization validation application service is framework and persistence neutral
type = forbidden
@@ -479,7 +615,7 @@ forbidden_modules =
name = Data source OAuth application service is framework and persistence neutral
type = forbidden
source_modules =
services.data_source_oauth_service
services.data_source.oauth_service
forbidden_modules =
configs
controllers
@@ -800,6 +936,46 @@ forbidden_modules =
tasks
werkzeug
[importlinter:contract:knowledge-segment-presentation-boundary]
name = Knowledge operations do not serialize HTTP responses
type = forbidden
allow_indirect_imports = True
source_modules =
services.knowledge.segments
services.knowledge.documents
services.knowledge.datasets
services.knowledge.external
forbidden_modules =
controllers
fields
[importlinter:contract:knowledge-response-persistence-boundary]
name = Knowledge response schemas do not load database state
type = forbidden
allow_indirect_imports = True
source_modules =
fields.dataset_fields
fields.document_fields
fields.segment_fields
forbidden_modules =
sqlalchemy
models
repositories
services.knowledge
[importlinter:contract:dataset-model-infrastructure-boundary]
name = Dataset models do not execute infrastructure or application operations
type = forbidden
allow_indirect_imports = True
source_modules =
models.dataset
forbidden_modules =
repositories
services
extensions
core.tools
models.engine
[importlinter:contract:app-mcp-server-service-boundary]
name = App MCP server application service is framework and persistence neutral
type = forbidden
api/AGENTS.md
+2 -1
View File
@@ -15,7 +15,8 @@ Run direct Python commands through `uv run --project api`. Docker-backed integra
## Architecture And Boundaries
- Keep transport parsing and serialization in controllers, orchestration in services, and domain policy in `core/` or its domain owner. Keep `libs/` business-agnostic and reuse existing owners before adding abstractions.
- Treat `api/core/` as a migration-only area: do not add files or move implementations into it. Move responsibilities out; limit changes to existing callers to the adaptations needed for that migration.
- Keep transport parsing and serialization in controllers, orchestration in services, and domain policy in its domain owner. New or extracted implementations must stay outside `core/`. Keep `libs/` business-agnostic and reuse existing owners before adding abstractions.
- Before changing controller schemas, generated API contracts, or `SystemFeatureModel`, read `controllers/API_SCHEMA_GUIDE.md`. Treat `/system-features` as a minimal unauthenticated bootstrap allowlist, not a general configuration registry.
- Scope tenant-owned reads and writes by the complete owner chain, and propagate `tenant_id` across every affected layer. Reconstruct trusted internal references from validated database state after payload or async boundaries.
- Keep write transactions explicit and bounded. Do not perform external I/O inside an open transaction unless a documented consistency contract requires it.
api/commands/vector.py
+2 -1
View File
@@ -19,6 +19,7 @@ from models.dataset import Dataset, DatasetCollectionBinding, DatasetMetadata, D
from models.dataset import Document as DatasetDocument
from models.enums import DatasetMetadataType, IndexingStatus, SegmentStatus
from models.model import App, AppAnnotationSetting, MessageAnnotation
from repositories.knowledge.dataset_read_repository import get_segment_child_chunks
@click.command("vdb-migrate", help="Migrate vector db.")
@@ -278,7 +279,7 @@ def migrate_knowledge_vector_database():
},
)
if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
child_chunks = segment.get_child_chunks(session=db_session)
child_chunks = get_segment_child_chunks(segment, session=db_session)
if child_chunks:
child_documents = []
for child_chunk in child_chunks:
api/controllers/common/rbac/__init__.py
+2
View File
@@ -4,6 +4,7 @@ from .checks import RBAC_CHECKS_ATTR, RBACCheck, enforce_rbac_checks
from .locators import (
AgentBehindApp,
AgentId,
DatasetByDocument,
DatasetByPipeline,
DatasetId,
PlainApp,
@@ -16,6 +17,7 @@ __all__ = [
"RBAC_CHECKS_ATTR",
"AgentBehindApp",
"AgentId",
"DatasetByDocument",
"DatasetByPipeline",
"DatasetId",
"PlainApp",
api/controllers/common/rbac/locators.py
+13
View File
@@ -15,6 +15,7 @@ if TYPE_CHECKING:
__all__ = [
"AgentBehindApp",
"AgentId",
"DatasetByDocument",
"DatasetByPipeline",
"DatasetId",
"PlainApp",
@@ -133,6 +134,18 @@ class DatasetId(_ParamLocator):
return RBACResourceService.get_dataset_maintainer(db.session, tenant_id, identity.id)
class DatasetByDocument(DatasetId):
default_param = "document_id"
@override
def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
document_id = _required(path_args, self.param)
dataset_id = RBACResourceService.get_dataset_id_by_document(db.session, tenant_id, document_id)
if dataset_id is None:
raise NotFound("Document not found")
return ResourceIdentity(self.scope, dataset_id)
class DatasetByPipeline(DatasetId):
default_param = "pipeline_id"
api/controllers/console/app/app.py
+2 -2
View File
@@ -79,7 +79,8 @@ from services.entities.app_entities import (
UpdateAppParams,
)
from services.entities.dsl_entities import ImportStatus
from services.entities.knowledge_entities.knowledge_entities import (
from services.errors.account import NoPermissionError
from services.knowledge.entities.knowledge_entities import (
DataSource,
InfoList,
NotionIcon,
@@ -91,7 +92,6 @@ from services.entities.knowledge_entities.knowledge_entities import (
WeightModel,
WeightVectorSetting,
)
from services.errors.account import NoPermissionError
register_enum_models(console_ns, IconType)
api/controllers/console/auth/data_source_bearer_auth.py
+4 -4
View File
@@ -27,7 +27,7 @@ from services.auth.errors import (
InvalidDataSourceApiKeyAuthCredentialsError,
UnsupportedDataSourceApiKeyAuthProviderError,
)
from services.entities.data_source_api_key_auth_entities import (
from services.data_source.entities.api_key_auth import (
DataSourceApiKeyAuthBindingCreate,
DataSourceApiKeyAuthCredentials,
)
@@ -80,7 +80,7 @@ class ApiKeyAuthDataSource(Resource):
@console_ns.response(HTTPStatus.OK, "Success", console_ns.models[ApiKeyAuthDataSourceListResponse.__name__])
@console_account_admission()
def get(self, request_context: RequestContext):
bindings = application_services().data_source_api_key_auth.list_bindings(request_context)
bindings = application_services().data_sources.api_key_auth.list_bindings(request_context)
return dump_response(
ApiKeyAuthDataSourceListResponse,
{
@@ -120,7 +120,7 @@ class ApiKeyAuthDataSourceBinding(Resource):
),
)
try:
application_services().data_source_api_key_auth.create_binding(request_context, command)
application_services().data_sources.api_key_auth.create_binding(request_context, command)
except UnsupportedDataSourceApiKeyAuthProviderError as exc:
raise DataSourceApiKeyAuthProviderNotSupportedError() from exc
except InvalidDataSourceApiKeyAuthCredentialsError as exc:
@@ -140,5 +140,5 @@ class ApiKeyAuthDataSourceBindingDelete(Resource):
rbac_checks=[RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace())],
)
def delete(self, request_context: RequestContext, binding_id: UUID):
application_services().data_source_api_key_auth.delete_binding(request_context, str(binding_id))
application_services().data_sources.api_key_auth.delete_binding(request_context, str(binding_id))
return "", HTTPStatus.NO_CONTENT
api/controllers/console/auth/data_source_oauth.py
+6 -6
View File
@@ -24,13 +24,13 @@ from fields.base import ResponseModel
from libs.helper import dump_response
from machinery.context import RequestContext
from models.account import TenantAccountRole
from services.data_source_oauth_service import (
from services.data_source.entities.oauth import DataSourceOAuthCallback
from services.data_source.oauth_service import (
DataSourceOAuthConfigurationError,
DataSourceOAuthError,
InvalidDataSourceOAuthCodeError,
InvalidDataSourceOAuthProviderError,
)
from services.entities.data_source_oauth_entities import DataSourceOAuthCallback
from .. import console_ns
@@ -116,7 +116,7 @@ class OAuthDataSource(Resource):
)
def get(self, request_context: RequestContext, provider: str):
try:
service = application_services().resolve_data_source_oauth(provider)
service = application_services().data_sources.resolve_oauth(provider)
authorization = service.start_authorization(request_context)
except InvalidDataSourceOAuthProviderError:
return _invalid_provider_response()
@@ -141,7 +141,7 @@ class OAuthDataSourceCallback(Resource):
def get(self, provider: str):
query = query_params_from_request(OAuthDataSourceCallbackQuery)
try:
service = application_services().resolve_data_source_oauth(provider)
service = application_services().data_sources.resolve_oauth(provider)
callback = service.complete_callback(
code=query.code,
error=query.error,
@@ -172,7 +172,7 @@ class OAuthDataSourceBinding(Resource):
if not query.code:
return {"error": "Invalid code"}, HTTPStatus.BAD_REQUEST
try:
service = application_services().resolve_data_source_oauth(provider)
service = application_services().data_sources.resolve_oauth(provider)
service.bind(request_context, code=query.code)
except InvalidDataSourceOAuthProviderError:
return _invalid_provider_response()
@@ -198,7 +198,7 @@ class OAuthDataSourceSync(Resource):
@console_account_admission()
def get(self, request_context: RequestContext, provider: str, binding_id: UUID):
try:
service = application_services().resolve_data_source_oauth(provider)
service = application_services().data_sources.resolve_oauth(provider)
service.sync(
request_context,
binding_id=str(binding_id),
api/controllers/console/datasets/data_source.py
+190 -305
View File
@@ -1,53 +1,59 @@
import json
from collections.abc import Generator
from datetime import datetime
from typing import Any, Literal, cast
from typing import Never
from uuid import UUID
from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_serializer
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from configs import dify_config
from controllers.common.errors import InvalidArgumentError, NotFoundError
from controllers.common.fields import SimpleResultResponse, TextContentResponse
from controllers.common.rbac import DatasetId, RBACCheck, Workspace
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from core.datasource.entities.datasource_entities import DatasourceProviderType, OnlineDocumentPagesMessage
from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin
from core.entities.knowledge_entities import IndexingEstimate
from core.indexing_runner import IndexingRunner
from core.rag.extractor.entity.datasource_type import DatasourceType
from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo
from core.rag.extractor.notion_extractor import NotionExtractor
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.datetime_utils import naive_utc_now
from libs.helper import dump_response, to_timestamp
from libs.login import login_required
from models import Account, DataSourceOauthBinding, Document
from services.dataset_service import DatasetService, DocumentService
from services.datasource_provider_service import DatasourceProviderService
from tasks.document_indexing_sync_task import document_indexing_sync_task
from .. import console_ns
from ..wraps import (
RBACPermission,
account_initialization_required,
is_admin_or_owner_required,
model_validate,
rbac_permission_required,
setup_required,
with_current_tenant_id,
with_current_user,
from controllers.console import console_ns
from controllers.console.app.error import ProviderNotInitializeError
from controllers.console.datasets.error import (
DatasetAccessDeniedRequestError,
IndexingEstimateError,
InvalidActionError,
)
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import model_validate
from core.entities.knowledge_entities import IndexingEstimate
from core.rbac import RBACPermission
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from fields.dataset_fields import NotionEstimateWorkspacePayload
from libs.helper import dump_response, to_timestamp
from machinery.context import RequestContext
from models.account import TenantAccountRole
from services.data_source.binding_application_service import (
BindingAction,
DataSourceBindingNotFoundError,
DataSourceBindingStateError,
)
from services.data_source.entities.notion_import import NotionPageType
from services.data_source.notion_import_application_service import (
DatasetIsNotNotionSourceError,
NotionImportCredentialUnavailableError,
)
from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError
from services.knowledge.document_sync import SyncDocumentNotFoundError, SyncDocumentSourceError
from services.knowledge.entities.indexing_estimate import NewSourcesEstimateCommand, NotionEstimateSource
from services.knowledge.indexing.estimate import (
EstimateSourceNotFoundError,
IndexingEstimateCredentialUnavailableError,
IndexingEstimateExecutionError,
IndexingEstimateProviderUnavailableError,
)
_ADMIN_OR_OWNER_ROLES = frozenset({TenantAccountRole.ADMIN, TenantAccountRole.OWNER})
_DATA_SOURCE_OAUTH_BASE_PATH = "/console/api/oauth/data-source"
class NotionEstimatePayload(BaseModel):
notion_info_list: list[dict[str, Any]]
process_rule: dict[str, Any]
notion_info_list: list[NotionEstimateWorkspacePayload] = Field(min_length=1)
process_rule: dict[str, object]
doc_form: str = Field(default="text_model")
doc_language: str = Field(default="English")
@@ -72,7 +78,7 @@ class DataSourceIntegratePageResponse(ResponseModel):
page_id: str
page_icon: DataSourceIntegrateIconResponse | None
parent_id: str
type: str
type: NotionPageType
class DataSourceIntegrateWorkspaceResponse(ResponseModel):
@@ -106,7 +112,7 @@ class NotionIntegratePageResponse(ResponseModel):
page_id: str
page_icon: DataSourceIntegrateIconResponse | None
parent_id: str | None
type: str
page_type: NotionPageType = Field(alias="type")
is_bound: bool
@@ -132,325 +138,204 @@ register_response_schema_models(
)
@console_ns.route(
"/data-source/integrates",
"/data-source/integrates/<uuid:binding_id>/<string:action>",
)
class DataSourceApi(Resource):
@setup_required
@login_required
@is_admin_or_owner_required
@rbac_permission_required(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace()))
@account_initialization_required
def _raise_dataset_access_error(error: DatasetNotFoundError | DatasetAccessDeniedError) -> Never:
if isinstance(error, DatasetNotFoundError):
raise NotFoundError(description="Dataset not found.") from None
raise DatasetAccessDeniedRequestError(description=str(error)) from None
@console_ns.route("/data-source/integrates")
class DataSourceIntegrationListApi(Resource):
@console_ns.response(200, "Success", console_ns.models[DataSourceIntegrateListResponse.__name__])
@with_current_tenant_id
def get(self, current_tenant_id: str) -> tuple[dict[str, Any], int]:
# get workspace data source integrates
data_source_integrates = db.session.scalars(
select(DataSourceOauthBinding).where(
DataSourceOauthBinding.tenant_id == current_tenant_id,
DataSourceOauthBinding.disabled == False,
)
).all()
@console_account_admission(
allowed_roles=_ADMIN_OR_OWNER_ROLES,
rbac_checks=(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace()),),
)
def get(self, request_context: RequestContext) -> tuple[dict[str, object], int]:
bindings = application_services().data_sources.bindings.list_integrations(request_context)
base_url = dify_config.CONSOLE_API_URL.rstrip("/")
data = [
{
"id": binding.id,
"provider": binding.provider,
"created_at": binding.created_at,
"is_bound": True,
"disabled": binding.disabled,
"source_info": binding.source_info,
"link": f"{base_url}{_DATA_SOURCE_OAUTH_BASE_PATH}/{binding.provider}",
}
for binding in bindings
]
return dump_response(DataSourceIntegrateListResponse, {"data": data}), 200
base_url = request.url_root.rstrip("/")
data_source_oauth_base_path = "/console/api/oauth/data-source"
providers = ["notion"]
integrate_data: list[dict[str, bool | datetime | dict[str, Any] | str | None]] = []
for provider in providers:
# existing_integrate = next((ai for ai in data_source_integrates if ai.provider == provider), None)
existing_integrates = filter(lambda item: item.provider == provider, data_source_integrates)
if existing_integrates:
for existing_integrate in list(existing_integrates):
integrate_data.append(
{
"id": existing_integrate.id,
"provider": provider,
"created_at": existing_integrate.created_at,
"is_bound": True,
"disabled": existing_integrate.disabled,
"source_info": existing_integrate.source_info,
"link": f"{base_url}{data_source_oauth_base_path}/{provider}",
}
)
else:
integrate_data.append(
{
"id": None,
"provider": provider,
"created_at": None,
"source_info": None,
"is_bound": False,
"disabled": None,
"link": f"{base_url}{data_source_oauth_base_path}/{provider}",
}
)
return dump_response(DataSourceIntegrateListResponse, {"data": integrate_data}), 200
@setup_required
@login_required
@is_admin_or_owner_required
@rbac_permission_required(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace()))
@account_initialization_required
@console_ns.route("/data-source/integrates/<uuid:binding_id>/<string:action>")
class DataSourceIntegrationApi(Resource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@with_current_tenant_id
@with_session
@console_account_admission(
allowed_roles=_ADMIN_OR_OWNER_ROLES,
rbac_checks=(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace()),),
)
def patch(
self, session: Session, current_tenant_id: str, binding_id: UUID, action: Literal["enable", "disable"]
) -> tuple[dict[str, str], int]:
binding_id_str = str(binding_id)
data_source_binding = session.scalar(
select(DataSourceOauthBinding).where(
DataSourceOauthBinding.id == binding_id_str, DataSourceOauthBinding.tenant_id == current_tenant_id
self,
request_context: RequestContext,
binding_id: UUID,
action: str,
) -> tuple[dict[str, object], int]:
try:
validated_action = BindingAction(action)
except ValueError as error:
raise InvalidActionError() from error
try:
application_services().data_sources.bindings.change_state(
request_context,
str(binding_id),
validated_action,
)
)
if data_source_binding is None:
raise NotFound("Data source binding not found.")
# enable binding
match action:
case "enable":
if data_source_binding.disabled:
data_source_binding.disabled = False
data_source_binding.updated_at = naive_utc_now()
else:
raise ValueError("Data source is not disabled.")
# disable binding
case "disable":
if not data_source_binding.disabled:
data_source_binding.disabled = True
data_source_binding.updated_at = naive_utc_now()
else:
raise ValueError("Data source is disabled.")
return {"result": "success"}, 200
except DataSourceBindingNotFoundError as error:
raise NotFoundError(description="Data source binding not found.") from error
except DataSourceBindingStateError as error:
raise InvalidArgumentError(description=str(error)) from error
return dump_response(SimpleResultResponse, {"result": "success"}), 200
@console_ns.route("/notion/pre-import/pages")
class DataSourceNotionListApi(Resource):
@setup_required
@login_required
@account_initialization_required
@console_ns.doc(params=query_params_from_model(DataSourceNotionListQuery))
@console_ns.response(200, "Success", console_ns.models[NotionIntegrateInfoListResponse.__name__])
@with_current_user
@with_current_tenant_id
@with_session(write=False)
@console_account_admission()
@model_validate(DataSourceNotionListQuery)
def get(
self,
req_data: DataSourceNotionListQuery,
session: Session,
current_tenant_id: str,
current_user: Account,
) -> tuple[dict[str, Any], int]:
datasource_provider_service = DatasourceProviderService()
credential = datasource_provider_service.get_datasource_credentials(
tenant_id=current_tenant_id,
credential_id=req_data.credential_id,
provider="notion_datasource",
plugin_id="langgenius/notion_datasource",
)
if not credential:
raise NotFound("Credential not found.")
exist_page_ids = []
# import notion in the exist dataset
if req_data.dataset_id:
dataset = DatasetService.get_dataset(req_data.dataset_id, session)
if not dataset:
raise NotFound("Dataset not found.")
if dataset.data_source_type != "notion_import":
raise ValueError("Dataset is not notion type.")
documents = session.scalars(
select(Document).where(
Document.dataset_id == req_data.dataset_id,
Document.tenant_id == current_tenant_id,
Document.data_source_type == "notion_import",
Document.enabled.is_(True),
)
).all()
if documents:
for document in documents:
data_source_info = json.loads(document.data_source_info)
exist_page_ids.append(data_source_info["notion_page_id"])
# get all authorized pages
from core.datasource.datasource_manager import DatasourceManager
datasource_runtime = DatasourceManager.get_datasource_runtime(
provider_id="langgenius/notion_datasource/notion_datasource",
datasource_name="notion_datasource",
tenant_id=current_tenant_id,
datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
)
datasource_provider_service = DatasourceProviderService()
if credential:
datasource_runtime.runtime.credentials = credential
datasource_runtime = cast(OnlineDocumentDatasourcePlugin, datasource_runtime)
online_document_result: Generator[OnlineDocumentPagesMessage, None, None] = (
datasource_runtime.get_online_document_pages(
user_id=current_user.id,
datasource_parameters={},
provider_type=datasource_runtime.datasource_provider_type(),
)
)
request_context: RequestContext,
) -> tuple[dict[str, object], int]:
try:
pages = []
workspace_info = {}
for message in online_document_result:
result = message.result
for info in result:
workspace_info = {
"workspace_id": info.workspace_id,
"workspace_name": info.workspace_name,
"workspace_icon": info.workspace_icon,
}
for page in info.pages:
page_info = {
"page_id": page.page_id,
"page_name": page.page_name,
"type": page.type,
"parent_id": page.parent_id,
"is_bound": page.page_id in exist_page_ids,
"page_icon": page.page_icon,
}
pages.append(page_info)
except Exception as e:
raise e
notion_info = [{**workspace_info, "pages": pages}] if workspace_info else []
return dump_response(NotionIntegrateInfoListResponse, {"notion_info": notion_info}), 200
result = application_services().data_sources.notion_imports.list_pages(
request_context,
credential_id=req_data.credential_id,
dataset_id=req_data.dataset_id,
)
except NotionImportCredentialUnavailableError as error:
raise NotFoundError(description="Credential not found.") from error
except (DatasetNotFoundError, DatasetAccessDeniedError) as error:
_raise_dataset_access_error(error)
except DatasetIsNotNotionSourceError as error:
raise InvalidArgumentError(description="Dataset is not notion type.") from error
return dump_response(NotionIntegrateInfoListResponse, {"notion_info": result.workspaces}), 200
@console_ns.route("/notion/pages/<uuid:page_id>/<string:page_type>/preview")
class DataSourceNotionPreviewApi(Resource):
"""Preview one authorized Notion page through the datasource credential."""
@setup_required
@login_required
@account_initialization_required
@console_ns.doc(params=query_params_from_model(DataSourceNotionPreviewQuery))
@console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__])
@with_current_tenant_id
@console_account_admission()
@model_validate(DataSourceNotionPreviewQuery)
def get(
self,
req_data: DataSourceNotionPreviewQuery,
current_tenant_id: str,
request_context: RequestContext,
page_id: UUID,
page_type: str,
) -> tuple[dict[str, str], int]:
datasource_provider_service = DatasourceProviderService()
credential = datasource_provider_service.get_datasource_credentials(
tenant_id=current_tenant_id,
credential_id=req_data.credential_id,
provider="notion_datasource",
plugin_id="langgenius/notion_datasource",
)
page_id_str = str(page_id)
extractor = NotionExtractor(
notion_workspace_id="",
notion_obj_id=page_id_str,
notion_page_type=page_type,
notion_access_token=credential.get("integration_secret"),
tenant_id=current_tenant_id,
)
text_docs = extractor.extract()
return {"content": "\n".join([doc.page_content for doc in text_docs])}, 200
) -> tuple[dict[str, object], int]:
try:
validated_page_type = NotionPageType(page_type)
except ValueError as error:
raise InvalidArgumentError(description="Invalid Notion page type.") from error
try:
content = application_services().data_sources.notion_imports.preview_page(
request_context,
credential_id=req_data.credential_id,
page_id=str(page_id),
page_type=validated_page_type,
)
except NotionImportCredentialUnavailableError as error:
raise NotFoundError(description="Credential not found.") from error
return dump_response(TextContentResponse, {"content": content}), 200
@console_ns.route("/datasets/notion-indexing-estimate")
class DataSourceNotionIndexingEstimateApi(Resource):
"""Estimate indexing work for selected Notion pages."""
@setup_required
@login_required
@account_initialization_required
@console_ns.expect(console_ns.models[NotionEstimatePayload.__name__])
@console_ns.response(200, "Success", console_ns.models[IndexingEstimate.__name__])
@with_current_tenant_id
@with_session
@console_account_admission()
@model_validate(NotionEstimatePayload)
def post(
self,
req_data: NotionEstimatePayload,
session: Session,
current_tenant_id: str,
) -> tuple[dict[str, Any], int]:
args = req_data.model_dump()
# validate args
DocumentService.estimate_args_validate(args)
notion_info_list = req_data.notion_info_list
extract_settings = []
for notion_info in notion_info_list:
workspace_id = notion_info["workspace_id"]
credential_id = notion_info.get("credential_id")
for page in notion_info["pages"]:
extract_setting = ExtractSetting(
datasource_type=DatasourceType.NOTION,
notion_info=NotionInfo.model_validate(
{
"credential_id": credential_id,
"notion_workspace_id": workspace_id,
"notion_obj_id": page["page_id"],
"notion_page_type": page["type"],
"tenant_id": current_tenant_id,
}
),
document_model=args["doc_form"],
request_context: RequestContext,
) -> tuple[dict[str, object], int]:
command = NewSourcesEstimateCommand(
sources=tuple(
NotionEstimateSource(
workspace_id=workspace.workspace_id,
page_id=page.page_id,
page_type=page.page_type,
credential_id=workspace.credential_id,
)
extract_settings.append(extract_setting)
indexing_runner = IndexingRunner()
response = indexing_runner.indexing_estimate(
tenant_id=current_tenant_id,
extract_settings=extract_settings,
tmp_processing_rule=args["process_rule"],
doc_form=args["doc_form"],
doc_language=args["doc_language"],
session=session,
for workspace in req_data.notion_info_list
for page in workspace.pages
),
process_rule=req_data.process_rule,
doc_form=req_data.doc_form,
doc_language=req_data.doc_language,
)
try:
response = application_services().knowledge.indexing_estimates.estimate_new_sources(
request_context,
command,
)
except IndexingEstimateCredentialUnavailableError as error:
raise NotFoundError(description="Credential not found.") from error
except EstimateSourceNotFoundError as error:
raise NotFoundError(description=str(error)) from error
except IndexingEstimateProviderUnavailableError as error:
raise ProviderNotInitializeError(str(error)) from error
except IndexingEstimateExecutionError as error:
raise IndexingEstimateError(str(error)) from error
return dump_response(IndexingEstimate, response), 200
@console_ns.route("/datasets/<uuid:dataset_id>/notion/sync")
class DataSourceNotionDatasetSyncApi(Resource):
@setup_required
@login_required
@account_initialization_required
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID) -> tuple[dict[str, str], int]:
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
documents = DocumentService.get_document_by_dataset_id(dataset_id_str, session)
for document in documents:
document_indexing_sync_task.delay(dataset_id_str, document.id)
return {"result": "success"}, 200
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
def get(self, request_context: RequestContext, dataset_id: UUID) -> tuple[dict[str, object], int]:
try:
application_services().knowledge.document_sync.sync_dataset(request_context, str(dataset_id))
except (DatasetNotFoundError, DatasetAccessDeniedError) as error:
_raise_dataset_access_error(error)
return dump_response(SimpleResultResponse, {"result": "success"}), 200
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/notion/sync")
class DataSourceNotionDocumentSyncApi(Resource):
@setup_required
@login_required
@account_initialization_required
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID, document_id: UUID) -> tuple[dict[str, str], int]:
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if document is None:
raise NotFound("Document not found.")
document_indexing_sync_task.delay(dataset_id_str, document_id_str)
return {"result": "success"}, 200
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
def get(
self,
request_context: RequestContext,
dataset_id: UUID,
document_id: UUID,
) -> tuple[dict[str, object], int]:
try:
application_services().knowledge.document_sync.sync_document(
request_context,
str(dataset_id),
str(document_id),
)
except (DatasetNotFoundError, DatasetAccessDeniedError) as error:
_raise_dataset_access_error(error)
except SyncDocumentNotFoundError as error:
raise NotFoundError(description="Document not found.") from error
except SyncDocumentSourceError as error:
raise InvalidArgumentError(description=str(error)) from error
return dump_response(SimpleResultResponse, {"result": "success"}), 200
api/controllers/console/datasets/datasets.py
+260 -763
View File
File diff suppressed because it is too large. Load diff
api/controllers/console/datasets/datasets_document.py
+411 -1304
View File
File diff suppressed because it is too large. Load diff
api/controllers/console/datasets/datasets_segments.py
+265 -649
View File
File diff suppressed because it is too large. Load diff
api/controllers/console/datasets/error.py
+6
View File
@@ -25,6 +25,12 @@ class InvalidActionError(BaseHTTPException):
code = 400
class DatasetAccessDeniedRequestError(BaseHTTPException):
error_code = "dataset_access_denied"
description = "Dataset access denied."
code = 403
class DocumentAlreadyFinishedError(BaseHTTPException):
error_code = "document_already_finished"
description = "The document has been processed. Please refresh the page or go to the document details."
api/controllers/console/datasets/external.py
+93 -234
View File
@@ -1,16 +1,13 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from typing import Any, Never
from uuid import UUID
from flask_restx import Resource
from pydantic import AliasChoices, BaseModel, Field, field_validator
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
import services
from controllers.common.fields import UsageCountResponse
from controllers.common.rbac import DatasetId, RBACCheck, Workspace
from controllers.common.schema import (
@@ -18,31 +15,42 @@ from controllers.common.schema import (
register_response_schema_models,
register_schema_models,
)
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.datasets.error import DatasetNameDuplicateError
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import (
RBACPermission,
account_initialization_required,
edit_permission_required,
model_validate,
rbac_permission_required,
setup_required,
with_current_tenant_id,
with_current_user,
)
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from fields.dataset_fields import (
DatasetDetailResponse,
)
from libs.helper import dump_response
from libs.login import login_required
from libs.pagination import clamp_pagination
from models import Account
from models.dataset import ExternalKnowledgeApis
from services.dataset_service import DatasetService
from services.enterprise import rbac_service as enterprise_rbac_service
from machinery.context import RequestContext
from models.account import TenantAccountRole
from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload
from services.external_knowledge_service import ExternalDatasetService
from services.hit_testing_service import HitTestingService
from services.errors.dataset import DatasetNameDuplicateError as DatasetNameDuplicateFailure
from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError
from services.knowledge.external.application import ExternalHitTestingError, ExternalTemplateNotFoundError
_DATASET_EDIT_ROLES = frozenset(
{TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.DATASET_OPERATOR}
)
_EXTERNAL_CONNECT_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR})
def _raise_external_error(error: Exception) -> Never:
if isinstance(error, (DatasetNotFoundError, ExternalTemplateNotFoundError)):
raise NotFound(str(error)) from error
if isinstance(error, DatasetAccessDeniedError):
raise Forbidden(str(error)) from error
if isinstance(error, DatasetNameDuplicateFailure):
raise DatasetNameDuplicateError() from error
if isinstance(error, ExternalHitTestingError):
raise InternalServerError(str(error)) from error
raise error
class ExternalKnowledgeApiPayload(BaseModel):
@@ -85,28 +93,6 @@ class ExternalKnowledgeApiResponse(ResponseModel):
return value
@dataclass(frozen=True)
class ExternalKnowledgeApiResponseSource:
external_knowledge_api: ExternalKnowledgeApis
session: Session
@property
def dataset_bindings(self) -> Any:
return self.external_knowledge_api.get_dataset_bindings(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.external_knowledge_api, name) # guard-ignore: no-new-getattr -- delegates model fields
def external_knowledge_api_response(
external_knowledge_api: ExternalKnowledgeApis, *, session: Session
) -> ExternalKnowledgeApiResponse:
return ExternalKnowledgeApiResponse.model_validate(
ExternalKnowledgeApiResponseSource(external_knowledge_api=external_knowledge_api, session=session),
from_attributes=True,
)
class ExternalKnowledgeApiListResponse(ResponseModel):
data: list[ExternalKnowledgeApiResponse]
has_more: bool
@@ -161,25 +147,13 @@ class ExternalApiTemplateListApi(Resource):
"External API templates retrieved successfully",
console_ns.models[ExternalKnowledgeApiListResponse.__name__],
)
@setup_required
@login_required
@with_current_tenant_id
@account_initialization_required
@with_session(write=False)
@console_account_admission()
@model_validate(ExternalApiTemplateListQuery)
def get(self, req_data: ExternalApiTemplateListQuery, session: Session, current_tenant_id: str):
effective_page, effective_limit = clamp_pagination(req_data.page, req_data.limit, 100)
external_knowledge_apis, total = ExternalDatasetService.get_external_knowledge_apis(
effective_page, effective_limit, current_tenant_id, req_data.keyword, session=session
def get(self, req_data: ExternalApiTemplateListQuery, request_context: RequestContext):
result = application_services().knowledge.external.list_templates(
request_context, page=req_data.page, limit=req_data.limit, keyword=req_data.keyword
)
return ExternalKnowledgeApiListResponse(
data=[external_knowledge_api_response(item, session=session) for item in external_knowledge_apis],
has_more=effective_page * effective_limit < total,
limit=effective_limit,
total=total,
page=effective_page,
).model_dump(mode="json"), 200
return dump_response(ExternalKnowledgeApiListResponse, result), 200
@console_ns.doc("create_external_api_template")
@console_ns.doc(description="Create external knowledge API template")
@@ -190,48 +164,16 @@ class ExternalApiTemplateListApi(Resource):
console_ns.models[ExternalKnowledgeApiResponse.__name__],
)
@console_ns.response(403, "Permission denied")
@setup_required
@login_required
@account_initialization_required
@with_current_user
@with_current_tenant_id
@with_session
@console_account_admission(allowed_roles=_DATASET_EDIT_ROLES)
@model_validate(ExternalKnowledgeApiPayload)
def post(
self,
req_data: ExternalKnowledgeApiPayload,
session: Session,
current_tenant_id: str,
current_user: Account,
):
ExternalDatasetService.validate_api_list(req_data.settings)
# The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
if not current_user.is_dataset_editor:
raise Forbidden()
def post(self, req_data: ExternalKnowledgeApiPayload, request_context: RequestContext):
try:
external_knowledge_api = ExternalDatasetService.create_external_knowledge_api(
tenant_id=current_tenant_id,
user_id=current_user.id,
args=req_data.model_dump(),
session=session,
result = application_services().knowledge.external.create_template(
request_context, name=req_data.name, settings=req_data.settings
)
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
# if dify_config.RBAC_ENABLED:
# enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist(
# current_tenant_id,
# current_user.id,
# external_knowledge_api.id,
# enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=True),
# )
# initialize_created_app_rbac_access_task.delay(
# current_tenant_id, current_user.id, dataset_id=external_knowledge_api.id)
return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 201
except Exception as error:
_raise_external_error(error)
return dump_response(ExternalKnowledgeApiResponse, result), 201
@console_ns.route("/datasets/external-knowledge-api/<uuid:external_knowledge_api_id>")
@@ -245,20 +187,15 @@ class ExternalApiTemplateApi(Resource):
console_ns.models[ExternalKnowledgeApiResponse.__name__],
)
@console_ns.response(404, "Template not found")
@setup_required
@login_required
@account_initialization_required
@with_current_tenant_id
@with_session
def get(self, session: Session, current_tenant_id: str, external_knowledge_api_id: UUID):
external_knowledge_api_id_str = str(external_knowledge_api_id)
external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(
external_knowledge_api_id=external_knowledge_api_id_str, tenant_id=current_tenant_id, session=session
)
if external_knowledge_api is None:
raise NotFound("API template not found.")
return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 200
@console_account_admission()
def get(self, request_context: RequestContext, external_knowledge_api_id: UUID):
try:
result = application_services().knowledge.external.get_template(
request_context, template_id=str(external_knowledge_api_id)
)
except Exception as error:
_raise_external_error(error)
return dump_response(ExternalKnowledgeApiResponse, result), 200
@console_ns.doc("update_external_api_template")
@console_ns.doc(description="Update external knowledge API template")
@@ -270,51 +207,31 @@ class ExternalApiTemplateApi(Resource):
console_ns.models[ExternalKnowledgeApiResponse.__name__],
)
@console_ns.response(404, "Template not found")
@setup_required
@login_required
@account_initialization_required
@with_current_user
@with_current_tenant_id
@with_session
@console_account_admission()
@model_validate(ExternalKnowledgeApiPayload)
def patch(
self,
req_data: ExternalKnowledgeApiPayload,
session: Session,
current_tenant_id: str,
current_user: Account,
external_knowledge_api_id: UUID,
self, req_data: ExternalKnowledgeApiPayload, request_context: RequestContext, external_knowledge_api_id: UUID
):
external_knowledge_api_id_str = str(external_knowledge_api_id)
try:
result = application_services().knowledge.external.update_template(
request_context,
template_id=str(external_knowledge_api_id),
name=req_data.name,
settings=req_data.settings,
)
except Exception as error:
_raise_external_error(error)
return dump_response(ExternalKnowledgeApiResponse, result), 200
ExternalDatasetService.validate_api_list(req_data.settings)
external_knowledge_api = ExternalDatasetService.update_external_knowledge_api(
tenant_id=current_tenant_id,
user_id=current_user.id,
external_knowledge_api_id=external_knowledge_api_id_str,
args=req_data.model_dump(),
session=session,
)
return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 200
@setup_required
@login_required
@account_initialization_required
@console_ns.response(204, "External knowledge API deleted successfully")
@with_current_user
@with_current_tenant_id
@with_session
def delete(self, session: Session, current_tenant_id: str, current_user: Account, external_knowledge_api_id: UUID):
external_knowledge_api_id_str = str(external_knowledge_api_id)
if not (current_user.has_edit_permission or current_user.is_dataset_operator):
raise Forbidden()
ExternalDatasetService.delete_external_knowledge_api(
current_tenant_id, external_knowledge_api_id_str, session=session
)
@console_account_admission(allowed_roles=_DATASET_EDIT_ROLES)
def delete(self, request_context: RequestContext, external_knowledge_api_id: UUID):
try:
application_services().knowledge.external.delete_template(
request_context, template_id=str(external_knowledge_api_id)
)
except Exception as error:
_raise_external_error(error)
return "", 204
@@ -324,18 +241,12 @@ class ExternalApiUseCheckApi(Resource):
@console_ns.doc(description="Check if external knowledge API is being used")
@console_ns.doc(params={"external_knowledge_api_id": "External knowledge API ID"})
@console_ns.response(200, "Usage check completed successfully", console_ns.models[UsageCountResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@with_current_tenant_id
@with_session
def get(self, session: Session, current_tenant_id: str, external_knowledge_api_id: UUID):
external_knowledge_api_id_str = str(external_knowledge_api_id)
external_knowledge_api_is_using, count = ExternalDatasetService.external_knowledge_api_use_check(
external_knowledge_api_id_str, current_tenant_id, session=session
@console_account_admission()
def get(self, request_context: RequestContext, external_knowledge_api_id: UUID):
is_using, count = application_services().knowledge.external.template_usage(
request_context, template_id=str(external_knowledge_api_id)
)
return UsageCountResponse(is_using=external_knowledge_api_is_using, count=count).model_dump(mode="json"), 200
return dump_response(UsageCountResponse, {"is_using": is_using, "count": count}), 200
@console_ns.route("/datasets/external")
@@ -348,50 +259,17 @@ class ExternalDatasetCreateApi(Resource):
)
@console_ns.response(400, "Invalid parameters")
@console_ns.response(403, "Permission denied")
@setup_required
@login_required
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EXTERNAL_CONNECT, Workspace()))
@with_current_user
@with_current_tenant_id
@with_session
@console_account_admission(
allowed_roles=_EXTERNAL_CONNECT_ROLES,
rbac_checks=(RBACCheck(RBACPermission.DATASET_EXTERNAL_CONNECT, Workspace()),),
)
@model_validate(ExternalDatasetCreatePayload)
def post(
self,
req_data: ExternalDatasetCreatePayload,
session: Session,
current_tenant_id: str,
current_user: Account,
):
# The role of the current user in the ta table must be admin, owner, or editor
# The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
if not current_user.is_dataset_editor:
raise Forbidden()
def post(self, req_data: ExternalDatasetCreatePayload, request_context: RequestContext):
try:
dataset = ExternalDatasetService.create_external_dataset(
tenant_id=current_tenant_id,
user_id=current_user.id,
args=req_data,
session=session,
)
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
dataset_id_str = str(dataset.id)
permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get(
str(current_tenant_id),
current_user.id,
[dataset_id_str],
session=session,
)
data = DatasetDetailResponse.model_validate(
dataset_detail_response_source(dataset, session=session)
).model_dump(mode="json")
data["permission_keys"] = permission_keys_map.get(dataset_id_str, [])
return data, 201
result = application_services().knowledge.external.create_dataset(request_context, payload=req_data)
except Exception as error:
_raise_external_error(error)
return dump_response(DatasetDetailResponse, result), 201
@console_ns.route("/datasets/<uuid:dataset_id>/external-hit-testing")
@@ -407,36 +285,17 @@ class ExternalKnowledgeHitTestingApi(Resource):
)
@console_ns.response(404, "Dataset not found")
@console_ns.response(400, "Invalid parameters")
@setup_required
@login_required
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_PIPELINE_TEST, DatasetId()))
@with_session
@console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_PIPELINE_TEST, DatasetId()),))
@model_validate(ExternalHitTestingPayload)
def post(self, req_data: ExternalHitTestingPayload, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
def post(self, req_data: ExternalHitTestingPayload, request_context: RequestContext, dataset_id: UUID):
try:
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
HitTestingService.hit_testing_args_check(req_data.model_dump())
try:
response = HitTestingService.external_retrieve(
session=session,
dataset=dataset,
result = application_services().knowledge.external.hit_testing(
request_context,
dataset_id=str(dataset_id),
query=req_data.query,
account=current_user,
external_retrieval_model=req_data.external_retrieval_model,
metadata_filtering_conditions=req_data.metadata_filtering_conditions,
retrieval_model=req_data.external_retrieval_model,
metadata_filters=req_data.metadata_filtering_conditions,
)
return dump_response(ExternalHitTestingResponse, response)
except Exception as e:
raise InternalServerError(str(e))
except Exception as error:
_raise_external_error(error)
return dump_response(ExternalHitTestingResponse, result)
api/controllers/console/datasets/hit_testing_base.py
+2 -2
View File
@@ -23,9 +23,9 @@ from graphon.model_runtime.errors.invoke import InvokeError
from libs.login import resolve_account_fallback
from models.account import Account
from models.dataset import Dataset
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.knowledge_entities import ExternalRetrievalModel, RetrievalModel
from services.hit_testing_service import HitTestingService
from services.knowledge.dataset_service import DatasetService
from services.knowledge.entities.knowledge_entities import ExternalRetrievalModel, RetrievalModel
logger = logging.getLogger(__name__)
api/controllers/console/datasets/metadata.py
+84 -168
View File
@@ -1,43 +1,47 @@
from typing import Literal
from collections.abc import Callable
from functools import wraps
from uuid import UUID
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, NotFound
import services
from controllers.common.controller_schemas import MetadataUpdatePayload
from controllers.common.rbac import DatasetId, RBACCheck
from controllers.common.schema import register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.wraps import (
RBACPermission,
account_initialization_required,
enterprise_license_required,
model_validate,
rbac_permission_required,
setup_required,
with_current_tenant_id,
with_current_user,
)
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import model_validate
from core.rbac import RBACPermission
from extensions.ext_application_services import application_services
from fields.dataset_fields import (
DatasetMetadataBuiltInFieldsResponse,
DatasetMetadataListResponse,
DatasetMetadataResponse,
)
from libs.helper import dump_response
from libs.login import login_required
from models.account import Account
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.knowledge_entities import (
from machinery.context import RequestContext
from services.errors.metadata import MetadataResourceNotFoundError
from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError
from services.knowledge.entities.knowledge_entities import (
DocumentMetadataOperation,
MetadataArgs,
MetadataDetail,
MetadataOperationData,
)
from services.errors.metadata import MetadataResourceNotFoundError
from services.metadata_service import MetadataService
def metadata_errors[**P, T](method: Callable[P, T]) -> Callable[P, T]:
@wraps(method)
def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
try:
return method(*args, **kwargs)
except (DatasetNotFoundError, MetadataResourceNotFoundError) as error:
raise NotFound(str(error)) from error
except DatasetAccessDeniedError as error:
raise Forbidden(str(error)) from error
return wrapped
register_schema_models(
console_ns, MetadataArgs, MetadataOperationData, MetadataUpdatePayload, DocumentMetadataOperation, MetadataDetail
@@ -52,199 +56,111 @@ register_response_schema_models(
@console_ns.route("/datasets/<uuid:dataset_id>/metadata")
class DatasetMetadataCreateApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.response(201, "Metadata created successfully", console_ns.models[DatasetMetadataResponse.__name__])
@console_ns.expect(console_ns.models[MetadataArgs.__name__])
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
@metadata_errors
@model_validate(MetadataArgs)
def post(
self,
req_data: MetadataArgs,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.create_metadata(
dataset_id_str, req_data, current_user, current_tenant_id, session=session
)
def post(self, req_data: MetadataArgs, request_context: RequestContext, dataset_id: UUID):
service = application_services().knowledge.metadata
ref = service.require_dataset(request_context, str(dataset_id))
metadata = service.create_metadata(ref, req_data, actor_id=request_context.account_id)
return dump_response(DatasetMetadataResponse, metadata), 201
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.response(
200, "Metadata retrieved successfully", console_ns.models[DatasetMetadataListResponse.__name__]
)
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()))
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
metadata = MetadataService.get_dataset_metadatas(dataset, session)
return dump_response(DatasetMetadataListResponse, metadata), 200
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),),
)
@metadata_errors
def get(self, request_context: RequestContext, dataset_id: UUID):
service = application_services().knowledge.metadata
ref = service.require_dataset(request_context, str(dataset_id))
return dump_response(DatasetMetadataListResponse, service.get_dataset_metadatas(ref)), 200
@console_ns.route("/datasets/<uuid:dataset_id>/metadata/<uuid:metadata_id>")
class DatasetMetadataApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.response(200, "Metadata updated successfully", console_ns.models[DatasetMetadataResponse.__name__])
@console_ns.expect(console_ns.models[MetadataUpdatePayload.__name__])
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
@metadata_errors
@model_validate(MetadataUpdatePayload)
def patch(
self,
req_data: MetadataUpdatePayload,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
metadata_id: UUID,
self, req_data: MetadataUpdatePayload, request_context: RequestContext, dataset_id: UUID, metadata_id: UUID
):
name = req_data.name
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.update_metadata_name(dataset, metadata_id_str, name, current_user, session=session)
service = application_services().knowledge.metadata
ref = service.require_dataset(request_context, str(dataset_id))
metadata = service.update_metadata_name(
ref, str(metadata_id), req_data.name, actor_id=request_context.account_id
)
return dump_response(DatasetMetadataResponse, metadata), 200
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.response(204, "Metadata deleted successfully")
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session
def delete(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
metadata_id: UUID,
):
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
MetadataService.delete_metadata(dataset, metadata_id_str, session)
# Frontend callers only await success and invalidate metadata caches; no response body is consumed.
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
@metadata_errors
def delete(self, request_context: RequestContext, dataset_id: UUID, metadata_id: UUID):
service = application_services().knowledge.metadata
ref = service.require_dataset(request_context, str(dataset_id))
service.delete_metadata(ref, str(metadata_id))
return "", 204
@console_ns.route("/datasets/metadata/built-in")
class DatasetMetadataBuiltInFieldApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.response(
200,
"Built-in fields retrieved successfully",
console_ns.models[DatasetMetadataBuiltInFieldsResponse.__name__],
)
def get(self):
built_in_fields = MetadataService.get_built_in_fields()
return dump_response(DatasetMetadataBuiltInFieldsResponse, {"fields": built_in_fields}), 200
@console_account_admission()
@metadata_errors
def get(self, request_context: RequestContext):
fields = application_services().knowledge.metadata.get_built_in_fields()
return dump_response(DatasetMetadataBuiltInFieldsResponse, {"fields": fields}), 200
@console_ns.route("/datasets/<uuid:dataset_id>/metadata/built-in/<string:action>")
class DatasetMetadataBuiltInFieldActionApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.response(204, "Action completed successfully")
@with_current_user
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session
def post(self, session: Session, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable"]):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
match action:
case "enable":
MetadataService.enable_built_in_field(dataset, session)
case "disable":
MetadataService.disable_built_in_field(dataset, session)
# Frontend callers only await success and invalidate metadata caches; no response body is consumed.
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
@metadata_errors
def post(self, request_context: RequestContext, dataset_id: UUID, action: str):
service = application_services().knowledge.metadata
ref = service.require_dataset(request_context, str(dataset_id))
if action == "enable":
service.enable_built_in_field(ref)
elif action == "disable":
service.disable_built_in_field(ref)
else:
raise ValueError("Invalid action.")
return "", 204
@console_ns.route("/datasets/<uuid:dataset_id>/documents/metadata")
class DocumentMetadataEditApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@console_ns.expect(console_ns.models[MetadataOperationData.__name__])
@console_ns.response(
204,
"Documents metadata updated successfully",
)
@console_ns.response(404, "Dataset, document, or metadata not found")
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()))
@with_session
@console_account_admission(
rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),),
)
@metadata_errors
@model_validate(MetadataOperationData)
def post(
self,
req_data: MetadataOperationData,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
):
dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), current_tenant_id, session=session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
try:
MetadataService.update_documents_metadata(dataset, req_data, current_user, session=session)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
# Frontend callers only await success and invalidate caches; no response body is consumed.
def post(self, req_data: MetadataOperationData, request_context: RequestContext, dataset_id: UUID):
service = application_services().knowledge.metadata
ref = service.require_dataset(request_context, str(dataset_id))
service.update_documents_metadata(ref, req_data, actor_id=request_context.account_id)
return "", 204
api/controllers/console/datasets/rag_pipeline/datasource_auth.py
+22 -16
View File
@@ -24,6 +24,7 @@ from core.entities.provider_entities import ProviderConfig
from core.plugin.entities.plugin_daemon import PluginOAuthAuthorizationUrlResponse
from core.plugin.impl.oauth import OAuthHandler
from core.tools.entities.common_entities import I18nObject
from extensions.ext_application_services import application_services
from extensions.ext_database import db
from fields.base import ResponseModel
from graphon.model_runtime.errors.validate import CredentialsValidateFailedError
@@ -32,7 +33,6 @@ from libs.login import login_required
from models import Account
from models.enums import PermissionEnum
from models.provider_ids import DatasourceProviderID
from services.datasource_provider_service import DatasourceProviderService
from services.plugin.oauth_service import OAuthProxyService
@@ -176,7 +176,7 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource):
datasource_provider_id = DatasourceProviderID(provider_id)
provider_name = datasource_provider_id.provider_name
plugin_id = datasource_provider_id.plugin_id
oauth_config = DatasourceProviderService().get_oauth_client(
oauth_config = application_services().data_sources.providers.get_oauth_client(
tenant_id=tenant_id,
datasource_provider_id=datasource_provider_id,
)
@@ -246,7 +246,7 @@ class DatasourceOAuthCallback(Resource):
tenant_id: str = context["tenant_id"]
datasource_provider_id = DatasourceProviderID(provider_id)
plugin_id = datasource_provider_id.plugin_id
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
oauth_client_params = datasource_provider_service.get_oauth_client(
tenant_id=tenant_id,
datasource_provider_id=datasource_provider_id,
@@ -315,7 +315,7 @@ class DatasourceAuth(Resource):
@model_validate(DatasourceCredentialPayload)
def post(self, req_data: DatasourceCredentialPayload, current_tenant_id: str, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
try:
datasource_provider_service.add_datasource_api_key_provider(
@@ -342,14 +342,14 @@ class DatasourceAuth(Resource):
@with_current_tenant_id
def get(self, current_tenant_id: str, user: Account, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasources = datasource_provider_service.list_datasource_credentials(
tenant_id=current_tenant_id,
provider=datasource_provider_id.provider_name,
plugin_id=datasource_provider_id.plugin_id,
user=user,
session=db.session(),
credential_query=application_services().credential_queries,
)
return dump_response(DatasourceCredentialListResponse, {"result": datasources}), 200
@@ -370,7 +370,7 @@ class DatasourceAuthDeleteApi(Resource):
plugin_id = datasource_provider_id.plugin_id
provider_name = datasource_provider_id.provider_name
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasource_provider_service.remove_datasource_credentials(
tenant_id=current_tenant_id,
auth_id=req_data.credential_id,
@@ -397,7 +397,7 @@ class DatasourceAuthUpdateApi(Resource):
def post(self, req_data: DatasourceCredentialUpdatePayload, current_tenant_id: str, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasource_provider_service.update_datasource_credentials(
tenant_id=current_tenant_id,
auth_id=req_data.credential_id,
@@ -422,9 +422,12 @@ class DatasourceAuthListApi(Resource):
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, user: Account):
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasources = datasource_provider_service.get_all_datasource_credentials(
tenant_id=current_tenant_id, session=db.session(), user=user
tenant_id=current_tenant_id,
session=db.session(),
user=user,
credential_query=application_services().credential_queries,
)
return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200
@@ -442,9 +445,12 @@ class DatasourceHardCodeAuthListApi(Resource):
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, user: Account):
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasources = datasource_provider_service.get_hard_code_datasource_credentials(
tenant_id=current_tenant_id, session=db.session(), user=user
tenant_id=current_tenant_id,
session=db.session(),
user=user,
credential_query=application_services().credential_queries,
)
return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200
@@ -464,7 +470,7 @@ class DatasourceAuthOauthCustomClient(Resource):
@model_validate(DatasourceCustomClientPayload)
def post(self, req_data: DatasourceCustomClientPayload, current_tenant_id: str, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasource_provider_service.setup_oauth_custom_client_params(
tenant_id=current_tenant_id,
datasource_provider_id=datasource_provider_id,
@@ -482,7 +488,7 @@ class DatasourceAuthOauthCustomClient(Resource):
@with_current_tenant_id
def delete(self, current_tenant_id: str, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasource_provider_service.remove_oauth_custom_client_params(
tenant_id=current_tenant_id,
datasource_provider_id=datasource_provider_id,
@@ -503,7 +509,7 @@ class DatasourceAuthDefaultApi(Resource):
@model_validate(DatasourceDefaultPayload)
def post(self, req_data: DatasourceDefaultPayload, current_tenant_id: str, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasource_provider_service.set_default_datasource_provider(
tenant_id=current_tenant_id,
datasource_provider_id=datasource_provider_id,
@@ -525,7 +531,7 @@ class DatasourceUpdateProviderNameApi(Resource):
@model_validate(DatasourceUpdateNamePayload)
def post(self, req_data: DatasourceUpdateNamePayload, current_tenant_id: str, provider_id: str):
datasource_provider_id = DatasourceProviderID(provider_id)
datasource_provider_service = DatasourceProviderService()
datasource_provider_service = application_services().data_sources.providers
datasource_provider_service.update_datasource_provider_name(
tenant_id=current_tenant_id,
datasource_provider_id=datasource_provider_id,
api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py
+2
View File
@@ -9,6 +9,7 @@ from controllers.common.schema import register_schema_models
from controllers.console import console_ns
from controllers.console.datasets.wraps import get_rag_pipeline
from controllers.console.wraps import account_initialization_required, model_validate, setup_required, with_current_user
from extensions.ext_application_services import application_services
from extensions.ext_database import db
from libs.login import login_required
from models import Account
@@ -51,5 +52,6 @@ class DataSourceContentPreviewApi(Resource):
datasource_type=datasource_type,
is_published=True,
credential_id=req_data.credential_id,
datasource_providers=application_services().data_sources.providers,
)
return preview_content, 200
api/controllers/console/datasets/rag_pipeline/rag_pipeline.py
+3 -2
View File
@@ -35,10 +35,11 @@ from libs.helper import dump_response
from libs.login import login_required
from models.account import Account
from models.dataset import Pipeline
from services.dataset_service import DatasetService
from repositories.knowledge.dataset_read_repository import get_pipeline_dataset
from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity
from services.errors.account import NoPermissionError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.knowledge.dataset_service import DatasetService
from services.rag_pipeline.rag_pipeline import RagPipelineService
logger: logging.Logger = logging.getLogger(__name__)
@@ -232,7 +233,7 @@ class PublishCustomizedPipelineTemplateApi(Resource):
pipeline: Pipeline,
) -> tuple[str, int]:
session = db.session()
dataset = pipeline.retrieve_dataset(session=session)
dataset = get_pipeline_dataset(pipeline, session=session)
if dataset is None:
raise NotFound("Dataset not found")
api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py
+6 -3
View File
@@ -17,14 +17,17 @@ from controllers.console.wraps import (
with_current_user,
)
from extensions.ext_database import db
from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from fields.dataset_fields import (
DatasetDetailResponse,
)
from libs.helper import dump_response
from libs.login import login_required
from models import Account
from models.dataset import DatasetPermissionEnum
from services.dataset_service import DatasetPermissionService, DatasetService
from services.enterprise import rbac_service as enterprise_rbac_service
from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, RagPipelineDatasetCreateEntity
from services.knowledge.dataset_read_service import load_dataset_detail
from services.knowledge.dataset_service import DatasetPermissionService, DatasetService
from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService
@@ -136,4 +139,4 @@ class CreateEmptyRagPipelineDatasetApi(Resource):
),
session=session,
)
return dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session)), 201
return dump_response(DatasetDetailResponse, load_dataset_detail(dataset, session=session)), 201
api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py
+23 -4
View File
@@ -46,6 +46,7 @@ from core.app.apps.base_app_queue_manager import AppQueueManager
from core.app.apps.pipeline.pipeline_generator import PipelineGenerator
from core.app.entities.app_invoke_entities import InvokeFrom
from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity
from extensions.ext_application_services import application_services
from extensions.ext_database import db
from factories import variable_factory
from fields.base import ResponseModel
@@ -66,10 +67,10 @@ from models.dataset import Pipeline
from models.model import EndUser
from models.workflow import Workflow
from services.agent.retirement_service import WorkflowAgentRetirementService
from services.dataset_service import DatasetService
from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
from services.errors.llm import InvokeRateLimitError
from services.errors.rag_pipeline import RagPipelineResourceNotFoundError
from services.knowledge.dataset_service import DatasetService
from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService
from services.rag_pipeline.rag_pipeline import RagPipelineService
from services.rag_pipeline.rag_pipeline_manage_service import RagPipelineManageService
@@ -303,7 +304,13 @@ class RagPipelineDraftRunIterationNodeApi(Resource):
try:
response = PipelineGenerateService.generate_single_iteration(
pipeline=pipeline, user=current_user, node_id=node_id, args=args, session=db.session(), streaming=True
generator=application_services().knowledge.pipeline_generator,
pipeline=pipeline,
user=current_user,
node_id=node_id,
args=args,
session=db.session(),
streaming=True,
)
return helper.compact_generate_response(response)
@@ -338,7 +345,13 @@ class RagPipelineDraftRunLoopNodeApi(Resource):
try:
response = PipelineGenerateService.generate_single_loop(
pipeline=pipeline, user=current_user, node_id=node_id, args=args, session=db.session(), streaming=True
generator=application_services().knowledge.pipeline_generator,
pipeline=pipeline,
user=current_user,
node_id=node_id,
args=args,
session=db.session(),
streaming=True,
)
return helper.compact_generate_response(response)
@@ -374,6 +387,7 @@ class DraftRagPipelineRunApi(Resource):
try:
response = PipelineGenerateService.generate(
generator=application_services().knowledge.pipeline_generator,
session=session,
pipeline=pipeline,
user=current_user,
@@ -409,6 +423,7 @@ class PublishedRagPipelineRunApi(Resource):
try:
response = PipelineGenerateService.generate(
generator=application_services().knowledge.pipeline_generator,
session=session,
pipeline=pipeline,
user=current_user,
@@ -450,6 +465,7 @@ class RagPipelinePublishedDatasourceNodeRunApi(Resource):
datasource_type=req_data.datasource_type,
is_published=False,
credential_id=req_data.credential_id,
datasource_providers=application_services().data_sources.providers,
)
)
)
@@ -483,6 +499,7 @@ class RagPipelineDraftDatasourceNodeRunApi(Resource):
datasource_type=req_data.datasource_type,
is_published=False,
credential_id=req_data.credential_id,
datasource_providers=application_services().data_sources.providers,
)
)
)
@@ -1016,7 +1033,9 @@ class DatasourceListApi(Resource):
@account_initialization_required
@with_current_tenant_id
def get(self, current_tenant_id: str):
providers = RagPipelineManageService.list_rag_pipeline_datasources(current_tenant_id)
providers = RagPipelineManageService.list_rag_pipeline_datasources(
current_tenant_id, datasource_providers=application_services().data_sources.providers
)
return RagPipelineDatasourceListResponse.model_validate(providers, from_attributes=True).model_dump(mode="json")
api/controllers/console/datasets/website.py
+6 -3
View File
@@ -7,8 +7,9 @@ from controllers.common.schema import query_params_from_model, register_response
from controllers.console import console_ns
from controllers.console.datasets.error import WebsiteCrawlError
from controllers.console.wraps import account_initialization_required, model_validate, setup_required
from extensions.ext_application_services import application_services
from libs.login import login_required
from services.website_service import WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, WebsiteService
from services.data_source.website_service import WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, WebsiteService
class WebsiteCrawlPayload(BaseModel):
@@ -49,7 +50,7 @@ class WebsiteCrawlApi(Resource):
# Crawl URL using typed request
try:
result = WebsiteService.crawl_url(api_request)
result = WebsiteService(providers=application_services().data_sources.providers).crawl_url(api_request)
except Exception as e:
raise WebsiteCrawlError(str(e))
return result, 200
@@ -77,7 +78,9 @@ class WebsiteCrawlStatusApi(Resource):
# Get crawl status using typed request
try:
result = WebsiteService.get_crawl_status_typed(api_request)
result = WebsiteService(providers=application_services().data_sources.providers).get_crawl_status_typed(
api_request
)
except Exception as e:
raise WebsiteCrawlError(str(e))
return result, 200
api/controllers/console/flask_admission.py
+1 -1
View File
@@ -71,7 +71,7 @@ def console_account_admission[T, **P, R](
) -> Callable[Concatenate[T, P], R | Response]:
admitted_view = cloud_edition_billing_resource_check(billing_resource)(view) if billing_resource else view
@wraps(view, updated=())
@wraps(view)
def inject_request_context(self: T, /, *args: P.args, **kwargs: P.kwargs) -> R:
account_with_tenant = current_account_with_tenant()
account = account_with_tenant.account
api/controllers/console/workspace/model_providers.py
+1 -1
View File
@@ -33,7 +33,7 @@ from services.entities.model_provider_entities import (
ModelProviderSummaryResponse,
ProviderResponse,
)
from services.model_provider_service import ModelProviderService
from services.model_provider.service import ModelProviderService
from services.workspace_service import WorkspaceService
api/controllers/console/workspace/models.py
+3 -1
View File
@@ -25,6 +25,7 @@ from controllers.console.wraps import (
)
from core.entities.model_entities import DefaultModelSetting
from core.entities.provider_entities import CredentialConfiguration
from extensions.ext_application_services import application_services
from extensions.ext_database import db
from fields.base import ResponseModel
from graphon.model_runtime.entities.model_entities import ModelType, ParameterRule
@@ -38,7 +39,7 @@ from services.entities.model_provider_entities import (
ProviderWithModelsResponse,
)
from services.model_load_balancing_service import ModelLoadBalancingService
from services.model_provider_service import ModelProviderService
from services.model_provider.service import ModelProviderService
logger = logging.getLogger(__name__)
@@ -376,6 +377,7 @@ class ModelProviderModelCredentialApi(Resource):
tenant_id=tenant_id,
provider=provider,
user=user,
credential_query=application_services().credential_queries,
)
else:
available_credentials = model_provider_service.get_provider_model_available_credentials(
api/controllers/console/workspace/tool_providers.py
+3 -2
View File
@@ -62,6 +62,7 @@ from core.tools.entities.tool_entities import (
WorkflowToolParameterConfiguration,
)
from enums import DeploymentEdition
from extensions.ext_application_services import application_services
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.helper import alphanumeric, dump_response, uuid_value
@@ -664,9 +665,9 @@ class ToolBuiltinProviderGetCredentialsApi(Resource):
BuiltinToolManageService.get_builtin_tool_provider_credentials(
tenant_id=tenant_id,
provider_name=provider,
session=db.session(),
user=user,
include_credential_ids=query.include_credential_ids or None,
credential_query=application_services().credential_queries,
),
)
@@ -1363,9 +1364,9 @@ class ToolBuiltinProviderGetCredentialInfoApi(Resource):
BuiltinToolManageService.get_builtin_tool_provider_credential_info(
tenant_id=tenant_id,
provider=provider,
session=db.session(),
user=user,
include_credential_ids=query.include_credential_ids or None,
credential_query=application_services().credential_queries,
),
)
api/controllers/console/workspace/trigger_providers.py
+2
View File
@@ -22,6 +22,7 @@ from core.trigger.entities.api_entities import (
)
from core.trigger.entities.entities import RequestLog, SubscriptionBuilderUpdater
from core.trigger.trigger_manager import TriggerManager
from extensions.ext_application_services import application_services
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.helper import dump_response
@@ -213,6 +214,7 @@ class TriggerSubscriptionListApi(Resource):
tenant_id=tenant_id,
provider_id=TriggerProviderID(provider),
user=user,
credential_query=application_services().credential_queries,
),
)
except ValueError as e:
api/controllers/service_api/dataset/dataset.py
+10 -18
View File
@@ -37,11 +37,6 @@ from core.plugin.impl.model_runtime_factory import create_plugin_provider_manage
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from fields.dataset_fields import (
DatasetDetailPrefetch,
build_dataset_detail_prefetch,
dataset_detail_response_source,
)
from fields.dataset_fields import DatasetDetailResponse as BaseDatasetDetailResponse
from graphon.model_runtime.entities.model_entities import ModelType
from libs.helper import dump_response
@@ -51,9 +46,11 @@ from models.account import Account
from models.dataset import DatasetPermissionEnum
from models.enums import TagType
from models.provider_ids import ModelProviderID
from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
from services.enterprise import rbac_service as enterprise_rbac_service
from services.entities.knowledge_entities.knowledge_entities import (
from services.knowledge.dataset_read_service import load_dataset_detail, load_dataset_details
from services.knowledge.dataset_service import DatasetPermissionService, DatasetService, DocumentService
from services.knowledge.entities.datasets import DatasetDetailRecord
from services.knowledge.entities.knowledge_entities import (
ExternalRetrievalModel,
KnowledgeProvider,
RetrievalModel,
@@ -104,12 +101,8 @@ _SERVICE_DATASET_DETAIL_EXCLUDE = {"permission_keys"}
_SERVICE_DATASET_LIST_EXCLUDE = {"data": {"__all__": _SERVICE_DATASET_DETAIL_EXCLUDE}}
def _dump_service_dataset_detail(
dataset: Any, *, session: Session, prefetch: DatasetDetailPrefetch | None = None
) -> dict[str, Any]:
return DatasetDetailResponse.model_validate(
dataset_detail_response_source(dataset, session=session, prefetch=prefetch), from_attributes=True
).model_dump(
def _dump_service_dataset_detail(detail: DatasetDetailRecord) -> dict[str, Any]:
return DatasetDetailResponse.model_validate(detail).model_dump(
mode="json",
exclude=_SERVICE_DATASET_DETAIL_EXCLUDE,
)
@@ -470,8 +463,7 @@ class DatasetListApi(DatasetApiResource):
for embedding_model in embedding_models:
model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
prefetch = build_dataset_detail_prefetch(datasets, session=session)
data = [_dump_service_dataset_detail(dataset, session=session, prefetch=prefetch) for dataset in datasets]
data = [_dump_service_dataset_detail(detail) for detail in load_dataset_details(datasets, session=session)]
for item in data:
if item["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY and item["embedding_model_provider"]:
item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"]))
@@ -571,7 +563,7 @@ class DatasetListApi(DatasetApiResource):
enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False),
)
return _dump_service_dataset_detail(dataset, session=session), 200
return _dump_service_dataset_detail(load_dataset_detail(dataset, session=session)), 200
@service_api_ns.route("/datasets/<uuid:dataset_id>")
@@ -617,7 +609,7 @@ class DatasetApi(DatasetApiResource):
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
data = _dump_service_dataset_detail(dataset, session=session)
data = _dump_service_dataset_detail(load_dataset_detail(dataset, session=session))
# check embedding setting
assert isinstance(current_user, Account)
cid = current_user.current_tenant_id
@@ -734,7 +726,7 @@ class DatasetApi(DatasetApiResource):
if dataset is None:
raise NotFound("Dataset not found.")
result_data = _dump_service_dataset_detail(dataset, session=session)
result_data = _dump_service_dataset_detail(load_dataset_detail(dataset, session=session))
assert isinstance(current_user, Account)
tenant_id = current_user.current_tenant_id
api/controllers/service_api/dataset/document.py
+39 -28
View File
@@ -69,25 +69,36 @@ from fields.document_fields import (
DocumentMetadataResponse,
DocumentResponse,
DocumentStatusListResponse,
document_response,
document_responses,
normalize_enum,
)
from libs.helper import dump_response
from libs.login import current_user
from libs.pagination import clamp_pagination, paginate_query
from models.dataset import Dataset, Document
from services.dataset_service import DatasetService, DocumentService
from services.entities.knowledge_entities.knowledge_entities import (
from repositories.knowledge.dataset_read_repository import (
get_dataset_creator,
get_document_hit_count,
get_document_process_rule,
get_document_segment_count,
get_latest_dataset_process_rule,
)
from services.feature_service import FeatureService
from services.file_service import FileService
from services.knowledge.dataset_read_service import (
get_document_metadata_details,
get_document_source_detail,
load_document_detail,
load_document_details,
)
from services.knowledge.dataset_service import DatasetService, DocumentService
from services.knowledge.entities.knowledge_entities import (
DocForm,
IndexingTechnique,
KnowledgeConfig,
ProcessRule,
RetrievalModel,
)
from services.feature_service import FeatureService
from services.file_service import FileService
from services.summary_index_service import SummaryIndexService
from services.knowledge.summaries.adapters import SummaryIndexAdapter
class DocumentTextCreatePayload(BaseModel):
@@ -304,7 +315,7 @@ class DocumentAndBatchResponse(ResponseModel):
def _document_and_batch_response(document: Document, batch: str, *, session: Session) -> dict[str, Any]:
return dump_response(
DocumentAndBatchResponse,
{"document": document_response(document, session=session), "batch": batch},
{"document": load_document_detail(document, session=session), "batch": batch},
)
@@ -445,7 +456,7 @@ def _create_document_by_text(session: Session, tenant_id: str, dataset_id: UUID)
dataset=dataset,
knowledge_config=knowledge_config,
account=current_user,
dataset_process_rule=dataset.get_latest_process_rule(session=session)
dataset_process_rule=get_latest_dataset_process_rule(dataset, session=session)
if "process_rule" not in args
else None,
created_from="api",
@@ -509,7 +520,7 @@ def _update_document_by_text(
dataset=dataset,
knowledge_config=knowledge_config,
account=current_user,
dataset_process_rule=dataset.get_latest_process_rule(session=session)
dataset_process_rule=get_latest_dataset_process_rule(dataset, session=session)
if "process_rule" not in args
else None,
created_from="api",
@@ -822,7 +833,9 @@ class DocumentAddByFileApi(DatasetApiResource):
knowledge_config = KnowledgeConfig.model_validate(args)
DocumentService.document_create_args_validate(knowledge_config)
dataset_process_rule = dataset.get_latest_process_rule(session=session) if "process_rule" not in args else None
dataset_process_rule = (
get_latest_dataset_process_rule(dataset, session=session) if "process_rule" not in args else None
)
if not knowledge_config.original_document_id and not dataset_process_rule and not knowledge_config.process_rule:
raise ValueError("process_rule is required.")
@@ -830,7 +843,7 @@ class DocumentAddByFileApi(DatasetApiResource):
documents, batch = DocumentService.save_document_with_dataset_id(
dataset=dataset,
knowledge_config=knowledge_config,
account=dataset.get_created_by_account(session=session),
account=get_dataset_creator(dataset, session=session),
dataset_process_rule=dataset_process_rule,
created_from="api",
session=session,
@@ -909,8 +922,8 @@ def _update_document_by_file(
documents, _ = DocumentService.save_document_with_dataset_id(
dataset=dataset,
knowledge_config=knowledge_config,
account=dataset.get_created_by_account(session=session),
dataset_process_rule=dataset.get_latest_process_rule(session=session)
account=get_dataset_creator(dataset, session=session),
dataset_process_rule=get_latest_dataset_process_rule(dataset, session=session)
if "process_rule" not in args
else None,
created_from="api",
@@ -1046,9 +1059,7 @@ class DocumentListApi(DatasetApiResource):
)
response = {
"data": document_responses(documents, session=session),
# The result object already knows: it was built from the page the query
# ran with, while the requested values are only ever a request.
"data": load_document_details(documents, session=session),
"has_more": paginated_documents.has_next,
"limit": paginated_documents.per_page,
"total": paginated_documents.total,
@@ -1303,7 +1314,7 @@ class DocumentApi(DatasetApiResource):
summary_index_status = None
has_summary_index = dataset.summary_index_setting and dataset.summary_index_setting.get("enable") is True
if has_summary_index and document.need_summary is True:
summary_index_status = SummaryIndexService.get_document_summary_index_status(
summary_index_status = SummaryIndexAdapter.get_document_summary_index_status(
document_id=document_id_str,
dataset_id=dataset_id_str,
tenant_id=tenant_id,
@@ -1315,15 +1326,15 @@ class DocumentApi(DatasetApiResource):
response = {
"id": document.id,
"doc_type": document.doc_type,
"doc_metadata": document.get_doc_metadata_details(session=session),
"doc_metadata": get_document_metadata_details(document, session=session),
}
elif metadata == "without":
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session)
response_exclude = {"doc_type", "doc_metadata"}
document_process_rule = document.get_dataset_process_rule(session=session)
document_process_rule = get_document_process_rule(document, session=session)
document_process_rules: Mapping[str, Any] = document_process_rule.to_dict() if document_process_rule else {}
data_source_info = document.get_data_source_detail_dict(session=session)
segment_count = document.get_segment_count(session=session)
data_source_info = get_document_source_detail(document, session=session)
segment_count = get_document_segment_count(document, session=session)
response = {
"id": document.id,
"position": document.position,
@@ -1348,7 +1359,7 @@ class DocumentApi(DatasetApiResource):
"archived": document.archived,
"segment_count": segment_count,
"average_segment_length": (document.word_count or 0) // segment_count if segment_count else 0,
"hit_count": document.get_hit_count(session=session),
"hit_count": get_document_hit_count(document, session=session),
"display_status": document.display_status,
"doc_form": document.doc_form,
"doc_language": document.doc_language,
@@ -1357,10 +1368,10 @@ class DocumentApi(DatasetApiResource):
}
else:
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session)
document_process_rule = document.get_dataset_process_rule(session=session)
document_process_rule = get_document_process_rule(document, session=session)
document_process_rules = document_process_rule.to_dict() if document_process_rule else {}
data_source_info = document.get_data_source_detail_dict(session=session)
segment_count = document.get_segment_count(session=session)
data_source_info = get_document_source_detail(document, session=session)
segment_count = get_document_segment_count(document, session=session)
response = {
"id": document.id,
"position": document.position,
@@ -1384,10 +1395,10 @@ class DocumentApi(DatasetApiResource):
"disabled_by": document.disabled_by,
"archived": document.archived,
"doc_type": document.doc_type,
"doc_metadata": document.get_doc_metadata_details(session=session),
"doc_metadata": get_document_metadata_details(document, session=session),
"segment_count": segment_count,
"average_segment_length": (document.word_count or 0) // segment_count if segment_count else 0,
"hit_count": document.get_hit_count(session=session),
"hit_count": get_document_hit_count(document, session=session),
"display_status": document.display_status,
"doc_form": document.doc_form,
"doc_language": document.doc_language,
api/controllers/service_api/dataset/metadata.py
+54 -23
View File
@@ -1,4 +1,4 @@
from typing import Literal, cast
from typing import Literal
from uuid import UUID
from flask_login import current_user
@@ -11,6 +11,7 @@ from controllers.common.session import with_session
from controllers.console.wraps import model_validate
from controllers.service_api import service_api_ns
from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_rate_limit_check
from extensions.ext_application_services import application_services
from fields.dataset_fields import (
DatasetMetadataActionResponse,
DatasetMetadataBuiltInFieldsResponse,
@@ -18,16 +19,15 @@ from fields.dataset_fields import (
DatasetMetadataResponse,
)
from libs.helper import dump_response
from models import Account
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.knowledge_entities import (
from services.errors.metadata import MetadataResourceNotFoundError
from services.knowledge.dataset_service import DatasetService
from services.knowledge.entities.knowledge_entities import (
DocumentMetadataOperation,
MetadataArgs,
MetadataDetail,
MetadataOperationData,
)
from services.errors.metadata import MetadataResourceNotFoundError
from services.metadata_service import MetadataService
from services.knowledge.resource_scope import DatasetRef
BUILT_IN_METADATA_ACTION_PARAM = {
"description": "`enable` to activate built-in metadata fields, `disable` to deactivate them.",
@@ -87,12 +87,18 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource):
"""Create metadata for a dataset."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.create_metadata(dataset_id_str, metadata_args, session=session)
session.commit()
try:
metadata = application_services().knowledge.metadata.create_metadata(
DatasetRef(str(tenant_id), dataset_id_str), metadata_args, actor_id=current_user.id
)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return dump_response(DatasetMetadataResponse, metadata), 201
@service_api_ns.doc(
@@ -123,10 +129,16 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource):
def get(self, session: Session, tenant_id, dataset_id: UUID):
"""Get all metadata for a dataset."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session)
if dataset is None:
raise NotFound("Dataset not found.")
metadata = MetadataService.get_dataset_metadatas(dataset, session)
session.commit()
try:
metadata = application_services().knowledge.metadata.get_dataset_metadatas(
DatasetRef(str(tenant_id), dataset_id_str)
)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return dump_response(DatasetMetadataListResponse, metadata), 200
@@ -168,9 +180,13 @@ class DatasetMetadataServiceApi(DatasetApiResource):
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.update_metadata_name(
dataset, metadata_id_str, payload.name, cast(Account, current_user), session=session
)
session.commit()
try:
metadata = application_services().knowledge.metadata.update_metadata_name(
DatasetRef(str(tenant_id), dataset_id_str), metadata_id_str, payload.name, actor_id=current_user.id
)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return dump_response(DatasetMetadataResponse, metadata), 200
@service_api_ns.doc(
@@ -206,7 +222,13 @@ class DatasetMetadataServiceApi(DatasetApiResource):
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
MetadataService.delete_metadata(dataset, metadata_id_str, session)
session.commit()
try:
application_services().knowledge.metadata.delete_metadata(
DatasetRef(str(tenant_id), dataset_id_str), metadata_id_str
)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return "", 204
@@ -239,7 +261,7 @@ class DatasetMetadataBuiltInFieldServiceApi(DatasetApiResource):
)
def get(self, tenant_id, dataset_id: UUID):
"""Get all built-in metadata fields."""
built_in_fields = MetadataService.get_built_in_fields()
built_in_fields = application_services().knowledge.metadata.get_built_in_fields()
return dump_response(DatasetMetadataBuiltInFieldsResponse, {"fields": built_in_fields}), 200
@@ -271,16 +293,24 @@ class DatasetMetadataBuiltInFieldActionServiceApi(DatasetApiResource):
def post(self, session: Session, tenant_id, dataset_id: UUID, action: Literal["enable", "disable"]):
"""Enable or disable built-in metadata field."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, session)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
match action:
case "enable":
MetadataService.enable_built_in_field(dataset, session)
case "disable":
MetadataService.disable_built_in_field(dataset, session)
session.commit()
try:
match action:
case "enable":
application_services().knowledge.metadata.enable_built_in_field(
DatasetRef(str(tenant_id), dataset_id_str)
)
case "disable":
application_services().knowledge.metadata.disable_built_in_field(
DatasetRef(str(tenant_id), dataset_id_str)
)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200
@@ -324,9 +354,10 @@ class DocumentMetadataEditServiceApi(DatasetApiResource):
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, session)
session.commit()
try:
MetadataService.update_documents_metadata(
dataset, metadata_args, cast(Account, current_user), session=session
application_services().knowledge.metadata.update_documents_metadata(
DatasetRef(str(tenant_id), str(dataset_id)), metadata_args, actor_id=current_user.id
)
except MetadataResourceNotFoundError as exc:
raise NotFound(str(exc)) from exc
api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py
+8 -1
View File
@@ -32,6 +32,7 @@ from controllers.service_api.wraps import DatasetApiResource
from core.app.apps.pipeline.pipeline_generator import PipelineGenerator
from core.app.entities.app_invoke_entities import InvokeFrom
from core.entities.knowledge_entities import PipelineDataset, PipelineDocument
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from libs import helper
from libs.helper import dump_response
@@ -183,7 +184,11 @@ class DatasourcePluginsApi(DatasetApiResource):
rag_pipeline_service = RagPipelineService(db.session())
datasource_plugins: list[dict[Any, Any]] = rag_pipeline_service.get_datasource_plugins(
tenant_id=tenant_id, dataset_id=dataset_id_str, is_published=query.is_published
tenant_id=tenant_id,
dataset_id=dataset_id_str,
is_published=query.is_published,
credential_query=application_services().credential_queries,
datasource_providers=application_services().data_sources.providers,
)
return dump_response(DatasourcePluginListResponse, datasource_plugins), 200
@@ -247,6 +252,7 @@ class DatasourceNodeRunApi(DatasetApiResource):
datasource_type=datasource_node_run_api_entity.datasource_type,
is_published=datasource_node_run_api_entity.is_published,
credential_id=datasource_node_run_api_entity.credential_id,
datasource_providers=application_services().data_sources.providers,
)
)
)
@@ -310,6 +316,7 @@ class PipelineRunApi(DatasetApiResource):
pipeline = rag_pipeline_service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset_id_str)
try:
response: dict[Any, Any] | Generator[str, Any, None] = PipelineGenerateService.generate(
generator=application_services().knowledge.pipeline_generator,
session=session,
pipeline=pipeline,
user=current_user,
api/controllers/service_api/dataset/segment.py
+60 -31
View File
@@ -27,26 +27,26 @@ from controllers.service_api.wraps import (
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.model_manager import ModelManager
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from fields.segment_fields import (
ChildChunkDetailResponse,
ChildChunkListResponse,
SegmentDetailResponse,
SegmentResponse,
segment_response_with_summary,
segment_responses_with_summaries,
)
from graphon.model_runtime.entities.model_entities import ModelType
from libs.helper import dump_response
from libs.login import current_account_with_tenant
from models.dataset import Dataset, Document, DocumentSegment
from services.dataset_ref_service import DatasetRefService, SegmentRef
from services.dataset_service import DatasetService, DocumentService, SegmentService
from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs
from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError
from services.errors.chunk import ChildChunkDeleteIndexError as ChildChunkDeleteIndexServiceError
from services.errors.chunk import ChildChunkIndexingError as ChildChunkIndexingServiceError
from services.summary_index_service import SummaryIndexService
from services.knowledge.dataset_read_service import load_segment_detail, load_segment_details
from services.knowledge.dataset_service import DatasetService, DocumentService, SegmentService
from services.knowledge.entities.segments import SegmentUpdateArgs
from services.knowledge.resource_scope import DatasetRef, SegmentRef
from services.knowledge.summaries.adapters import SummaryIndexAdapter
class SegmentCreateItemPayload(BaseModel):
@@ -133,12 +133,11 @@ register_response_schema_models(
def _get_segment_for_document(
session: Session, dataset: Dataset, document: Document, segment_id: str
) -> tuple[SegmentRef, DocumentSegment]:
dataset_ref = DatasetRefService.create_dataset_ref(dataset)
document_ref = DatasetRefService.create_document_ref(dataset_ref, document)
if document_ref is None:
if document.tenant_id != dataset.tenant_id or document.dataset_id != dataset.id:
raise NotFound("Document not found.")
segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id)
dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id)
segment_ref = dataset_ref.document(document.id).segment(segment_id)
segment = SegmentService.get_segment_by_ref(segment_ref, session=session)
if not segment:
raise NotFound("Segment not found.")
@@ -230,17 +229,24 @@ class SegmentApi(DatasetApiResource):
for args_item in segment_items:
SegmentService.segment_create_args_validate(args_item, document)
segments = cast(
list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset, session)
list[DocumentSegment],
SegmentService.multi_create_segment(
segment_items,
document,
dataset,
session,
mutations=application_services().knowledge.segments.mutations,
),
)
segment_ids = [segment.id for segment in segments]
summaries: dict[str, str | None] = {}
if segment_ids:
summary_records = SummaryIndexService.get_segments_summaries(
summary_records = SummaryIndexAdapter.get_segments_summaries(
segment_ids=segment_ids, dataset_id=dataset_id_str, session=session
)
summaries = {chunk_id: record.summary_content for chunk_id, record in summary_records.items()}
response = {
"data": segment_responses_with_summaries(segments, summaries, session=session),
"data": load_segment_details(segments, summaries, session=session),
"doc_form": document.doc_form,
}
return dump_response(SegmentCreateListResponse, response), 200
@@ -322,13 +328,13 @@ class SegmentApi(DatasetApiResource):
segment_ids = [segment.id for segment in segments]
summaries: dict[str, str | None] = {}
if segment_ids:
summary_records = SummaryIndexService.get_segments_summaries(
summary_records = SummaryIndexAdapter.get_segments_summaries(
segment_ids=segment_ids, dataset_id=dataset_id_str, session=session
)
summaries = {chunk_id: record.summary_content for chunk_id, record in summary_records.items()}
response = {
"data": segment_responses_with_summaries(segments, summaries, session=session),
"data": load_segment_details(segments, summaries, session=session),
"doc_form": document.doc_form,
"total": total,
"has_more": page * limit < total,
@@ -380,7 +386,9 @@ class DatasetSegmentApi(DatasetApiResource):
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
SegmentService.delete_segment(segment, document, dataset, session)
SegmentService.delete_segment(
segment, document, dataset, session, mutations=application_services().knowledge.segments.mutations
)
return "", 204
@service_api_ns.doc(
@@ -417,7 +425,7 @@ class DatasetSegmentApi(DatasetApiResource):
document_id: UUID,
segment_id: UUID,
):
_, current_tenant_id = current_account_with_tenant()
current_account, current_tenant_id = current_account_with_tenant()
dataset_id_str = str(dataset_id)
# check dataset
dataset = session.scalar(
@@ -451,14 +459,20 @@ class DatasetSegmentApi(DatasetApiResource):
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, session)
summary = SummaryIndexService.get_segment_summary(
updated_segment = SegmentService.update_segment(
payload.segment,
segment,
document,
dataset,
session,
actor_id=current_account.id,
mutations=application_services().knowledge.segments.mutations,
)
summary = SummaryIndexAdapter.get_segment_summary(
segment_id=updated_segment.id, dataset_id=dataset_id_str, session=session
)
response = {
"data": segment_response_with_summary(
updated_segment, summary.summary_content if summary else None, session=session
),
"data": load_segment_detail(updated_segment, summary.summary_content if summary else None, session=session),
"doc_form": document.doc_form,
}
return dump_response(SegmentDetailResponse, response), 200
@@ -510,13 +524,11 @@ class DatasetSegmentApi(DatasetApiResource):
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
summary = SummaryIndexService.get_segment_summary(
summary = SummaryIndexAdapter.get_segment_summary(
segment_id=segment.id, dataset_id=dataset_id_str, session=session
)
response = {
"data": segment_response_with_summary(
segment, summary.summary_content if summary else None, session=session
),
"data": load_segment_detail(segment, summary.summary_content if summary else None, session=session),
"doc_form": document.doc_form,
}
return dump_response(SegmentDetailResponse, response), 200
@@ -567,7 +579,7 @@ class ChildChunkApi(DatasetApiResource):
document_id: UUID,
segment_id: UUID,
):
_, current_tenant_id = current_account_with_tenant()
current_account, current_tenant_id = current_account_with_tenant()
"""Create child chunk."""
dataset_id_str = str(dataset_id)
# check dataset
@@ -604,7 +616,15 @@ class ChildChunkApi(DatasetApiResource):
raise ProviderNotInitializeError(ex.description)
try:
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, session)
child_chunk = SegmentService.create_child_chunk(
payload.content,
segment,
document,
dataset,
session,
actor_id=current_account.id,
mutations=application_services().knowledge.segments.mutations,
)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
@@ -744,7 +764,9 @@ class DatasetChildChunkApi(DatasetApiResource):
raise NotFound("Child chunk not found.")
try:
SegmentService.delete_child_chunk(child_chunk, dataset, session)
SegmentService.delete_child_chunk(
child_chunk, dataset, session, mutations=application_services().knowledge.segments.mutations
)
except ChildChunkDeleteIndexServiceError as e:
raise ChildChunkDeleteIndexError(str(e))
@@ -790,7 +812,7 @@ class DatasetChildChunkApi(DatasetApiResource):
segment_id: UUID,
child_chunk_id: UUID,
):
current_account_with_tenant()
current_account, _ = current_account_with_tenant()
"""Update child chunk."""
dataset_id_str = str(dataset_id)
# check dataset
@@ -817,7 +839,14 @@ class DatasetChildChunkApi(DatasetApiResource):
try:
child_chunk = SegmentService.update_child_chunk(
payload.content, child_chunk, segment, document, dataset, session
payload.content,
child_chunk,
segment,
document,
dataset,
session,
actor_id=current_account.id,
mutations=application_services().knowledge.segments.mutations,
)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
api/controllers/service_api/workspace/models.py
+1 -1
View File
@@ -7,7 +7,7 @@ from controllers.service_api.wraps import validate_dataset_token
from fields.base import ResponseModel
from graphon.model_runtime.utils.encoders import jsonable_encoder
from services.entities.model_provider_entities import ProviderWithModelsResponse
from services.model_provider_service import ModelProviderService
from services.model_provider.service import ModelProviderService
MODEL_TYPE_PARAM = {
"description": "Type of model to retrieve.",
api/core/app/app_config/easy_ui_based_app/dataset/manager.py
+1 -1
View File
@@ -12,7 +12,7 @@ from core.app.app_config.entities import (
from core.entities.agent_entities import PlanningStrategy
from core.rag.data_post_processor.data_post_processor import RerankingModelDict, WeightsDict
from models.model import AppMode, AppModelConfigDict
from services.dataset_service import DatasetService
from services.knowledge.dataset_service import DatasetService
class DatasetConfigManager:
api/core/app/apps/pipeline/pipeline_generator.py
+12 -8
View File
@@ -22,7 +22,7 @@ from core.app.apps.draft_variable_saver import DraftVariableSaverFactory
from core.app.apps.exc import GenerateTaskStoppedError
from core.app.apps.pipeline.pipeline_config_manager import PipelineConfigManager
from core.app.apps.pipeline.pipeline_queue_manager import PipelineQueueManager
from core.app.apps.pipeline.pipeline_runner import PipelineRunner
from core.app.apps.pipeline.pipeline_runner import PipelineDocumentStore, PipelineRunner
from core.app.apps.workflow.generate_response_converter import WorkflowAppGenerateResponseConverter
from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline
from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity
@@ -52,7 +52,7 @@ from models import Account, EndUser, Workflow, WorkflowNodeExecutionTriggeredFro
from models.dataset import Document, DocumentPipelineExecutionLog, Pipeline
from models.enums import WorkflowRunTriggeredFrom
from models.model import AppMode
from services.datasource_provider_service import DatasourceProviderService
from services.data_source.provider_service import DatasourceProviderService
from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy
from services.workflow_draft_variable_service import DraftVarLoader, WorkflowDraftVariableService
@@ -60,6 +60,10 @@ logger = logging.getLogger(__name__)
class PipelineGenerator(BaseAppGenerator):
def __init__(self, *, documents: PipelineDocumentStore, datasource_providers: DatasourceProviderService) -> None:
self._documents = documents
self._datasource_providers = datasource_providers
@overload
def generate(
self,
@@ -124,7 +128,7 @@ class PipelineGenerator(BaseAppGenerator):
) -> Mapping[str, Any] | Generator[Mapping | str, None, None] | None:
# Add null check for dataset
dataset = pipeline.retrieve_dataset(session)
dataset = self._documents.get_pipeline_dataset(pipeline, session=session)
if not dataset:
raise ValueError("Pipeline dataset is required")
inputs: Mapping[str, Any] = args["inputs"]
@@ -140,8 +144,8 @@ class PipelineGenerator(BaseAppGenerator):
)
documents: list[Document] = []
if invoke_from == InvokeFrom.PUBLISHED_PIPELINE and not is_retry and not args.get("original_document_id"):
from services.dataset_service import DocumentService
from services.feature_service import FeatureService
from services.knowledge.dataset_service import DocumentService
features = FeatureService.get_features(pipeline.tenant_id)
DocumentService.check_document_creation_limits(len(datasource_info_list), features)
@@ -407,7 +411,7 @@ class PipelineGenerator(BaseAppGenerator):
pipeline=pipeline, workflow=workflow, start_node_id=args.get("start_node_id", "shared")
)
dataset = pipeline.retrieve_dataset(session)
dataset = self._documents.get_pipeline_dataset(pipeline, session=session)
if not dataset:
raise ValueError("Pipeline dataset is required")
@@ -504,7 +508,7 @@ class PipelineGenerator(BaseAppGenerator):
if args.get("inputs") is None:
raise ValueError("inputs is required")
dataset = pipeline.retrieve_dataset(session)
dataset = self._documents.get_pipeline_dataset(pipeline, session=session)
if not dataset:
raise ValueError("Pipeline dataset is required")
@@ -635,6 +639,7 @@ class PipelineGenerator(BaseAppGenerator):
system_user_id=system_user_id,
workflow_execution_repository=workflow_execution_repository,
workflow_node_execution_repository=workflow_node_execution_repository,
documents=self._documents,
)
runner.run()
@@ -779,8 +784,7 @@ class PipelineGenerator(BaseAppGenerator):
tenant_id=pipeline.tenant_id,
datasource_type=DatasourceProviderType(datasource_type),
)
datasource_provider_service = DatasourceProviderService()
credentials = datasource_provider_service.get_datasource_credentials(
credentials = self._datasource_providers.get_datasource_credentials(
tenant_id=pipeline.tenant_id,
provider=datasource_node_data.get("provider_name"),
plugin_id=datasource_node_data.get("plugin_id"),
api/core/app/apps/pipeline/pipeline_runner.py
+48 -28
View File
@@ -1,6 +1,6 @@
import logging
import time
from typing import cast
from typing import Protocol, cast
from sqlalchemy import select
from sqlalchemy.orm import Session
@@ -28,14 +28,23 @@ from graphon.graph_events import GraphEngineEvent, GraphRunFailedEvent
from graphon.runtime import GraphRuntimeState, VariablePool
from graphon.variable_loader import VariableLoader
from graphon.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput
from models.dataset import Pipeline
from models.dataset import Dataset, Pipeline
from models.model import EndUser
from models.workflow import Workflow
from services.dataset_ref_service import DatasetRefService, DocumentRef
logger = logging.getLogger(__name__)
class PipelineDocumentStore(Protocol):
"""Knowledge state required by a pipeline run, with explicit session ownership."""
def get_pipeline_dataset(self, pipeline: Pipeline, *, session: Session) -> Dataset | None: ...
def exists(self, *, workspace_id: str, dataset_id: str, document_id: str) -> bool: ...
def mark_failed(self, *, workspace_id: str, dataset_id: str, document_id: str, error: str) -> None: ...
class PipelineRunner(WorkflowBasedAppRunner):
"""
Pipeline Application Runner
@@ -50,6 +59,7 @@ class PipelineRunner(WorkflowBasedAppRunner):
system_user_id: str,
workflow_execution_repository: WorkflowExecutionRepository,
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
documents: PipelineDocumentStore,
workflow_thread_pool_id: str | None = None,
) -> None:
"""
@@ -68,6 +78,7 @@ class PipelineRunner(WorkflowBasedAppRunner):
self._sys_user_id = system_user_id
self._workflow_execution_repository = workflow_execution_repository
self._workflow_node_execution_repository = workflow_node_execution_repository
self._documents = documents
def _get_app_id(self) -> str:
return self.application_generate_entity.app_config.app_id
@@ -98,7 +109,7 @@ class PipelineRunner(WorkflowBasedAppRunner):
if not pipeline or pipeline.tenant_id != app_config.tenant_id:
raise ValueError("Pipeline not found")
dataset = pipeline.retrieve_dataset(session)
dataset = self._documents.get_pipeline_dataset(pipeline, session=session)
if (
not dataset
or dataset.tenant_id != pipeline.tenant_id
@@ -108,23 +119,20 @@ class PipelineRunner(WorkflowBasedAppRunner):
document_id = self.application_generate_entity.document_id
original_document_id = self.application_generate_entity.original_document_id
dataset_ref = DatasetRefService.create_dataset_ref(dataset)
document_ref = (
DatasetRefService.create_document_ref_from_id(
dataset_ref,
document_id,
)
if document_id
else None
)
if document_ref and DatasetRefService.get_document_by_ref(document_ref, session=session) is None:
dataset_workspace_id = dataset.tenant_id
dataset_id = dataset.id
if document_id and not self._documents.exists(
workspace_id=dataset_workspace_id,
dataset_id=dataset_id,
document_id=document_id,
):
raise ValueError("Pipeline document not found")
if original_document_id and original_document_id != document_id:
original_document_ref = DatasetRefService.create_document_ref_from_id(
dataset_ref,
original_document_id,
)
if DatasetRefService.get_document_by_ref(original_document_ref, session=session) is None:
if not self._documents.exists(
workspace_id=dataset_workspace_id,
dataset_id=dataset_id,
document_id=original_document_id,
):
raise ValueError("Pipeline original document not found")
workflow = self.get_workflow(session=session, pipeline=pipeline, workflow_id=app_config.workflow_id)
@@ -237,7 +245,12 @@ class PipelineRunner(WorkflowBasedAppRunner):
generator = workflow_entry.run()
for event in generator:
self._update_document_status(event, document_ref)
self._update_document_status(
event,
workspace_id=dataset_workspace_id,
dataset_id=dataset_id,
document_id=document_id,
)
self._handle_event(workflow_entry, event)
def get_workflow(self, session: Session, pipeline: Pipeline, workflow_id: str) -> Workflow | None:
@@ -325,14 +338,21 @@ class PipelineRunner(WorkflowBasedAppRunner):
return graph
def _update_document_status(self, event: GraphEngineEvent, document_ref: DocumentRef | None) -> None:
def _update_document_status(
self,
event: GraphEngineEvent,
*,
workspace_id: str,
dataset_id: str,
document_id: str | None,
) -> None:
"""Set an owner-bound document to error after a failed graph run, if it exists."""
if not isinstance(event, GraphRunFailedEvent) or document_ref is None:
if not isinstance(event, GraphRunFailedEvent) or document_id is None:
return
with create_session() as session, session.begin():
document = DatasetRefService.get_document_by_ref(document_ref, session=session)
if document:
document.indexing_status = "error"
document.error = event.error or "Unknown error"
session.add(document)
self._documents.mark_failed(
workspace_id=workspace_id,
dataset_id=dataset_id,
document_id=document_id,
error=event.error or "Unknown error",
)
api/core/app/features/annotation_reply/annotation_reply.py
+1 -1
View File
@@ -10,7 +10,7 @@ from models.dataset import Dataset
from models.enums import CollectionBindingType, ConversationFromSource
from models.model import AnnotationReplyEnabledConfig, App, Message, MessageAnnotation, load_annotation_reply_config
from services.annotation_service import AppAnnotationService
from services.dataset_service import DatasetCollectionBindingService
from services.knowledge.dataset_service import DatasetCollectionBindingService
logger = logging.getLogger(__name__)
api/core/datasource/datasource_manager.py
+6 -21
View File
@@ -1,5 +1,5 @@
import logging
from collections.abc import Generator
from collections.abc import Generator, Mapping
from threading import Lock
from typing import Any, cast
@@ -33,7 +33,6 @@ from graphon.file import File, FileTransferMethod, FileType, get_file_type_by_mi
from graphon.node_events import NodeRunResult, StreamChunkEvent, StreamCompletedEvent
from models.model import UploadFile
from models.tools import ToolFile
from services.datasource_provider_service import DatasourceProviderService
logger = logging.getLogger(__name__)
_file_access_controller = DatabaseFileAccessController()
@@ -147,9 +146,7 @@ class DatasourceManager:
datasource_type: str,
provider_id: str,
tenant_id: str,
provider: str,
plugin_id: str,
credential_id: str,
credentials: Mapping[str, Any],
datasource_param: DatasourceParameter | None = None,
online_drive_request: OnlineDriveDownloadFileParam | None = None,
) -> Generator[DatasourceMessage, None, Any]:
@@ -166,18 +163,10 @@ class DatasourceManager:
datasource_type=ds_type,
)
dsp_service = DatasourceProviderService()
credentials = dsp_service.get_datasource_credentials(
tenant_id=tenant_id,
provider=provider,
plugin_id=plugin_id,
credential_id=credential_id,
)
if ds_type == DatasourceProviderType.ONLINE_DOCUMENT:
doc_runtime = cast(OnlineDocumentDatasourcePlugin, runtime)
if credentials:
doc_runtime.runtime.credentials = credentials
doc_runtime.runtime.credentials = dict(credentials)
if datasource_param is None:
raise ValueError("datasource_param is required for ONLINE_DOCUMENT streaming")
inner_gen: Generator[DatasourceMessage, None, None] = doc_runtime.get_online_document_page_content(
@@ -192,7 +181,7 @@ class DatasourceManager:
elif ds_type == DatasourceProviderType.ONLINE_DRIVE:
drive_runtime = cast(OnlineDriveDatasourcePlugin, runtime)
if credentials:
drive_runtime.runtime.credentials = credentials
drive_runtime.runtime.credentials = dict(credentials)
if online_drive_request is None:
raise ValueError("online_drive_request is required for ONLINE_DRIVE streaming")
inner_gen = drive_runtime.online_drive_download_file(
@@ -221,9 +210,7 @@ class DatasourceManager:
datasource_type: str,
provider_id: str,
tenant_id: str,
provider: str,
plugin_id: str,
credential_id: str,
credentials: Mapping[str, Any],
parameters_for_log: dict[str, Any],
datasource_info: dict[str, Any],
variable_pool: Any,
@@ -238,9 +225,7 @@ class DatasourceManager:
datasource_type=datasource_type,
provider_id=provider_id,
tenant_id=tenant_id,
provider=provider,
plugin_id=plugin_id,
credential_id=credential_id,
credentials=credentials,
datasource_param=datasource_param,
online_drive_request=online_drive_request,
)
api/core/indexing_runner.py
-938
View File
@@ -1,938 +0,0 @@
import concurrent.futures
import json
import logging
import re
import threading
import time
import uuid
from collections.abc import Mapping
from typing import Any
from flask import Flask, current_app
from sqlalchemy import delete, func, select, update
from sqlalchemy.orm import Session
from sqlalchemy.orm.exc import ObjectDeletedError
from configs import dify_config
from core.credit_usage import CreditUsageCreatedBy
from core.db.session_factory import session_factory
from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail
from core.errors.error import ProviderTokenNotInitError
from core.model_context import with_credit_usage_created_by
from core.model_manager import ModelInstance, ModelManager
from core.rag.cleaner.clean_processor import CleanProcessor
from core.rag.datasource.keyword.keyword_factory import Keyword
from core.rag.docstore.dataset_docstore import DatasetDocumentStore
from core.rag.embedding.token_counter import calculate_segment_token_counts
from core.rag.extractor.entity.datasource_type import DatasourceType
from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from core.rag.index_processor.index_processor_base import BaseIndexProcessor
from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
from core.rag.models.document import ChildDocument, Document
from core.rag.splitter.fixed_text_splitter import (
EnhanceRecursiveCharacterTextSplitter,
FixedRecursiveCharacterTextSplitter,
)
from core.rag.splitter.text_splitter import TextSplitter
from core.tools.utils.web_reader_tool import get_image_upload_file_ids
from enums import DeploymentEdition
from extensions.ext_redis import redis_client
from extensions.ext_storage import storage
from extensions.otel import propagate_context
from graphon.model_runtime.entities.model_entities import ModelType
from libs import helper
from libs.datetime_utils import naive_utc_now
from models import Account
from models.dataset import AutomaticRulesConfig, ChildChunk, Dataset, DatasetProcessRule, DocumentSegment
from models.dataset import Document as DatasetDocument
from models.enums import DataSourceType, IndexingStatus, ProcessRuleMode, SegmentStatus
from models.model import UploadFile
from services.vector_space_admission_service import VectorSpaceAdmissionService
logger = logging.getLogger(__name__)
class IndexingRunner:
def __init__(
self,
*,
enforce_vector_space_admission: bool = False,
):
self.storage = storage
self.enforce_vector_space_admission = enforce_vector_space_admission
@staticmethod
def _get_model_manager(tenant_id: str) -> ModelManager:
return ModelManager.for_tenant(tenant_id=tenant_id)
def _handle_indexing_error(self, document_id: str, error: Exception, session: Session) -> None:
"""Handle indexing errors by updating document status."""
logger.exception("consume document failed")
document = session.get(DatasetDocument, document_id)
if document:
document.indexing_status = IndexingStatus.ERROR
error_message = getattr(error, "description", str(error))
document.error = str(error_message)
document.stopped_at = naive_utc_now()
session.flush()
@with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
def run(self, dataset_documents: list[DatasetDocument], session: Session):
"""Run indexing with commits before slow transforms and parallel index workers.
The phase commits keep document locks short and make newly created segments
visible to the worker sessions used for keyword and vector indexing.
"""
vector_space_admission = VectorSpaceAdmissionService()
for dataset_document in dataset_documents:
document_id = dataset_document.id
try:
# Re-query the document to ensure it's bound to the current session
requeried_document = session.get(DatasetDocument, document_id)
if not requeried_document:
logger.warning("Document not found, skipping document id: %s", document_id)
continue
# get dataset
dataset = session.get(Dataset, requeried_document.dataset_id)
if not dataset:
raise ValueError("no dataset found")
# get the process rule
stmt = select(DatasetProcessRule).where(
DatasetProcessRule.id == requeried_document.dataset_process_rule_id
)
processing_rule = session.scalar(stmt)
if not processing_rule:
raise ValueError("no process rule found")
index_type = requeried_document.doc_form
index_processor = IndexProcessorFactory(index_type).init_index_processor()
# extract
text_docs = self._extract(index_processor, requeried_document, processing_rule.to_dict(), session)
session.commit()
# transform
current_user = session.get(Account, requeried_document.created_by)
if not current_user:
raise ValueError("no current user found")
current_user.set_tenant_id_with_session(dataset.tenant_id, session=session)
documents = self._transform(
index_processor,
dataset,
text_docs,
requeried_document.doc_language,
processing_rule.to_dict(),
current_user=current_user,
session=session,
)
if self.enforce_vector_space_admission:
vector_space_admission.ensure_document_can_be_indexed(
dataset=dataset,
document_id=requeried_document.id,
doc_form=requeried_document.doc_form,
documents=documents,
include_summaries=bool(requeried_document.need_summary),
session=session,
)
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
total_tokens = sum(token_counts)
# save segment
self._load_segments(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
token_counts=token_counts,
)
session.commit()
# load
self._load(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
total_tokens=total_tokens,
)
except DocumentIsPausedError:
raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
except ProviderTokenNotInitError as e:
self._handle_indexing_error(document_id, e, session)
except ObjectDeletedError:
logger.warning("Document deleted, document id: %s", document_id)
except Exception as e:
self._handle_indexing_error(document_id, e, session)
@with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
def run_in_splitting_status(self, dataset_document: DatasetDocument, session: Session):
"""Run the indexing process when the index_status is splitting."""
document_id = dataset_document.id
try:
# Re-query the document to ensure it's bound to the current session
requeried_document = session.get(DatasetDocument, document_id)
if not requeried_document:
logger.warning("Document not found: %s", document_id)
return
# get dataset
dataset = session.get(Dataset, requeried_document.dataset_id)
if not dataset:
raise ValueError("no dataset found")
# get exist document_segment list and delete
document_segments = session.scalars(
select(DocumentSegment).where(
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.document_id == requeried_document.id,
)
).all()
for document_segment in document_segments:
session.delete(document_segment)
if requeried_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
# delete child chunks
session.execute(delete(ChildChunk).where(ChildChunk.segment_id == document_segment.id))
session.commit()
# get the process rule
stmt = select(DatasetProcessRule).where(DatasetProcessRule.id == requeried_document.dataset_process_rule_id)
processing_rule = session.scalar(stmt)
if not processing_rule:
raise ValueError("no process rule found")
index_type = requeried_document.doc_form
index_processor = IndexProcessorFactory(index_type).init_index_processor()
# extract
text_docs = self._extract(index_processor, requeried_document, processing_rule.to_dict(), session)
session.commit()
# transform
current_user = session.get(Account, requeried_document.created_by)
if not current_user:
raise ValueError("no current user found")
current_user.set_tenant_id_with_session(dataset.tenant_id, session=session)
documents = self._transform(
index_processor,
dataset,
text_docs,
requeried_document.doc_language,
processing_rule.to_dict(),
current_user=current_user,
session=session,
)
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
total_tokens = sum(token_counts)
# save segment
self._load_segments(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
token_counts=token_counts,
)
session.commit()
# load
self._load(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
total_tokens=total_tokens,
)
except DocumentIsPausedError:
raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
except ProviderTokenNotInitError as e:
self._handle_indexing_error(document_id, e, session)
except Exception as e:
self._handle_indexing_error(document_id, e, session)
@with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
def run_in_indexing_status(self, dataset_document: DatasetDocument, session: Session):
"""Run the indexing process when the index_status is indexing."""
document_id = dataset_document.id
try:
# Re-query the document to ensure it's bound to the current session
requeried_document = session.get(DatasetDocument, document_id)
if not requeried_document:
logger.warning("Document not found: %s", document_id)
return
# get dataset
dataset = session.get(Dataset, requeried_document.dataset_id)
if not dataset:
raise ValueError("no dataset found")
# get existing document segments
document_segments = session.scalars(
select(DocumentSegment).where(
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.document_id == requeried_document.id,
)
).all()
documents = []
if document_segments:
for document_segment in document_segments:
# transform segment to node
if document_segment.status != SegmentStatus.COMPLETED:
document = Document(
page_content=document_segment.content,
metadata={
"doc_id": document_segment.index_node_id,
"doc_hash": document_segment.index_node_hash,
"document_id": document_segment.document_id,
"dataset_id": document_segment.dataset_id,
},
)
if requeried_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
child_chunks = document_segment.get_child_chunks(session=session)
if child_chunks:
child_documents = []
for child_chunk in child_chunks:
child_document = ChildDocument(
page_content=child_chunk.content,
metadata={
"doc_id": child_chunk.index_node_id,
"doc_hash": child_chunk.index_node_hash,
"document_id": document_segment.document_id,
"dataset_id": document_segment.dataset_id,
},
)
child_documents.append(child_document)
document.children = child_documents
documents.append(document)
# Preserve the full document total even when only incomplete segments are re-indexed.
total_tokens = sum(document_segment.tokens for document_segment in document_segments)
# build index
self._load(
session=session,
dataset=dataset,
dataset_document=requeried_document,
documents=documents,
total_tokens=total_tokens,
)
except DocumentIsPausedError:
raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
except ProviderTokenNotInitError as e:
self._handle_indexing_error(document_id, e, session)
except Exception as e:
self._handle_indexing_error(document_id, e, session)
@with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
def indexing_estimate(
self,
tenant_id: str,
extract_settings: list[ExtractSetting],
tmp_processing_rule: Mapping[str, Any],
doc_form: str | None = None,
doc_language: str = "English",
dataset_id: str | None = None,
indexing_technique: str = IndexTechniqueType.ECONOMY,
*,
session: Session,
) -> IndexingEstimate:
"""
Estimate the indexing for the document.
"""
# check document limit
if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
count = len(extract_settings)
batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT
if count > batch_upload_limit:
raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
embedding_model_instance = None
if dataset_id:
dataset = session.get(Dataset, dataset_id)
if not dataset:
raise ValueError("Dataset not found.")
if IndexTechniqueType.HIGH_QUALITY in {dataset.indexing_technique, indexing_technique}:
if dataset.embedding_model_provider:
embedding_model_instance = self._get_model_manager(tenant_id).get_model_instance(
tenant_id=tenant_id,
provider=dataset.embedding_model_provider,
model_type=ModelType.TEXT_EMBEDDING,
model=dataset.embedding_model,
)
else:
embedding_model_instance = self._get_model_manager(tenant_id).get_default_model_instance(
tenant_id=tenant_id,
model_type=ModelType.TEXT_EMBEDDING,
)
else:
if indexing_technique == IndexTechniqueType.HIGH_QUALITY:
embedding_model_instance = self._get_model_manager(tenant_id).get_default_model_instance(
tenant_id=tenant_id,
model_type=ModelType.TEXT_EMBEDDING,
)
# keep separate, avoid union-list ambiguity
preview_texts: list[PreviewDetail] = []
qa_preview_texts: list[QAPreviewDetail] = []
total_segments = 0
# doc_form represents the segmentation method (general, parent-child, QA)
index_type = doc_form
index_processor = IndexProcessorFactory(index_type).init_index_processor()
# one extract_setting is one source document
for extract_setting in extract_settings:
# extract
processing_rule = {
"mode": tmp_processing_rule["mode"],
"rules": tmp_processing_rule.get("rules"),
}
# Extract document content
text_docs = index_processor.extract(
extract_setting, process_rule_mode=tmp_processing_rule["mode"], session=session
)
# Cleaning and segmentation
documents = index_processor.transform(
text_docs,
current_user=None,
embedding_model_instance=embedding_model_instance,
process_rule=processing_rule,
tenant_id=tenant_id,
doc_language=doc_language,
preview=True,
session=session,
)
total_segments += len(documents)
for document in documents:
preview_count = len(qa_preview_texts) if doc_form == "qa_model" else len(preview_texts)
if preview_count < 10:
if doc_form and doc_form == "qa_model":
qa_detail = QAPreviewDetail(
question=document.page_content, answer=document.metadata.get("answer") or ""
)
qa_preview_texts.append(qa_detail)
else:
preview_detail = PreviewDetail(content=document.page_content)
if document.children:
preview_detail.child_chunks = [child.page_content for child in document.children]
preview_texts.append(preview_detail)
# delete image files and related db records
image_upload_file_ids = get_image_upload_file_ids(document.page_content)
for upload_file_id in image_upload_file_ids:
stmt = select(UploadFile).where(UploadFile.id == upload_file_id)
image_file = session.scalar(stmt)
if image_file is None:
continue
try:
storage.delete(image_file.key)
except Exception:
logger.exception(
"Delete image_files failed while indexing_estimate, \
image_upload_file_is: %s",
upload_file_id,
)
session.delete(image_file)
# Persist preview cleanup and release the caller transaction before
# summary workers query through their own sessions.
session.commit()
if doc_form and doc_form == "qa_model":
return IndexingEstimate(total_segments=total_segments * 20, qa_preview=qa_preview_texts, preview=[])
# Generate summary preview
summary_index_setting = tmp_processing_rule.get("summary_index_setting")
if summary_index_setting and summary_index_setting.get("enable") and preview_texts:
preview_texts = index_processor.generate_summary_preview(
tenant_id, preview_texts, summary_index_setting, doc_language, session=session
)
return IndexingEstimate(total_segments=total_segments, preview=preview_texts)
def _extract(
self,
index_processor: BaseIndexProcessor,
dataset_document: DatasetDocument,
process_rule: Mapping[str, Any],
session: Session,
) -> list[Document]:
data_source_info = dataset_document.data_source_info_dict
text_docs = []
match dataset_document.data_source_type:
case DataSourceType.UPLOAD_FILE:
if not data_source_info or "upload_file_id" not in data_source_info:
raise ValueError("no upload file found")
stmt = select(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"])
file_detail = session.scalars(stmt).one_or_none()
if file_detail:
extract_setting = ExtractSetting(
datasource_type=DatasourceType.FILE,
upload_file=file_detail,
document_model=dataset_document.doc_form,
)
text_docs = index_processor.extract(
extract_setting, process_rule_mode=process_rule["mode"], session=session
)
case DataSourceType.NOTION_IMPORT:
if (
not data_source_info
or "notion_workspace_id" not in data_source_info
or "notion_page_id" not in data_source_info
):
raise ValueError("no notion import info found")
extract_setting = ExtractSetting(
datasource_type=DatasourceType.NOTION,
notion_info=NotionInfo.model_validate(
{
"credential_id": data_source_info.get("credential_id"),
"notion_workspace_id": data_source_info["notion_workspace_id"],
"notion_obj_id": data_source_info["notion_page_id"],
"notion_page_type": data_source_info["type"],
"document": dataset_document,
"tenant_id": dataset_document.tenant_id,
}
),
document_model=dataset_document.doc_form,
)
text_docs = index_processor.extract(
extract_setting, process_rule_mode=process_rule["mode"], session=session
)
case DataSourceType.WEBSITE_CRAWL:
if (
not data_source_info
or "provider" not in data_source_info
or "url" not in data_source_info
or "job_id" not in data_source_info
):
raise ValueError("no website import info found")
extract_setting = ExtractSetting(
datasource_type=DatasourceType.WEBSITE,
website_info=WebsiteInfo.model_validate(
{
"provider": data_source_info["provider"],
"job_id": data_source_info["job_id"],
"tenant_id": dataset_document.tenant_id,
"url": data_source_info["url"],
"mode": data_source_info["mode"],
"only_main_content": data_source_info["only_main_content"],
}
),
document_model=dataset_document.doc_form,
)
text_docs = index_processor.extract(
extract_setting, process_rule_mode=process_rule["mode"], session=session
)
case _:
return []
# update document status to splitting
self._update_document_index_status(
session=session,
document_id=dataset_document.id,
after_indexing_status=IndexingStatus.SPLITTING,
extra_update_params={
DatasetDocument.parsing_completed_at: naive_utc_now(),
},
)
# replace doc id to document model id
for text_doc in text_docs:
if text_doc.metadata is not None:
text_doc.metadata["document_id"] = dataset_document.id
text_doc.metadata["dataset_id"] = dataset_document.dataset_id
return text_docs
@staticmethod
def filter_string(text: str):
text = re.sub(r"<\|", "<", text)
text = re.sub(r"\|>", ">", text)
text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]", "", text)
# Unicode U+FFFE
text = re.sub("\ufffe", "", text)
return text
@staticmethod
def _get_splitter(
processing_rule_mode: str,
max_tokens: int,
chunk_overlap: int,
separator: str,
embedding_model_instance: ModelInstance | None,
) -> TextSplitter:
"""
Get the NodeParser object according to the processing rule.
"""
character_splitter: TextSplitter
if processing_rule_mode in ["custom", "hierarchical"]:
# The user-defined segmentation rule
max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH
if max_tokens < 50 or max_tokens > max_segmentation_tokens_length:
raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.")
if separator:
separator = separator.replace("\\n", "\n")
character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder(
chunk_size=max_tokens,
chunk_overlap=chunk_overlap,
fixed_separator=separator,
separators=["\n\n", "。", ". ", " ", ""],
embedding_model_instance=embedding_model_instance,
)
else:
# Automatic segmentation
automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"])
character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder(
chunk_size=automatic_rules["max_tokens"],
chunk_overlap=automatic_rules["chunk_overlap"],
separators=["\n\n", "。", ". ", " ", ""],
embedding_model_instance=embedding_model_instance,
)
return character_splitter
def _split_to_documents_for_estimate(
self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule
) -> list[Document]:
"""
Split the text documents into nodes.
"""
all_documents: list[Document] = []
for text_doc in text_docs:
# document clean
document_text = self._document_clean(text_doc.page_content, processing_rule)
text_doc.page_content = document_text
# parse document to nodes
documents = splitter.split_documents([text_doc])
split_documents = []
for document in documents:
if document.page_content is None or not document.page_content.strip():
continue
if document.metadata is not None:
doc_id = str(uuid.uuid4())
hash = helper.generate_text_hash(document.page_content)
document.metadata["doc_id"] = doc_id
document.metadata["doc_hash"] = hash
split_documents.append(document)
all_documents.extend(split_documents)
return all_documents
@staticmethod
def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str:
"""
Clean the document text according to the processing rules.
"""
rules: AutomaticRulesConfig | dict[str, Any]
if processing_rule.mode == ProcessRuleMode.AUTOMATIC:
rules = DatasetProcessRule.AUTOMATIC_RULES
else:
rules = json.loads(processing_rule.rules) if processing_rule.rules else {}
document_text = CleanProcessor.clean(text, {"rules": rules})
return document_text
@staticmethod
def format_split_text(text: str) -> list[QAPreviewDetail]:
regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)"
matches = re.findall(regex, text, re.UNICODE)
return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a]
def _load(
self,
session: Session,
dataset: Dataset,
dataset_document: DatasetDocument,
documents: list[Document],
total_tokens: int,
) -> None:
"""Build indexes and mark the document complete using the token total computed before hash sharding."""
# Build indexes using the existing hash-based worker groups.
indexing_start_at = time.perf_counter()
create_keyword_thread = None
if (
dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
and dataset.indexing_technique == IndexTechniqueType.ECONOMY
):
# create keyword index
create_keyword_thread = threading.Thread(
target=propagate_context(self._process_keyword_index),
args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents), # type: ignore
)
create_keyword_thread.start()
max_workers = 10
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = []
# Distribute documents into multiple groups based on the hash values of page_content
# This is done to prevent multiple threads from processing the same document,
# Thereby avoiding potential database insertion deadlocks
document_groups: list[list[Document]] = [[] for _ in range(max_workers)]
for document in documents:
hash = helper.generate_text_hash(document.page_content)
group_index = int(hash, 16) % max_workers
document_groups[group_index].append(document)
for chunk_documents in document_groups:
if len(chunk_documents) == 0:
continue
futures.append(
executor.submit(
propagate_context(self._process_chunk),
current_app._get_current_object(), # type: ignore
dataset_document.doc_form,
chunk_documents,
dataset.id,
dataset_document.id,
)
)
for future in futures:
future.result()
if (
dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
and dataset.indexing_technique == IndexTechniqueType.ECONOMY
and create_keyword_thread is not None
):
create_keyword_thread.join()
indexing_end_at = time.perf_counter()
# update document status to completed
self._update_document_index_status(
session=session,
document_id=dataset_document.id,
after_indexing_status=IndexingStatus.COMPLETED,
extra_update_params={
DatasetDocument.tokens: total_tokens,
DatasetDocument.completed_at: naive_utc_now(),
DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
DatasetDocument.error: None,
},
)
@staticmethod
def _process_keyword_index(flask_app: Flask, dataset_id: str, document_id: str, documents: list[Document]):
with flask_app.app_context():
with session_factory.create_session() as session:
dataset = session.get(Dataset, dataset_id)
if not dataset:
raise ValueError("no dataset found")
keyword = Keyword(dataset)
keyword.create(documents, session)
if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
document_ids = [document.metadata["doc_id"] for document in documents]
session.execute(
update(DocumentSegment)
.where(
DocumentSegment.document_id == document_id,
DocumentSegment.dataset_id == dataset_id,
DocumentSegment.index_node_id.in_(document_ids),
DocumentSegment.status == SegmentStatus.INDEXING,
)
.values(
status=SegmentStatus.COMPLETED,
enabled=True,
completed_at=naive_utc_now(),
)
)
session.commit()
def _process_chunk(
self,
flask_app: Flask,
index_type: str,
chunk_documents: list[Document],
dataset_id: str,
dataset_document_id: str,
) -> None:
with flask_app.app_context():
with session_factory.create_session() as session:
dataset = session.get(Dataset, dataset_id)
if not dataset:
raise ValueError("no dataset found")
dataset_document = session.get(DatasetDocument, dataset_document_id)
if not dataset_document:
raise ValueError("no document found")
# check document is paused
self._check_document_paused_status(dataset_document.id)
multimodal_documents = []
for document in chunk_documents:
if document.attachments and dataset.is_multimodal:
multimodal_documents.extend(document.attachments)
# load index
index_processor = IndexProcessorFactory(index_type).init_index_processor()
index_processor.load(
dataset,
chunk_documents,
multimodal_documents=multimodal_documents,
with_keywords=False,
session=session,
)
document_ids = [document.metadata["doc_id"] for document in chunk_documents]
session.execute(
update(DocumentSegment)
.where(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.index_node_id.in_(document_ids),
DocumentSegment.status == SegmentStatus.INDEXING,
)
.values(
status=SegmentStatus.COMPLETED,
enabled=True,
completed_at=naive_utc_now(),
)
)
session.commit()
@staticmethod
def _check_document_paused_status(document_id: str):
indexing_cache_key = f"document_{document_id}_is_paused"
result = redis_client.get(indexing_cache_key)
if result:
raise DocumentIsPausedError()
@staticmethod
def _update_document_index_status(
document_id: str,
after_indexing_status: IndexingStatus,
extra_update_params: Mapping[Any, Any] | None = None,
*,
session: Session,
):
"""
Update the document indexing status.
"""
count = (
session.scalar(
select(func.count())
.select_from(DatasetDocument)
.where(DatasetDocument.id == document_id, DatasetDocument.is_paused == True)
)
or 0
)
if count > 0:
raise DocumentIsPausedError()
document = session.get(DatasetDocument, document_id)
if not document:
raise DocumentIsDeletedPausedError()
update_params = {DatasetDocument.indexing_status: after_indexing_status}
if extra_update_params:
update_params.update(extra_update_params)
session.execute(update(DatasetDocument).where(DatasetDocument.id == document_id).values(update_params)) # type: ignore
session.flush()
@staticmethod
def _update_segments_by_document(dataset_document_id: str, update_params: Mapping[Any, Any], session: Session):
"""
Update the document segment by document id.
"""
session.execute(
update(DocumentSegment).where(DocumentSegment.document_id == dataset_document_id).values(update_params)
)
session.flush()
def _transform(
self,
index_processor: BaseIndexProcessor,
dataset: Dataset,
text_docs: list[Document],
doc_language: str,
process_rule: Mapping[str, Any],
current_user: Account | None = None,
*,
session: Session,
) -> list[Document]:
# get embedding model instance
embedding_model_instance = None
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
if dataset.embedding_model_provider:
embedding_model_instance = self._get_model_manager(dataset.tenant_id).get_model_instance(
tenant_id=dataset.tenant_id,
provider=dataset.embedding_model_provider,
model_type=ModelType.TEXT_EMBEDDING,
model=dataset.embedding_model,
)
else:
embedding_model_instance = self._get_model_manager(dataset.tenant_id).get_default_model_instance(
tenant_id=dataset.tenant_id,
model_type=ModelType.TEXT_EMBEDDING,
)
documents = index_processor.transform(
text_docs,
current_user,
embedding_model_instance=embedding_model_instance,
process_rule=process_rule,
tenant_id=dataset.tenant_id,
doc_language=doc_language,
session=session,
)
return documents
def _load_segments(
self,
session: Session,
dataset: Dataset,
dataset_document: DatasetDocument,
documents: list[Document],
token_counts: list[int],
) -> None:
"""Persist transformed documents and their precomputed token counts before indexing starts."""
# save node to document segment
doc_store = DatasetDocumentStore(
dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
)
# add document segments
doc_store.add_documents(
session=session,
docs=documents,
save_child=dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX,
token_counts=token_counts,
)
# update document status to indexing
cur_time = naive_utc_now()
self._update_document_index_status(
session=session,
document_id=dataset_document.id,
after_indexing_status=IndexingStatus.INDEXING,
extra_update_params={
DatasetDocument.cleaning_completed_at: cur_time,
DatasetDocument.splitting_completed_at: cur_time,
DatasetDocument.word_count: sum(len(doc.page_content) for doc in documents),
},
)
# update segment status to indexing
self._update_segments_by_document(
session=session,
dataset_document_id=dataset_document.id,
update_params={
DocumentSegment.status: SegmentStatus.INDEXING,
DocumentSegment.indexing_at: naive_utc_now(),
},
)
class DocumentIsPausedError(Exception):
pass
class DocumentIsDeletedPausedError(Exception):
pass
api/core/provider_manager.py
-45
View File
@@ -65,7 +65,6 @@ from services.feature_service import FeatureService
if TYPE_CHECKING:
from graphon.model_runtime.protocols.runtime import ModelRuntime
from models.account import Account
logger = logging.getLogger(__name__)
@@ -1162,50 +1161,6 @@ class ProviderManager:
provider_names.append(str(model_provider_id))
return provider_names
@staticmethod
def get_provider_available_credentials(
tenant_id: str,
provider_name: str,
user: Account | None = None,
) -> list[CredentialConfiguration]:
"""
Get provider all credentials, filtered by visibility.
:param tenant_id: workspace id
:param provider_name: provider name
:param user: current user (id + admin flag drive the visibility filter)
:return:
"""
from models.credential_permission import CredentialType as CredPermType
from services.credential_permission_service import CredentialPermissionService
with session_factory.create_session() as session:
stmt = (
select(ProviderCredential)
.where(
ProviderCredential.tenant_id == tenant_id,
ProviderCredential.provider_name.in_(ProviderManager._get_provider_names(provider_name)),
)
.order_by(ProviderCredential.created_at.desc())
)
if user is not None:
stmt = CredentialPermissionService.apply_visibility_filter(
stmt,
model_id_column=ProviderCredential.id,
model_user_id_column=ProviderCredential.user_id,
model_visibility_column=ProviderCredential.visibility,
credential_type=CredPermType.PROVIDER_CREDENTIAL,
user=user,
)
available_credentials = session.scalars(stmt).all()
return [
CredentialConfiguration(credential_id=credential.id, credential_name=credential.credential_name)
for credential in available_credentials
]
@staticmethod
def get_provider_model_available_credentials(
tenant_id: str, provider_name: str, model_name: str, model_type: str
api/core/rag/datasource/keyword/jieba/jieba.py
+100 -165
View File
@@ -1,7 +1,8 @@
from collections import defaultdict
from typing import Any, TypedDict, override
from collections.abc import Callable, Mapping, Sequence
from contextlib import AbstractContextManager
from typing import Any, override
import orjson
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.orm import Session
@@ -12,12 +13,15 @@ from core.rag.datasource.keyword.keyword_base import BaseKeyword
from core.rag.models.document import Document
from extensions.ext_redis import redis_client
from extensions.ext_storage import storage
from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment
class PreSegmentData(TypedDict):
segment: DocumentSegment
keywords: list[str]
from models.dataset import Dataset, DocumentSegment
from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table
from repositories.knowledge.keyword_table_repository import (
keyword_file_key,
load_dataset_keyword_table,
load_keyword_table,
persist_keyword_table,
save_keyword_table,
)
class KeywordTableConfig(BaseModel):
@@ -31,80 +35,114 @@ class Jieba(BaseKeyword):
@override
def create(self, texts: list[Document], session: Session, **kwargs: Any) -> BaseKeyword:
lock_name = f"keyword_indexing_lock_{self.dataset.id}"
with redis_client.lock(lock_name, timeout=600):
keyword_table_handler = JiebaKeywordTableHandler()
keyword_table = self._get_dataset_keyword_table(session=session)
keyword_number = self.dataset.keyword_number or self._config.max_keywords_per_chunk
for text in texts:
keywords = keyword_table_handler.extract_keywords(text.page_content, keyword_number)
if text.metadata is not None:
self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords), session)
keyword_table = self._add_text_to_keyword_table(
keyword_table or {}, text.metadata["doc_id"], list(keywords)
)
self._save_dataset_keyword_table(keyword_table, session)
return self
self.add_texts(texts, session)
return self
@override
def add_texts(self, texts: list[Document], session: Session, **kwargs: Any):
lock_name = f"keyword_indexing_lock_{self.dataset.id}"
with redis_client.lock(lock_name, timeout=600):
keyword_table_handler = JiebaKeywordTableHandler()
self.update_texts(
texts,
read=lambda: self._read_keyword_record(session),
write=lambda storage_type, data, keywords: persist_keyword_table(
session,
tenant_id=self.dataset.tenant_id,
dataset_id=self.dataset.id,
storage_type=storage_type,
data=data,
keywords=keywords,
),
lock=redis_client.lock(f"keyword_indexing_lock_{self.dataset.id}", timeout=600),
keywords_list=kwargs.get("keywords_list"),
)
keyword_table = self._get_dataset_keyword_table(session=session)
keywords_list = kwargs.get("keywords_list")
def update_texts(
self,
texts: Sequence[Document],
*,
read: Callable[[], tuple[str, str | None]],
write: Callable[[str, str, Mapping[str, Sequence[str]]], None],
lock: AbstractContextManager[object],
keywords_list: Sequence[Sequence[str] | None] | None = None,
replace_existing: bool = False,
deleted_ids: Sequence[str] = (),
) -> None:
"""Mutate the shared index with explicit persistence and one dataset lock.
Repository callbacks may close their sessions before returning, keeping
extraction and object-storage I/O outside database transactions.
"""
with lock:
storage_type, data = read()
payload = load_keyword_table(
storage=storage,
tenant_id=self.dataset.tenant_id,
dataset_id=self.dataset.id,
storage_type=storage_type,
data=data,
)
table = dict(payload["__data__"]["table"] or {}) if payload else {}
removed = list(deleted_ids)
if replace_existing:
removed.extend(text.metadata["doc_id"] for text in texts)
if removed:
table = self._delete_ids_from_keyword_table(table, removed)
handler: JiebaKeywordTableHandler | None = None
keyword_number = self.dataset.keyword_number or self._config.max_keywords_per_chunk
for i in range(len(texts)):
text = texts[i]
if keywords_list:
keywords = keywords_list[i]
if not keywords:
keywords = keyword_table_handler.extract_keywords(text.page_content, keyword_number)
else:
keywords = keyword_table_handler.extract_keywords(text.page_content, keyword_number)
if text.metadata is not None:
self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords), session)
keyword_table = self._add_text_to_keyword_table(
keyword_table or {}, text.metadata["doc_id"], list(keywords)
)
selected: dict[str, Sequence[str]] = {}
for position, text in enumerate(texts):
keywords = keywords_list[position] if keywords_list else None
if not keywords:
if handler is None:
handler = JiebaKeywordTableHandler()
keywords = list(handler.extract_keywords(text.page_content, keyword_number))
node_id = text.metadata["doc_id"]
selected[node_id] = list(keywords)
table = self._add_text_to_keyword_table(table, node_id, list(keywords))
encoded = save_keyword_table(
storage=storage,
tenant_id=self.dataset.tenant_id,
dataset_id=self.dataset.id,
storage_type=storage_type,
table=table,
)
write(storage_type, encoded, selected)
self._save_dataset_keyword_table(keyword_table, session)
def _read_keyword_record(self, session: Session) -> tuple[str, str | None]:
row = get_dataset_keyword_table(self.dataset, session=session)
return (row.data_source_type, row.keyword_table) if row else (dify_config.KEYWORD_DATA_SOURCE_TYPE, None)
@override
def text_exists(self, id: str, *, session: Session) -> bool:
dataset_keyword_table = self.dataset.get_dataset_keyword_table(session=session)
keyword_table = None
keyword_table_dict = (
dataset_keyword_table.get_keyword_table_dict(session=session) if dataset_keyword_table else None
)
keyword_table_dict = load_dataset_keyword_table(self.dataset, session=session, storage=storage)
if keyword_table_dict:
data: Any = keyword_table_dict["__data__"]
keyword_table = dict(data["table"])
if keyword_table is None:
return False
return id in set.union(*keyword_table.values())
return any(id in node_ids for node_ids in keyword_table.values())
@override
def delete_by_ids(self, ids: list[str], session: Session, **kwargs: Any):
lock_name = f"keyword_indexing_lock_{self.dataset.id}"
with redis_client.lock(lock_name, timeout=600):
keyword_table = self._get_dataset_keyword_table(session)
if keyword_table is not None:
keyword_table = self._delete_ids_from_keyword_table(keyword_table, ids)
self._save_dataset_keyword_table(keyword_table, session)
self.update_texts(
[],
read=lambda: self._read_keyword_record(session),
write=lambda storage_type, data, keywords: persist_keyword_table(
session,
tenant_id=self.dataset.tenant_id,
dataset_id=self.dataset.id,
storage_type=storage_type,
data=data,
keywords=keywords,
),
lock=redis_client.lock(f"keyword_indexing_lock_{self.dataset.id}", timeout=600),
deleted_ids=ids,
)
@override
def search(self, query: str, *, session: Session, **kwargs: Any) -> list[Document]:
dataset_keyword_table = self.dataset.get_dataset_keyword_table(session=session)
keyword_table = None
keyword_table_dict = (
dataset_keyword_table.get_keyword_table_dict(session=session) if dataset_keyword_table else None
)
keyword_table_dict = load_dataset_keyword_table(self.dataset, session=session, storage=storage)
if keyword_table_dict:
data: Any = keyword_table_dict["__data__"]
keyword_table = dict(data["table"])
@@ -145,62 +183,14 @@ class Jieba(BaseKeyword):
def delete(self, *, session: Session):
lock_name = f"keyword_indexing_lock_{self.dataset.id}"
with redis_client.lock(lock_name, timeout=600):
dataset_keyword_table = self.dataset.get_dataset_keyword_table(session=session)
dataset_keyword_table = get_dataset_keyword_table(self.dataset, session=session)
if dataset_keyword_table:
session.delete(dataset_keyword_table)
session.commit()
if dataset_keyword_table.data_source_type != "database":
file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt"
file_key = keyword_file_key(self.dataset.tenant_id, self.dataset.id)
storage.delete(file_key)
def _save_dataset_keyword_table(self, keyword_table: dict[str, set[str]] | None, session: Session):
keyword_table_dict = {
"__type__": "keyword_table",
"__data__": {"index_id": self.dataset.id, "summary": None, "table": keyword_table},
}
dataset_keyword_table = session.scalar(
select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.dataset.id)
)
keyword_data_source_type = dataset_keyword_table.data_source_type if dataset_keyword_table else "file"
if keyword_data_source_type == "database":
if dataset_keyword_table is None:
return
dataset_keyword_table.keyword_table = dumps_with_sets(keyword_table_dict)
session.flush()
else:
file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt"
if storage.exists(file_key):
storage.delete(file_key)
storage.save(file_key, dumps_with_sets(keyword_table_dict).encode("utf-8"))
def _get_dataset_keyword_table(self, session: Session) -> dict[str, set[str]] | None:
dataset_keyword_table = session.scalar(
select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.dataset.id)
)
if dataset_keyword_table:
keyword_table_dict = dataset_keyword_table.get_keyword_table_dict(session=session)
if keyword_table_dict:
data: Any = keyword_table_dict["__data__"]
return dict(data["table"])
else:
keyword_data_source_type = dify_config.KEYWORD_DATA_SOURCE_TYPE
dataset_keyword_table = DatasetKeywordTable(
dataset_id=self.dataset.id,
keyword_table="",
data_source_type=keyword_data_source_type,
)
if keyword_data_source_type == "database":
dataset_keyword_table.keyword_table = dumps_with_sets(
{
"__type__": "keyword_table",
"__data__": {"index_id": self.dataset.id, "summary": None, "table": {}},
}
)
session.add(dataset_keyword_table)
session.flush()
return {}
def _add_text_to_keyword_table(
self, keyword_table: dict[str, set[str]], id: str, keywords: list[str]
) -> dict[str, set[str]]:
@@ -245,58 +235,3 @@ class Jieba(BaseKeyword):
)
return sorted_chunk_indices[:k]
def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str], session: Session):
stmt = select(DocumentSegment).where(
DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id
)
document_segment = session.scalar(stmt)
if document_segment:
document_segment.keywords = keywords
session.add(document_segment)
session.flush()
def create_segment_keywords(self, node_id: str, keywords: list[str], session: Session):
keyword_table = self._get_dataset_keyword_table(session)
self._update_segment_keywords(self.dataset.id, node_id, keywords, session)
keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords)
self._save_dataset_keyword_table(keyword_table, session)
def multi_create_segment_keywords(self, pre_segment_data_list: list[PreSegmentData], session: Session):
keyword_table_handler = JiebaKeywordTableHandler()
keyword_table = self._get_dataset_keyword_table(session)
for pre_segment_data in pre_segment_data_list:
segment = pre_segment_data["segment"]
if pre_segment_data["keywords"]:
segment.keywords = pre_segment_data["keywords"]
assert segment.index_node_id
keyword_table = self._add_text_to_keyword_table(
keyword_table or {}, segment.index_node_id, pre_segment_data["keywords"]
)
else:
keyword_number = self.dataset.keyword_number or self._config.max_keywords_per_chunk
keywords = keyword_table_handler.extract_keywords(segment.content, keyword_number)
segment.keywords = list(keywords)
assert segment.index_node_id
keyword_table = self._add_text_to_keyword_table(
keyword_table or {}, segment.index_node_id, list(keywords)
)
self._save_dataset_keyword_table(keyword_table, session)
def update_segment_keywords_index(self, node_id: str, keywords: list[str], session: Session):
keyword_table = self._get_dataset_keyword_table(session)
keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords)
self._save_dataset_keyword_table(keyword_table, session)
def set_orjson_default(obj: Any):
"""Default function for orjson serialization of set types"""
if isinstance(obj, set):
return list(obj)
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
def dumps_with_sets(obj: Any) -> str:
"""JSON dumps with set support using orjson"""
return orjson.dumps(obj, default=set_orjson_default).decode("utf-8")
api/core/rag/datasource/retrieval_service.py
+36 -7
View File
@@ -5,7 +5,7 @@ from concurrent.futures import ThreadPoolExecutor
from typing import Any, NotRequired, TypedDict
from flask import Flask, current_app
from sqlalchemy import select
from sqlalchemy import select, tuple_
from sqlalchemy.orm import Session, load_only
from configs import dify_config
@@ -35,7 +35,7 @@ from models.dataset import (
)
from models.dataset import Document as DatasetDocument
from models.model import UploadFile
from services.external_knowledge_service import ExternalDatasetService
from services.knowledge.external.service import ExternalDatasetService
class SegmentAttachmentResult(TypedDict):
@@ -486,7 +486,14 @@ class RetrievalService:
for doc in session.scalars(
select(DatasetDocument)
.where(DatasetDocument.id.in_(document_ids))
.options(load_only(DatasetDocument.id, DatasetDocument.doc_form, DatasetDocument.dataset_id))
.options(
load_only(
DatasetDocument.id,
DatasetDocument.tenant_id,
DatasetDocument.doc_form,
DatasetDocument.dataset_id,
)
)
).all()
}
@@ -559,7 +566,14 @@ class RetrievalService:
doc_segment_map: dict[str, list[str]] = {}
segment_summary_map: dict[str, str] = {} # Map segment_id to summary content
attachments = cls.get_segment_attachment_infos(image_doc_ids, session)
dataset_refs = tuple(
{(document.tenant_id, document.dataset_id) for document in valid_dataset_documents.values()}
)
attachments = cls.get_segment_attachment_infos(
image_doc_ids,
dataset_refs=dataset_refs,
session=session,
)
for attachment in attachments:
segment_ids.append(attachment["segment_id"])
@@ -947,15 +961,30 @@ class RetrievalService:
@classmethod
def get_segment_attachment_infos(
cls, attachment_ids: list[str], session: Session
cls,
attachment_ids: list[str],
*,
dataset_refs: Sequence[tuple[str, str]],
session: Session,
) -> list[SegmentAttachmentInfoResult]:
if not attachment_ids or not dataset_refs:
return []
attachment_infos: list[SegmentAttachmentInfoResult] = []
granted_upload_file_ids: list[str] = []
upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(attachment_ids))).all()
tenant_ids = tuple({tenant_id for tenant_id, _ in dataset_refs})
upload_files = session.scalars(
select(UploadFile).where(
UploadFile.id.in_(attachment_ids),
UploadFile.tenant_id.in_(tenant_ids),
)
).all()
if upload_files:
upload_file_ids = [upload_file.id for upload_file in upload_files]
attachment_bindings = session.scalars(
select(SegmentAttachmentBinding).where(SegmentAttachmentBinding.attachment_id.in_(upload_file_ids))
select(SegmentAttachmentBinding).where(
SegmentAttachmentBinding.attachment_id.in_(upload_file_ids),
tuple_(SegmentAttachmentBinding.tenant_id, SegmentAttachmentBinding.dataset_id).in_(dataset_refs),
)
).all()
attachment_binding_map = {binding.attachment_id: binding for binding in attachment_bindings}
api/core/rag/datasource/vdb/vector_factory.py
+32 -6
View File
@@ -2,6 +2,7 @@ import base64
import logging
import time
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import Any, override
from sqlalchemy import select
@@ -105,7 +106,14 @@ class _LazyEmbeddings(Embeddings):
class Vector:
def __init__(self, dataset: Dataset, attributes: list | None = None, *, session: Session):
def __init__(
self,
dataset: Dataset,
attributes: list | None = None,
*,
session: Session | None,
vector_type: str | None = None,
):
if attributes is None:
# `is_summary` and `original_chunk_id` are stored on summary vectors
# by `SummaryIndexService` and read back by `RetrievalService` to
@@ -132,7 +140,14 @@ class Vector:
self._embeddings: Embeddings = _LazyEmbeddings(dataset)
self._attributes = attributes
self._session = session
self._vector_processor = self._init_vector(session=session)
if vector_type is not None:
self._vector_processor = self.get_vector_factory(vector_type)().init_vector(
dataset, self._attributes, self._embeddings
)
else:
if session is None:
raise ValueError("A resolved vector type is required without a database session")
self._vector_processor = self._init_vector(session=session)
@staticmethod
def resolve_vector_type(dataset: Dataset, *, session: Session) -> str:
@@ -192,7 +207,9 @@ class Vector:
self._vector_processor.create(texts=batch, embeddings=batch_embeddings, **kwargs)
logger.info("Embedding %s texts took %s s", len(texts), time.time() - start)
def create_multimodal(self, file_documents: list | None = None, **kwargs):
def create_multimodal(
self, file_documents: list | None = None, *, upload_files: Mapping[str, UploadFile] | None = None, **kwargs
):
if file_documents:
start = time.time()
logger.info("start embedding %s files %s", len(file_documents), start)
@@ -205,9 +222,16 @@ class Vector:
# Batch query all upload files to avoid N+1 queries
attachment_ids = [doc.metadata["doc_id"] for doc in batch]
stmt = select(UploadFile).where(UploadFile.id.in_(attachment_ids))
upload_files = self._session.scalars(stmt).all()
upload_file_map = {str(f.id): f for f in upload_files}
upload_file_map: Mapping[str, UploadFile]
if upload_files is None:
if self._session is None:
raise ValueError("Upload files must be loaded before indexing without a database session")
stmt = select(UploadFile).where(
UploadFile.id.in_(attachment_ids), UploadFile.tenant_id == self._dataset.tenant_id
)
upload_file_map = {str(f.id): f for f in self._session.scalars(stmt).all()}
else:
upload_file_map = upload_files
file_base64_list = []
real_batch = []
@@ -264,6 +288,8 @@ class Vector:
return self._vector_processor.search_by_vector(query_vector, **kwargs)
def search_by_file(self, file_id: str, **kwargs: Any) -> list[Document]:
if self._session is None:
raise ValueError("File search requires a database session")
upload_file: UploadFile | None = self._session.get(UploadFile, file_id)
if not upload_file:
api/core/rag/docstore/dataset_docstore.py
+5 -1
View File
@@ -235,8 +235,12 @@ class DatasetDocumentStore:
def get_document_segment(self, doc_id: str, session: Session) -> DocumentSegment | None:
stmt = select(DocumentSegment).where(
DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id
DocumentSegment.tenant_id == self._dataset.tenant_id,
DocumentSegment.dataset_id == self._dataset.id,
DocumentSegment.index_node_id == doc_id,
)
if self._document_id is not None:
stmt = stmt.where(DocumentSegment.document_id == self._document_id)
document_segment = session.scalar(stmt)
return document_segment
api/core/rag/entities/__init__.py
+8 -1
View File
@@ -3,7 +3,13 @@ from core.rag.entities.context_entities import DocumentContext
from core.rag.entities.event import DatasourceCompletedEvent, DatasourceErrorEvent, DatasourceProcessingEvent
from core.rag.entities.index_entities import EconomySetting, EmbeddingSetting, IndexMethod
from core.rag.entities.metadata_entities import Condition, MetadataFilteringCondition, SupportedComparisonOperator
from core.rag.entities.processing_entities import ParentMode, PreProcessingRule, Rule, Segmentation
from core.rag.entities.processing_entities import (
ParentMode,
PreProcessingRule,
PreProcessingRuleKey,
Rule,
Segmentation,
)
from core.rag.entities.retrieval_settings import (
KeywordSetting,
RerankingModelConfig,
@@ -24,6 +30,7 @@ __all__ = [
"MetadataFilteringCondition",
"ParentMode",
"PreProcessingRule",
"PreProcessingRuleKey",
"RerankingModelConfig",
"RetrievalSourceMetadata",
"Rule",
api/core/rag/entities/processing_entities.py
+4 -12
View File
@@ -1,7 +1,7 @@
from enum import StrEnum
from typing import Annotated, Literal
from typing import Literal
from pydantic import AliasChoices, BaseModel, Field, WithJsonSchema
from pydantic import AliasChoices, BaseModel, Field
class ParentMode(StrEnum):
@@ -9,19 +9,11 @@ class ParentMode(StrEnum):
PARAGRAPH = "paragraph"
PreProcessingRuleID = Annotated[
str,
WithJsonSchema(
{
"enum": ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"],
"type": "string",
}
),
]
type PreProcessingRuleKey = Literal["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
class PreProcessingRule(BaseModel):
id: PreProcessingRuleID = Field(description="Rule identifier.")
id: PreProcessingRuleKey = Field(description="Rule identifier.")
enabled: bool = Field(description="Whether this preprocessing rule is enabled.")
api/core/rag/extractor/entity/extract_setting.py
+23 -9
View File
@@ -1,7 +1,24 @@
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, Field
from models.dataset import Document
from models.model import UploadFile
class UploadFileExtractionInput(BaseModel):
"""Detached values required to extract one tenant-owned upload."""
id: str
tenant_id: str
key: str
created_by: str
model_config = ConfigDict(from_attributes=True)
class StoredDocumentExtractionInput(BaseModel):
"""Detached document values needed to record a source's last edited time."""
id: str
tenant_id: str
dataset_id: str
data_source_info_dict: dict[str, object] | None
model_config = ConfigDict(from_attributes=True)
class NotionInfo(BaseModel):
@@ -13,9 +30,9 @@ class NotionInfo(BaseModel):
notion_workspace_id: str | None = ""
notion_obj_id: str
notion_page_type: str
document: Document | None = None
document: StoredDocumentExtractionInput | None = None
tenant_id: str
model_config = ConfigDict(arbitrary_types_allowed=True)
notion_access_token: str | None = Field(default=None, exclude=True, repr=False)
class WebsiteInfo(BaseModel):
@@ -23,8 +40,6 @@ class WebsiteInfo(BaseModel):
website import info.
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
provider: str
job_id: str
url: str
@@ -39,8 +54,7 @@ class ExtractSetting(BaseModel):
"""
datasource_type: str
upload_file: UploadFile | None = None
upload_file: UploadFileExtractionInput | None = None
notion_info: NotionInfo | None = None
website_info: WebsiteInfo | None = None
document_model: str | None = None
model_config = ConfigDict(arbitrary_types_allowed=True)
api/core/rag/extractor/extract_processor.py
+20 -2
View File
@@ -7,10 +7,11 @@ from urllib.parse import unquote
from sqlalchemy.orm import Session
from configs import dify_config
from core.db.session_factory import get_session_maker
from core.file import remote_fetcher
from core.rag.extractor.csv_extractor import CSVExtractor
from core.rag.extractor.entity.datasource_type import DatasourceType
from core.rag.extractor.entity.extract_setting import ExtractSetting
from core.rag.extractor.entity.extract_setting import ExtractSetting, UploadFileExtractionInput
from core.rag.extractor.excel_extractor import ExcelExtractor
from core.rag.extractor.extractor_base import BaseExtractor
from core.rag.extractor.firecrawl.firecrawl_web_extractor import FirecrawlWebExtractor
@@ -59,7 +60,9 @@ class ExtractProcessor:
cls, upload_file: UploadFile, return_text: bool = False, is_automatic: bool = False
) -> list[Document] | str:
extract_setting = ExtractSetting(
datasource_type=DatasourceType.FILE, upload_file=upload_file, document_model="text_model"
datasource_type=DatasourceType.FILE,
upload_file=UploadFileExtractionInput.model_validate(upload_file),
document_model="text_model",
)
if return_text:
delimiter = "\n"
@@ -223,17 +226,29 @@ class ExtractProcessor:
extractor = TextExtractor(file_path, autodetect_encoding=True)
return extractor.extract()
elif extract_setting.datasource_type == DatasourceType.NOTION:
from extensions.application_services.data_sources import build_data_source_credentials
assert extract_setting.notion_info is not None, "notion_info is required"
notion_info = extract_setting.notion_info
extractor = NotionExtractor(
notion_workspace_id=extract_setting.notion_info.notion_workspace_id or "",
notion_obj_id=extract_setting.notion_info.notion_obj_id,
notion_page_type=extract_setting.notion_info.notion_page_type,
document_model=extract_setting.notion_info.document,
tenant_id=extract_setting.notion_info.tenant_id,
notion_access_token=extract_setting.notion_info.notion_access_token,
credential_id=extract_setting.notion_info.credential_id,
notion_token_loader=lambda: build_data_source_credentials(
database_client=get_session_maker()
).providers.get_stored_notion_access_token(
tenant_id=notion_info.tenant_id, credential_id=notion_info.credential_id
),
)
return extractor.extract()
elif extract_setting.datasource_type == DatasourceType.WEBSITE:
from extensions.application_services.data_sources import build_website_service
website_service = build_website_service(database_client=get_session_maker())
assert extract_setting.website_info is not None, "website_info is required"
match extract_setting.website_info.provider:
case "firecrawl":
@@ -243,6 +258,7 @@ class ExtractProcessor:
tenant_id=extract_setting.website_info.tenant_id,
mode=extract_setting.website_info.mode,
only_main_content=extract_setting.website_info.only_main_content,
website_service=website_service,
)
return extractor.extract()
case "watercrawl":
@@ -252,6 +268,7 @@ class ExtractProcessor:
tenant_id=extract_setting.website_info.tenant_id,
mode=extract_setting.website_info.mode,
only_main_content=extract_setting.website_info.only_main_content,
website_service=website_service,
)
return extractor.extract()
case "jinareader":
@@ -261,6 +278,7 @@ class ExtractProcessor:
tenant_id=extract_setting.website_info.tenant_id,
mode=extract_setting.website_info.mode,
only_main_content=extract_setting.website_info.only_main_content,
website_service=website_service,
)
return extractor.extract()
case _:
api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py
+6 -3
View File
@@ -2,7 +2,7 @@ from typing import override
from core.rag.extractor.extractor_base import BaseExtractor
from core.rag.models.document import Document
from services.website_service import WebsiteService
from services.data_source.website_service import WebsiteService
class FirecrawlWebExtractor(BaseExtractor):
@@ -24,8 +24,11 @@ class FirecrawlWebExtractor(BaseExtractor):
tenant_id: str,
mode: str = "crawl",
only_main_content: bool = True,
*,
website_service: WebsiteService,
):
"""Initialize with url, api_key, base_url and mode."""
self._website_service = website_service
self._url = url
self.job_id = job_id
self.tenant_id = tenant_id
@@ -37,7 +40,7 @@ class FirecrawlWebExtractor(BaseExtractor):
"""Extract content from the URL."""
documents = []
if self.mode == "crawl":
crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "firecrawl", self._url, self.tenant_id)
crawl_data = self._website_service.get_crawl_url_data(self.job_id, "firecrawl", self._url, self.tenant_id)
if crawl_data is None:
return []
document = Document(
@@ -50,7 +53,7 @@ class FirecrawlWebExtractor(BaseExtractor):
)
documents.append(document)
elif self.mode == "scrape":
scrape_data = WebsiteService.get_scrape_url_data(
scrape_data = self._website_service.get_scrape_url_data(
"firecrawl", self._url, self.tenant_id, self.only_main_content
)
api/core/rag/extractor/jina_reader_extractor.py
+5 -2
View File
@@ -2,7 +2,7 @@ from typing import override
from core.rag.extractor.extractor_base import BaseExtractor
from core.rag.models.document import Document
from services.website_service import WebsiteService
from services.data_source.website_service import WebsiteService
class JinaReaderWebExtractor(BaseExtractor):
@@ -17,8 +17,11 @@ class JinaReaderWebExtractor(BaseExtractor):
tenant_id: str,
mode: str = "crawl",
only_main_content: bool = False,
*,
website_service: WebsiteService,
):
"""Initialize with url, api_key, base_url and mode."""
self._website_service = website_service
self._url = url
self.job_id = job_id
self.tenant_id = tenant_id
@@ -30,7 +33,7 @@ class JinaReaderWebExtractor(BaseExtractor):
"""Extract content from the URL."""
documents = []
if self.mode == "crawl":
crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "jinareader", self._url, self.tenant_id)
crawl_data = self._website_service.get_crawl_url_data(self.job_id, "jinareader", self._url, self.tenant_id)
if crawl_data is None:
return []
document = Document(
api/core/rag/extractor/notion_extractor.py
+12 -40
View File
@@ -1,17 +1,17 @@
import json
import logging
import operator
from collections.abc import Callable
from typing import Any, cast, override
import httpx
from sqlalchemy import update
from configs import dify_config
from core.rag.extractor.entity.extract_setting import StoredDocumentExtractionInput
from core.rag.extractor.extractor_base import BaseExtractor
from core.rag.models.document import Document
from extensions.ext_database import db
from models.dataset import Document as DocumentModel
from services.datasource_provider_service import DatasourceProviderService
logger = logging.getLogger(__name__)
@@ -41,36 +41,21 @@ class NotionExtractor(BaseExtractor):
notion_obj_id: str,
notion_page_type: str,
tenant_id: str,
document_model: DocumentModel | None = None,
document_model: StoredDocumentExtractionInput | None = None,
notion_access_token: str | None = None,
credential_id: str | None = None,
*,
notion_token_loader: Callable[[], str],
):
self._notion_access_token = None
self._document_model = document_model
self._notion_workspace_id = notion_workspace_id
self._notion_obj_id = notion_obj_id
self._notion_page_type = notion_page_type
self._credential_id = credential_id
if notion_access_token:
self._notion_access_token = notion_access_token
else:
try:
self._notion_access_token = self._get_access_token(tenant_id, self._credential_id)
except Exception as e:
logger.warning(
(
"Failed to get Notion access token from datasource credentials: %s, "
"falling back to environment variable NOTION_INTEGRATION_TOKEN"
),
e,
)
integration_token = dify_config.NOTION_INTEGRATION_TOKEN
if integration_token is None:
raise ValueError(
"Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`."
) from e
self._notion_access_token = integration_token
self._notion_access_token = notion_token_loader()
@override
def extract(self) -> list[Document]:
@@ -344,7 +329,7 @@ class NotionExtractor(BaseExtractor):
result_lines = "\n".join(result_lines_arr)
return result_lines
def update_last_edited_time(self, document_model: DocumentModel | None):
def update_last_edited_time(self, document_model: StoredDocumentExtractionInput | None):
if not document_model:
return
@@ -355,7 +340,11 @@ class NotionExtractor(BaseExtractor):
db.session.execute(
update(DocumentModel)
.where(DocumentModel.id == document_model.id)
.where(
DocumentModel.id == document_model.id,
DocumentModel.tenant_id == document_model.tenant_id,
DocumentModel.dataset_id == document_model.dataset_id,
)
.values({DocumentModel.data_source_info: json.dumps(data_source_info)})
)
db.session.commit()
@@ -385,20 +374,3 @@ class NotionExtractor(BaseExtractor):
data = res.json()
return cast(str, data["last_edited_time"])
@classmethod
def _get_access_token(cls, tenant_id: str, credential_id: str | None) -> str:
# get credential from tenant_id and credential_id
if not credential_id:
raise Exception(f"No credential id found for tenant {tenant_id}")
datasource_provider_service = DatasourceProviderService()
credential = datasource_provider_service.get_datasource_credentials(
tenant_id=tenant_id,
credential_id=credential_id,
provider="notion_datasource",
plugin_id="langgenius/notion_datasource",
)
if not credential:
raise Exception(f"No notion credential found for tenant {tenant_id} and credential {credential_id}")
return cast(str, credential["integration_secret"])
api/core/rag/extractor/watercrawl/extractor.py
+6 -3
View File
@@ -2,7 +2,7 @@ from typing import override
from core.rag.extractor.extractor_base import BaseExtractor
from core.rag.models.document import Document
from services.website_service import WebsiteService
from services.data_source.website_service import WebsiteService
class WaterCrawlWebExtractor(BaseExtractor):
@@ -25,8 +25,11 @@ class WaterCrawlWebExtractor(BaseExtractor):
tenant_id: str,
mode: str = "crawl",
only_main_content: bool = True,
*,
website_service: WebsiteService,
):
"""Initialize with url, api_key, base_url and mode."""
self._website_service = website_service
self._url = url
self.job_id = job_id
self.tenant_id = tenant_id
@@ -38,7 +41,7 @@ class WaterCrawlWebExtractor(BaseExtractor):
"""Extract content from the URL."""
documents = []
if self.mode == "crawl":
crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "watercrawl", self._url, self.tenant_id)
crawl_data = self._website_service.get_crawl_url_data(self.job_id, "watercrawl", self._url, self.tenant_id)
if crawl_data is None:
return []
document = Document(
@@ -51,7 +54,7 @@ class WaterCrawlWebExtractor(BaseExtractor):
)
documents.append(document)
elif self.mode == "scrape":
scrape_data = WebsiteService.get_scrape_url_data(
scrape_data = self._website_service.get_scrape_url_data(
"watercrawl", self._url, self.tenant_id, self.only_main_content
)
api/core/rag/index_processor/index_processor_base.py
+29 -5
View File
@@ -147,7 +147,12 @@ class BaseIndexProcessor(ABC):
return character_splitter
def _get_content_files(
self, document: Document, current_user: Account | None = None, *, session: Session
self,
document: Document,
current_user: Account | None = None,
*,
tenant_id: str,
session: Session,
) -> list[AttachmentDocument]:
"""
Get the content files from the document.
@@ -185,7 +190,12 @@ class BaseIndexProcessor(ABC):
if match:
if current_user:
tool_file_id = match.group(1)
upload_file_id = self._download_tool_file(tool_file_id, current_user, session=session)
upload_file_id = self._download_tool_file(
tool_file_id,
current_user,
tenant_id=tenant_id,
session=session,
)
if upload_file_id:
upload_file_id_list.append(upload_file_id)
continue
@@ -208,7 +218,12 @@ class BaseIndexProcessor(ABC):
# Get unique IDs for database query
unique_upload_file_ids = list(set(upload_file_id_list))
upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(unique_upload_file_ids))).all()
upload_files = session.scalars(
select(UploadFile).where(
UploadFile.tenant_id == tenant_id,
UploadFile.id.in_(unique_upload_file_ids),
)
).all()
# Create a mapping from ID to UploadFile for quick lookup
upload_file_map = {upload_file.id: upload_file for upload_file in upload_files}
@@ -314,13 +329,22 @@ class BaseIndexProcessor(ABC):
logging.warning("Unexpected error downloading image from %s", image_url, exc_info=True)
return None
def _download_tool_file(self, tool_file_id: str, current_user: Account, *, session: Session) -> str | None:
def _download_tool_file(
self,
tool_file_id: str,
current_user: Account,
*,
tenant_id: str,
session: Session,
) -> str | None:
"""
Download the tool file from the ID.
"""
from services.file_service import FileService
tool_file = session.get(ToolFile, tool_file_id)
tool_file = session.scalar(
select(ToolFile).where(ToolFile.id == tool_file_id, ToolFile.tenant_id == tenant_id).limit(1)
)
if not tool_file:
return None
blob = storage.load_once(tool_file.file_key)
api/core/rag/index_processor/processor/paragraph_index_processor.py
+43 -16
View File
@@ -3,6 +3,7 @@
import logging
import re
import uuid
from collections.abc import Callable
from typing import Any, TypedDict, cast, override
from sqlalchemy import select
@@ -46,7 +47,7 @@ from models.account import Account
from models.dataset import Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding
from models.dataset import Document as DatasetDocument
from services.account_service import AccountService
from services.summary_index_service import SummaryIndexService
from services.knowledge.summaries.adapters import SummaryIndexAdapter
logger = logging.getLogger(__name__)
@@ -113,7 +114,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
document_node.metadata["doc_id"] = doc_id
document_node.metadata["doc_hash"] = hash
multimodal_documents = (
self._get_content_files(document_node, current_user, session=session)
self._get_content_files(
document_node,
current_user,
tenant_id=kwargs["tenant_id"],
session=session,
)
if document_node.metadata
else None
)
@@ -172,10 +178,10 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
).all()
segment_ids = [segment.id for segment in segments]
if segment_ids:
SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids, session=session)
SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids, session=session)
else:
# Delete all summaries for the dataset
SummaryIndexService.delete_summaries_for_segments(dataset, None, session=session)
SummaryIndexAdapter.delete_summaries_for_segments(dataset, None, session=session)
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
vector = Vector(dataset, session=session)
@@ -204,7 +210,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
"doc_hash": helper.generate_text_hash(content),
}
doc = Document(page_content=content, metadata=metadata)
attachments = self._get_content_files(doc, session=session)
attachments = self._get_content_files(doc, tenant_id=dataset.tenant_id, session=session)
if attachments:
doc.attachments = attachments
all_multimodal_documents.extend(attachments)
@@ -240,7 +246,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
account = AccountService.load_user(document.created_by, account_session)
if not account:
raise ValueError("Invalid account")
doc.attachments = self._get_content_files(doc, current_user=account, session=session)
doc.attachments = self._get_content_files(
doc,
current_user=account,
tenant_id=dataset.tenant_id,
session=session,
)
if doc.attachments:
all_multimodal_documents.extend(doc.attachments)
documents.append(doc)
@@ -378,7 +389,6 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
return preview_texts
@staticmethod
@with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
def generate_summary(
tenant_id: str,
text: str,
@@ -404,6 +414,30 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
Returns:
Tuple of (summary_content, llm_usage) where llm_usage is LLMUsage object
"""
def load_images() -> list[File]:
images = (
ParagraphIndexProcessor._extract_images_from_segment_attachments(tenant_id, segment_id, session)
if segment_id
else []
)
return images or ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, session)
return ParagraphIndexProcessor.generate_summary_from_inputs(
tenant_id, text, summary_index_setting, document_language=document_language, image_loader=load_images
)
@staticmethod
@with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING)
def generate_summary_from_inputs(
tenant_id: str,
text: str,
summary_index_setting: SummaryIndexSettingDict | None,
*,
document_language: str | None,
image_loader: Callable[[], list[File]],
) -> tuple[str, LLMUsage]:
"""Generate a summary after an explicit loader materializes any required images."""
if not summary_index_setting or not summary_index_setting.get("enable"):
raise ValueError("summary_index_setting is required and must be enabled to generate summary.")
@@ -447,15 +481,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
# Extract images if model supports vision
image_files = []
if supports_vision:
# First, try to get images from SegmentAttachmentBinding (preferred method)
if segment_id:
image_files = ParagraphIndexProcessor._extract_images_from_segment_attachments(
tenant_id, segment_id, session
)
# If no images from attachments, fall back to extracting from text
if not image_files:
image_files = ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, session)
image_files = image_loader()
# Build prompt messages
prompt_messages = []
@@ -603,6 +629,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
.where(
SegmentAttachmentBinding.segment_id == segment_id,
SegmentAttachmentBinding.tenant_id == tenant_id,
UploadFile.tenant_id == tenant_id,
)
).all()
api/core/rag/index_processor/processor/parent_child_index_processor.py
+23 -9
View File
@@ -29,7 +29,7 @@ from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegm
from models.dataset import Document as DatasetDocument
from models.enums import ProcessRuleMode
from services.account_service import AccountService
from services.summary_index_service import SummaryIndexService
from services.knowledge.summaries.adapters import SummaryIndexAdapter
logger = logging.getLogger(__name__)
@@ -99,11 +99,16 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
page_content = page_content
if len(page_content) > 0:
document_node.page_content = page_content
multimodel_documents = self._get_content_files(document_node, current_user, session=session)
multimodel_documents = self._get_content_files(
document_node,
current_user,
tenant_id=kwargs["tenant_id"],
session=session,
)
if multimodel_documents:
document_node.attachments = multimodel_documents
# parse document to child nodes
child_nodes = self._split_child_nodes(
child_nodes = self.split_child_nodes(
document_node, rules, process_rule.get("mode"), kwargs.get("embedding_model_instance")
)
document_node.children = child_nodes
@@ -112,11 +117,15 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
elif rules.parent_mode == ParentMode.FULL_DOC:
page_content = "\n".join([document.page_content for document in documents])
document = Document(page_content=page_content, metadata=documents[0].metadata)
multimodel_documents = self._get_content_files(document, session=session)
multimodel_documents = self._get_content_files(
document,
tenant_id=kwargs["tenant_id"],
session=session,
)
if multimodel_documents:
document.attachments = multimodel_documents
# parse document to child nodes
child_nodes = self._split_child_nodes(
child_nodes = self.split_child_nodes(
document, rules, process_rule.get("mode"), kwargs.get("embedding_model_instance")
)
if kwargs.get("preview"):
@@ -176,10 +185,10 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
).all()
segment_ids = [segment.id for segment in segments]
if segment_ids:
SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids, session=session)
SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids, session=session)
else:
# Delete all summaries for the dataset
SummaryIndexService.delete_summaries_for_segments(dataset, None, session=session)
SummaryIndexAdapter.delete_summaries_for_segments(dataset, None, session=session)
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
delete_child_chunks = kwargs.get("delete_child_chunks") or False
@@ -227,7 +236,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
)
session.flush()
def _split_child_nodes(
def split_child_nodes(
self,
document_node: Document,
rules: Rule,
@@ -302,7 +311,12 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
account = AccountService.load_user(document.created_by, account_session)
if not account:
raise ValueError("Invalid account")
doc.attachments = self._get_content_files(doc, current_user=account, session=session)
doc.attachments = self._get_content_files(
doc,
current_user=account,
tenant_id=dataset.tenant_id,
session=session,
)
documents.append(doc)
if documents:
token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
api/core/rag/index_processor/processor/qa_index_processor.py
+3 -3
View File
@@ -29,7 +29,7 @@ from libs import helper
from models.account import Account
from models.dataset import Dataset, DocumentSegment
from models.dataset import Document as DatasetDocument
from services.summary_index_service import SummaryIndexService
from services.knowledge.summaries.adapters import SummaryIndexAdapter
logger = logging.getLogger(__name__)
@@ -180,10 +180,10 @@ class QAIndexProcessor(BaseIndexProcessor):
).all()
segment_ids = [segment.id for segment in segments]
if segment_ids:
SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids, session=session)
SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids, session=session)
else:
# Delete all summaries for the dataset
SummaryIndexService.delete_summaries_for_segments(dataset, None, session=session)
SummaryIndexAdapter.delete_summaries_for_segments(dataset, None, session=session)
vector = Vector(dataset, session=session)
if node_ids:
api/core/rag/retrieval/dataset_retrieval.py
+45 -35
View File
@@ -28,7 +28,6 @@ from core.app.entities.app_invoke_entities import (
ModelConfigWithCredentialsEntity,
get_credit_usage_app_type,
)
from core.app.file_access import grant_retriever_segment_access, grant_upload_file_access
from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
from core.db.session_factory import session_factory
from core.entities.agent_entities import PlanningStrategy
@@ -83,7 +82,6 @@ from graphon.model_runtime.entities.model_entities import ModelFeature, ModelTyp
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
from libs.helper import parse_uuid_str_or_none
from libs.json_in_md_parser import parse_and_check_json_markdown
from models import UploadFile
from models.dataset import (
ChildChunk,
Dataset,
@@ -96,8 +94,11 @@ from models.dataset import (
from models.dataset import Document as DatasetDocument
from models.dataset import Document as DocumentModel
from models.enums import CreatorUserRole, DatasetQuerySource
from services.external_knowledge_service import ExternalDatasetService
from repositories.knowledge.dataset_read_repository import get_dataset_available_document_count
from repositories.knowledge.segment_read_adapter import sign_segment_content
from services.feature_service import FeatureService
from services.knowledge.external.service import ExternalDatasetService
from services.knowledge.retrieval.attachments import authorize_retrieved_segment
default_retrieval_model: DefaultRetrievalModelDict = {
"search_method": RetrievalMethod.SEMANTIC_SEARCH,
@@ -310,25 +311,36 @@ class DatasetRetrieval:
retrieval_resource_list.append(source)
# deal with dify documents
if dify_documents:
with Session(bind=session.get_bind()) as format_session:
records = RetrievalService.format_retrieval_documents(format_session, dify_documents)
dataset_ids = [i.segment.dataset_id for i in records]
document_ids = [i.segment.document_id for i in records]
# Materialize authorized, signed response values before closing the owned session.
with Session(bind=session.get_bind()) as retrieval_session:
records = RetrievalService.format_retrieval_documents(retrieval_session, dify_documents)
dataset_ids = [i.segment.dataset_id for i in records]
document_ids = [i.segment.document_id for i in records]
datasets = retrieval_session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all()
documents = retrieval_session.scalars(
select(DatasetDocument).where(DatasetDocument.id.in_(document_ids))
).all()
with session_factory.create_session() as session:
datasets = session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all()
documents = session.scalars(select(DatasetDocument).where(DatasetDocument.id.in_(document_ids))).all()
dataset_map = {i.id: i for i in datasets}
document_map = {i.id: i for i in documents}
dataset_map = {i.id: i for i in datasets}
document_map = {i.id: i for i in documents}
if records:
for record in records:
segment = record.segment
if (
authorize_retrieved_segment(
segment,
tenant_id=request.tenant_id,
dataset_ids=available_datasets_ids,
session=retrieval_session,
)
is None
):
continue
dataset = dataset_map.get(segment.dataset_id)
document = document_map.get(segment.document_id)
if dataset and document:
content = sign_segment_content(segment, session=retrieval_session)
source = Source(
metadata=SourceMetadata(
source="knowledge",
@@ -358,15 +370,14 @@ class DatasetRetrieval:
),
title=document.name,
files=list(record.files) if record.files else None,
content=segment.get_sign_content(),
content=content,
)
if segment.answer:
source.content = f"question:{segment.get_sign_content()} \nanswer:{segment.answer}"
source.content = f"question:{content} \nanswer:{segment.answer}"
if record.summary:
source.summary = record.summary
grant_retriever_segment_access([str(segment.id)])
retrieval_resource_list.append(source)
if retrieval_resource_list:
@@ -533,13 +544,22 @@ class DatasetRetrieval:
with Session(bind=session.get_bind()) as format_session:
records = RetrievalService.format_retrieval_documents(format_session, dify_documents)
if records:
authorized_records = []
for record in records:
segment = record.segment
attachments = authorize_retrieved_segment(
segment, tenant_id=tenant_id, dataset_ids=available_datasets_ids, session=session
)
if attachments is None:
continue
authorized_records.append(record)
# Build content: if summary exists, add it before the segment content
if segment.answer:
segment_content = f"question:{segment.get_sign_content()} answer:{segment.answer}"
segment_content = (
f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}"
)
else:
segment_content = segment.get_sign_content()
segment_content = sign_segment_content(segment, session=session)
# If summary exists, prepend it to the content
if record.summary:
@@ -554,18 +574,8 @@ class DatasetRetrieval:
)
)
if vision_enabled:
attachments_with_bindings = session.execute(
select(SegmentAttachmentBinding, UploadFile)
.join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
.where(
SegmentAttachmentBinding.segment_id == segment.id,
)
).all()
if attachments_with_bindings:
grant_upload_file_access(
str(upload_file.id) for _, upload_file in attachments_with_bindings
)
for _, upload_file in attachments_with_bindings:
if attachments:
for upload_file in attachments:
attachment_info = File(
file_id=upload_file.id,
filename=upload_file.name,
@@ -583,8 +593,8 @@ class DatasetRetrieval:
)
context_files.append(attachment_info)
if show_retrieve_source:
dataset_ids = [record.segment.dataset_id for record in records]
document_ids = [record.segment.document_id for record in records]
dataset_ids = [record.segment.dataset_id for record in authorized_records]
document_ids = [record.segment.document_id for record in authorized_records]
dataset_document_stmt = select(DatasetDocument).where(
DatasetDocument.id.in_(document_ids),
DatasetDocument.enabled == True,
@@ -597,7 +607,7 @@ class DatasetRetrieval:
datasets = session.execute(dataset_stmt).scalars().all() # type: ignore
dataset_map = {i.id: i for i in datasets}
document_map = {i.id: i for i in documents}
for record in records:
for record in authorized_records:
segment = record.segment
dataset_item = dataset_map.get(segment.dataset_id)
document_item = document_map.get(segment.document_id)
@@ -1333,7 +1343,7 @@ class DatasetRetrieval:
if (
dataset
and dataset.provider != "external"
and dataset.get_total_available_documents(session=session) == 0
and get_dataset_available_document_count(dataset, session=session) == 0
):
continue
api/core/rag/summary_index/summary_index.py
+2 -2
View File
@@ -8,7 +8,7 @@ from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict
from extensions.otel import propagate_context
from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary
from services.summary_index_service import SummaryIndexService
from services.knowledge.summaries.adapters import SummaryIndexAdapter
from tasks.generate_summary_index_task import generate_summary_index_task
logger = logging.getLogger(__name__)
@@ -82,7 +82,7 @@ class SummaryIndex:
if segment is None:
return
try:
SummaryIndexService.generate_and_vectorize_summary(
SummaryIndexAdapter.generate_and_vectorize_summary(
segment, dataset, summary_index_setting, session=session
)
except Exception:
api/core/tools/builtin_tool/providers/audio/tools/asr.py
+1 -1
View File
@@ -14,7 +14,7 @@ from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
from graphon.file import FileType
from graphon.file.file_manager import download
from graphon.model_runtime.entities.model_entities import ModelType
from services.model_provider_service import ModelProviderService
from services.model_provider.service import ModelProviderService
class ASRTool(BuiltinTool):
api/core/tools/builtin_tool/providers/audio/tools/tts.py
+1 -1
View File
@@ -12,7 +12,7 @@ from core.tools.builtin_tool.tool import BuiltinTool
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
from graphon.model_runtime.entities.model_entities import ModelPropertyKey, ModelType
from services.model_provider_service import ModelProviderService
from services.model_provider.service import ModelProviderService
class TTSTool(BuiltinTool):
api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py
+5 -2
View File
@@ -18,6 +18,7 @@ from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import Datas
from extensions.ext_database import db
from graphon.model_runtime.entities.model_entities import ModelType
from models.dataset import Dataset, Document, DocumentSegment
from repositories.knowledge.segment_read_adapter import sign_segment_content
default_retrieval_model: DefaultRetrievalModelDict = {
"search_method": RetrievalMethod.SEMANTIC_SEARCH,
@@ -105,9 +106,11 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
)
for segment in sorted_segments:
if segment.answer:
document_context_list.append(f"question:{segment.get_sign_content()} answer:{segment.answer}")
document_context_list.append(
f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}"
)
else:
document_context_list.append(segment.get_sign_content())
document_context_list.append(sign_segment_content(segment, session=session))
if self.return_resource:
context_list: list[RetrievalSourceMetadata] = []
resource_number = 1
api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
+6 -3
View File
@@ -14,7 +14,8 @@ from core.rag.retrieval.retrieval_methods import RetrievalMethod
from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
from models.dataset import Dataset
from models.dataset import Document as DatasetDocument
from services.external_knowledge_service import ExternalDatasetService
from repositories.knowledge.segment_read_adapter import sign_segment_content
from services.knowledge.external.service import ExternalDatasetService
default_retrieval_model: DefaultRetrievalModelDict = {
"search_method": RetrievalMethod.SEMANTIC_SEARCH,
@@ -175,9 +176,11 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
segment = record.segment
# Build content: if summary exists, add it before the segment content
if segment.answer:
segment_content = f"question:{segment.get_sign_content()} answer:{segment.answer}"
segment_content = (
f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}"
)
else:
segment_content = segment.get_sign_content()
segment_content = sign_segment_content(segment, session=session)
# If summary exists, prepend it to the content
if record.summary:
api/core/workflow/node_factory.py
+7 -1
View File
@@ -11,7 +11,7 @@ from configs import dify_config
from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext
from core.app.llm.model_access import build_dify_model_access, fetch_model_config
from core.credit_usage import created_by_from_app_type
from core.db.session_factory import session_factory
from core.db.session_factory import get_session_maker, session_factory
from core.file import remote_fetcher
from core.helper.code_executor.code_executor import (
CodeExecutionError,
@@ -55,6 +55,7 @@ from core.workflow.nodes.human_input.callback import DifyHITLCallback
from core.workflow.nodes.human_input.entities import HumanInputNodeData as DifyHumanInputNodeData
from core.workflow.system_variables import SystemVariableKey, get_system_text, system_variable_selector
from core.workflow.template_rendering import CodeExecutorJinja2TemplateRenderer
from extensions.application_services.data_sources import build_data_source_credentials
from graphon.entities.base_node_data import BaseNodeData
from graphon.entities.graph_config import NodeConfigDict, NodeConfigDictAdapter
from graphon.enums import BuiltinNodeTypes, NodeType
@@ -431,6 +432,11 @@ class DifyNodeFactory(NodeFactory):
resolved_node_data = self._resolve_llm_model_reference(cast(LLMNodeData, resolved_node_data))
node: Node | None = None
node_init_kwargs_factories: Mapping[NodeType, Callable[[], dict[str, object]]] = {
BuiltinNodeTypes.DATASOURCE: lambda: {
"datasource_credentials": build_data_source_credentials(
database_client=get_session_maker()
).providers.get_datasource_credentials,
},
BuiltinNodeTypes.CODE: lambda: {
"code_executor": self._code_executor,
"code_limits": self._code_limits,
api/core/workflow/nodes/datasource/datasource_node.py
+6 -4
View File
@@ -1,4 +1,4 @@
from collections.abc import Generator, Mapping, Sequence
from collections.abc import Callable, Generator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, override
from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext
@@ -40,6 +40,7 @@ class DatasourceNode(Node[DatasourceNodeData]):
*,
graph_init_params: "GraphInitParams",
graph_runtime_state: "GraphRuntimeState",
datasource_credentials: Callable[[str, str, str, str | None], dict[str, Any]],
) -> None:
super().__init__(
node_id=node_id,
@@ -48,6 +49,7 @@ class DatasourceNode(Node[DatasourceNodeData]):
graph_runtime_state=graph_runtime_state,
)
self.datasource_manager = DatasourceManager
self._datasource_credentials = datasource_credentials
@override
def populate_start_event(self, event) -> None:
@@ -117,9 +119,9 @@ class DatasourceNode(Node[DatasourceNodeData]):
datasource_type=datasource_type.value,
provider_id=provider_id,
tenant_id=dify_ctx.tenant_id,
provider=node_data.provider_name,
plugin_id=node_data.plugin_id,
credential_id=credential_id,
credentials=self._datasource_credentials(
dify_ctx.tenant_id, node_data.provider_name, node_data.plugin_id, credential_id
),
parameters_for_log=parameters_for_log,
datasource_info=datasource_info,
variable_pool=variable_pool,
api/core/workflow/nodes/datasource/protocols.py
-35
View File
@@ -1,35 +0,0 @@
from collections.abc import Generator
from typing import Any, Protocol
from graphon.file import File
from graphon.node_events import StreamChunkEvent, StreamCompletedEvent
from .entities import DatasourceParameter, OnlineDriveDownloadFileParam
class DatasourceManagerProtocol(Protocol):
@classmethod
def get_icon_url(cls, provider_id: str, tenant_id: str, datasource_name: str, datasource_type: str) -> str: ...
@classmethod
def stream_node_events(
cls,
*,
node_id: str,
user_id: str,
datasource_name: str,
datasource_type: str,
provider_id: str,
tenant_id: str,
provider: str,
plugin_id: str,
credential_id: str,
parameters_for_log: dict[str, Any],
datasource_info: dict[str, Any],
variable_pool: Any,
datasource_param: DatasourceParameter | None = None,
online_drive_request: OnlineDriveDownloadFileParam | None = None,
) -> Generator[StreamChunkEvent | StreamCompletedEvent, None, None]: ...
@classmethod
def get_upload_file_by_id(cls, file_id: str, tenant_id: str) -> File: ...
api/events/event_handlers/clean_when_dataset_deleted.py
+2 -1
View File
@@ -1,13 +1,14 @@
from events.dataset_event import dataset_was_deleted
from extensions.ext_database import db
from models import Dataset
from repositories.knowledge.dataset_read_repository import get_dataset_doc_form
from tasks.clean_dataset_task import clean_dataset_task
@dataset_was_deleted.connect
def handle(sender: Dataset, **kwargs):
dataset = sender
doc_form = dataset.get_doc_form(session=db.session())
doc_form = get_dataset_doc_form(dataset, session=db.session())
if not doc_form or not dataset.indexing_technique:
return
clean_dataset_task.delay(
api/events/event_handlers/create_document_index.py
+8 -5
View File
@@ -6,11 +6,13 @@ from sqlalchemy import select
from werkzeug.exceptions import NotFound
from core.db.session_factory import session_factory
from core.indexing_runner import DocumentIsPausedError, IndexingRunner
from events.document_index_event import document_index_created
from libs.datetime_utils import naive_utc_now
from models.dataset import Document
from models.enums import IndexingStatus
from services.knowledge.indexing.adapters.execution import build_document_indexing_service
from services.knowledge.indexing.errors import DocumentIsPausedError
from services.knowledge.resource_scope import DatasetRef
logger = logging.getLogger(__name__)
@@ -21,7 +23,9 @@ def handle(sender, **kwargs):
document_ids = kwargs.get("document_ids", [])
start_at = time.perf_counter()
try:
indexing_runner = IndexingRunner(enforce_vector_space_admission=True)
indexing_service = build_document_indexing_service(
session_factory=session_factory.get_session_maker(), enforce_vector_space_admission=True
)
with session_factory.create_session() as session:
documents = []
for document_id in document_ids:
@@ -41,11 +45,10 @@ def handle(sender, **kwargs):
document.processing_started_at = naive_utc_now()
documents.append(document)
session.add(document)
document_refs = [DatasetRef(doc.tenant_id, doc.dataset_id).document(doc.id) for doc in documents]
# Persist the status transition before extraction and indexing begin.
session.commit()
indexing_runner.run(documents, session)
session.commit()
indexing_service.run(document_refs)
end_at = time.perf_counter()
logger.info(click.style(f"Processed dataset: {dataset_id} latency: {end_at - start_at}", fg="green"))
except DocumentIsPausedError as ex:
api/extensions/application_services/data_sources.py
+124
View File
@@ -0,0 +1,124 @@
"""Composition root for data-source services and shared credential gateways."""
from collections.abc import Mapping
from dataclasses import dataclass
from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config
from core.datasource.datasource_manager import DatasourceManager
from core.helper.ssrf_proxy import ssrf_proxy
from core.plugin.impl.datasource import PluginDatasourceManager
from core.plugin.impl.oauth import OAuthHandler
from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository
from repositories.data_source.oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository
from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository
from services.data_source.auth.api_key_gateways import (
ProviderApiKeyAuthCredentialValidator,
TenantApiKeyAuthCredentialEncryptor,
)
from services.data_source.auth.api_key_service import DataSourceApiKeyAuthService
from services.data_source.binding_application_service import DataSourceBindingApplicationService
from services.data_source.credential_adapters import (
OAuthDatasourceCredentialRefresher,
PluginDatasourceCredentialCodec,
PluginDatasourceOAuthClientResolver,
)
from services.data_source.credential_gateway import (
ActorAwareDatasourceCredentialGateway,
TrustedStoredDatasourceCredentialGateway,
)
from services.data_source.notion_import_adapters import PluginNotionSourceGateway
from services.data_source.notion_import_application_service import NotionImportApplicationService
from services.data_source.notion_oauth_gateway import NotionDataSourceGateway
from services.data_source.oauth_service import DataSourceOAuthService, InvalidDataSourceOAuthProviderError
from services.data_source.provider_service import DatasourceProviderService
from services.data_source.website_service import WebsiteService
from services.knowledge.dataset_access import DatasetAccessService
@dataclass(frozen=True, slots=True)
class DataSourceServices:
providers: DatasourceProviderService
api_key_auth: DataSourceApiKeyAuthService
oauth: Mapping[str, DataSourceOAuthService]
bindings: DataSourceBindingApplicationService
notion_imports: NotionImportApplicationService
def resolve_oauth(self, provider: str) -> DataSourceOAuthService:
service = self.oauth.get(provider)
if service is None:
raise InvalidDataSourceOAuthProviderError("Invalid provider")
return service
@dataclass(frozen=True, slots=True)
class DataSourceCredentials:
providers: DatasourceProviderService
actor: ActorAwareDatasourceCredentialGateway
stored: TrustedStoredDatasourceCredentialGateway
def build_data_source_credentials(*, database_client: sessionmaker[Session]) -> DataSourceCredentials:
credentials = SQLAlchemyDatasourceCredentialRepository(session_factory=database_client)
provider_manager = PluginDatasourceManager()
codec = PluginDatasourceCredentialCodec(provider_manager=provider_manager)
refresher = OAuthDatasourceCredentialRefresher(
oauth_clients=PluginDatasourceOAuthClientResolver(configs=credentials, provider_manager=provider_manager),
oauth_handler=OAuthHandler(),
)
return DataSourceCredentials(
providers=DatasourceProviderService(credentials=credentials),
actor=ActorAwareDatasourceCredentialGateway(credentials=credentials, codec=codec, refresher=refresher),
stored=TrustedStoredDatasourceCredentialGateway(credentials=credentials, codec=codec, refresher=refresher),
)
def build_data_source_services(
*,
database_client: sessionmaker[Session],
dataset_access: DatasetAccessService,
datasets: SQLAlchemyDatasetRepository,
documents: SQLAlchemyDocumentRepository,
actor_credentials: ActorAwareDatasourceCredentialGateway,
providers: DatasourceProviderService,
) -> DataSourceServices:
oauth_bindings = SQLAlchemyDataSourceOAuthBindingRepository(session_factory=database_client)
notion_gateway = NotionDataSourceGateway(
client_id=dify_config.NOTION_CLIENT_ID or "",
client_secret=dify_config.NOTION_CLIENT_SECRET or "",
redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/data-source/callback/notion",
http_client=ssrf_proxy,
)
return DataSourceServices(
providers=providers,
api_key_auth=DataSourceApiKeyAuthService(
bindings=SQLAlchemyDataSourceApiKeyAuthBindingRepository(session_factory=database_client),
validator=ProviderApiKeyAuthCredentialValidator(),
encryptor=TenantApiKeyAuthCredentialEncryptor(),
),
oauth={
"notion": DataSourceOAuthService(
provider_name="notion",
provider_gateway=notion_gateway,
bindings=oauth_bindings,
is_internal_provider=dify_config.NOTION_INTEGRATION_TYPE == "internal",
internal_access_token=dify_config.NOTION_INTERNAL_SECRET,
)
},
bindings=DataSourceBindingApplicationService(bindings=oauth_bindings),
notion_imports=NotionImportApplicationService(
dataset_access=dataset_access,
datasets=datasets,
documents=documents,
source=PluginNotionSourceGateway(
credentials=actor_credentials, runtime_loader=DatasourceManager.get_datasource_runtime
),
),
)
def build_website_service(*, database_client: sessionmaker[Session]) -> WebsiteService:
return WebsiteService(providers=build_data_source_credentials(database_client=database_client).providers)
api/extensions/application_services/datasets.py
+32
View File
@@ -0,0 +1,32 @@
"""Shared dataset dependencies for knowledge and data-source use cases."""
from dataclasses import dataclass
from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config
from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository
from services.knowledge.dataset_access import DatasetAccessService, WorkspaceRoleReader
@dataclass(frozen=True, slots=True)
class DatasetDependencies:
datasets: SQLAlchemyDatasetRepository
documents: SQLAlchemyDocumentRepository
access: DatasetAccessService
def build_dataset_dependencies(
*, database_client: sessionmaker[Session], workspace_roles: WorkspaceRoleReader
) -> DatasetDependencies:
datasets = SQLAlchemyDatasetRepository(session_factory=database_client)
return DatasetDependencies(
datasets=datasets,
documents=SQLAlchemyDocumentRepository(session_factory=database_client),
access=DatasetAccessService(
datasets=datasets,
workspace_roles=workspace_roles,
legacy_permissions_enabled=not dify_config.RBAC_ENABLED,
),
)
api/extensions/application_services/knowledge.py
+160 -2
View File
@@ -1,16 +1,174 @@
"""Composition of knowledge base application services."""
"""Composition root for dataset-controller application services."""
from dataclasses import dataclass
from uuid import uuid4
from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config
from core.app.apps.pipeline.pipeline_generator import PipelineGenerator
from core.rag.extractor.entity.datasource_type import DatasourceType
from libs.helper import generate_text_hash
from repositories.knowledge.dataset_api_key_repository import DatasetApiKeyRepository
from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository
from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository
from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository
from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository
from services.api_token_service import ApiTokenCache
from services.app.query_service import AppQueryService
from services.data_source.credential_gateway import (
ActorAwareDatasourceCredentialGateway,
TrustedStoredDatasourceCredentialGateway,
)
from services.data_source.provider_service import DatasourceProviderService
from services.knowledge.api_key_service import DatasetApiKeyService
from services.knowledge.dataset_access import DatasetAccessService
from services.knowledge.dataset_service import DocumentService
from services.knowledge.datasets.adapters import SQLAlchemyDatasetOperations
from services.knowledge.datasets.application import DatasetApplicationService
from services.knowledge.document_sync import DocumentSyncApplicationService
from services.knowledge.document_sync_adapters import CeleryDocumentSyncDispatcher
from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations
from services.knowledge.documents.application import DatasetDocumentApplicationService
from services.knowledge.external.adapters import SQLAlchemyExternalKnowledgeOperations
from services.knowledge.external.application import ExternalKnowledgeApplicationService
from services.knowledge.indexing.adapters.estimate import IndexingEstimateAdapter, SQLAlchemyProcessRuleReader
from services.knowledge.indexing.adapters.sources import (
CompositeStoredSourceResolver,
FileSourceAdapter,
NotionSourceResolver,
WebsiteSourceAdapter,
)
from services.knowledge.indexing.estimate import IndexingEstimateApplicationService
from services.knowledge.metadata.application import MetadataService
from services.knowledge.segments.adapters import (
CelerySegmentBatchImportDispatcher,
ModelManagerSegmentGuard,
RedisSegmentClient,
RedisSegmentIndexingState,
)
from services.knowledge.segments.application import DatasetSegmentApplicationService
from services.knowledge.segments.indexing import SegmentIndexingGateway
from services.tag_application_service import TagTargetQuery
from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task
from tasks.delete_segment_from_index_task import delete_segment_from_index_task
from tasks.disable_segments_from_index_task import disable_segments_from_index_task
from tasks.document_indexing_sync_task import document_indexing_sync_task
from tasks.enable_segments_to_index_task import enable_segments_to_index_task
def build_dataset_api_key_service(*, database_client: sessionmaker[Session]) -> DatasetApiKeyService:
@dataclass(frozen=True, slots=True)
class KnowledgeServices:
metadata: MetadataService
datasets: DatasetApplicationService
external: ExternalKnowledgeApplicationService
documents: DatasetDocumentApplicationService
document_sync: DocumentSyncApplicationService
indexing_estimates: IndexingEstimateApplicationService
segments: DatasetSegmentApplicationService
pipeline_generator: PipelineGenerator
def build_dataset_api_key_service(
*, database_client: sessionmaker[Session], dataset_access: DatasetAccessService
) -> DatasetApiKeyService:
return DatasetApiKeyService(
keys=DatasetApiKeyRepository(session_factory=database_client),
cache=ApiTokenCache,
access=dataset_access,
rbac_enabled=lambda: dify_config.RBAC_ENABLED,
)
def build_knowledge_services(
*,
database_client: sessionmaker[Session],
dataset_access: DatasetAccessService,
datasets: SQLAlchemyDatasetRepository,
documents: SQLAlchemyDocumentRepository,
actor_credentials: ActorAwareDatasourceCredentialGateway,
stored_credentials: TrustedStoredDatasourceCredentialGateway,
providers: DatasourceProviderService,
redis: RedisSegmentClient,
tags: TagTargetQuery,
app_queries: AppQueryService,
) -> KnowledgeServices:
"""Build the dataset-controller knowledge use cases."""
uploads = SQLAlchemyKnowledgeUploadRepository(session_factory=database_client)
notion_sources = NotionSourceResolver(
actor_credentials=actor_credentials,
stored_credentials=stored_credentials,
)
file_sources = FileSourceAdapter(uploads=uploads)
website_sources = WebsiteSourceAdapter()
stored_sources = CompositeStoredSourceResolver(
adapters={
DatasourceType.FILE.value: file_sources,
DatasourceType.NOTION.value: notion_sources,
DatasourceType.WEBSITE.value: website_sources,
}
)
segments = SQLAlchemySegmentRepository(session_factory=database_client)
segment_index = SegmentIndexingGateway(
segments=segments,
uploads=uploads,
redis=redis,
new_session=database_client,
delete_task=delete_segment_from_index_task.delay,
enable_task=enable_segments_to_index_task.delay,
disable_task=disable_segments_from_index_task.delay,
)
return KnowledgeServices(
metadata=MetadataService(
store=SQLAlchemyMetadataRepository(session_factory=database_client), dataset_access=dataset_access
),
datasets=DatasetApplicationService(
dataset_access=dataset_access,
operations=SQLAlchemyDatasetOperations(session_factory=database_client, tags=tags, app_queries=app_queries),
rbac_enabled=dify_config.RBAC_ENABLED,
service_api_url=dify_config.SERVICE_API_URL,
vector_store=dify_config.VECTOR_STORE,
tidb_fulltext=dify_config.TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH,
),
external=ExternalKnowledgeApplicationService(
dataset_access=dataset_access,
operations=SQLAlchemyExternalKnowledgeOperations(session_factory=database_client),
),
documents=DatasetDocumentApplicationService(
dataset_access=dataset_access,
operations=SQLAlchemyDocumentOperations(session_factory=database_client),
metadata_schema=DocumentService.DOCUMENT_METADATA_SCHEMA,
),
pipeline_generator=PipelineGenerator(documents=documents, datasource_providers=providers),
document_sync=DocumentSyncApplicationService(
dataset_access=dataset_access,
documents=documents,
dispatcher=CeleryDocumentSyncDispatcher(delay=document_indexing_sync_task.delay),
),
indexing_estimates=IndexingEstimateApplicationService(
dataset_access=dataset_access,
datasets=datasets,
documents=documents,
files=file_sources,
websites=website_sources,
stored_sources=stored_sources,
notion=notion_sources,
process_rules=SQLAlchemyProcessRuleReader(session_factory=database_client),
runner=IndexingEstimateAdapter(session_factory=database_client),
),
segments=DatasetSegmentApplicationService(
dataset_access=dataset_access,
scopes=datasets,
store=segments,
index=segment_index,
limits=dify_config,
text_hash=generate_text_hash,
uploads=uploads,
model_guard=ModelManagerSegmentGuard(),
indexing_state=RedisSegmentIndexingState(redis),
batch_dispatcher=CelerySegmentBatchImportDispatcher(delay=batch_create_segment_to_index_task.delay),
job_id_factory=lambda: str(uuid4()),
),
)
api/extensions/ext_application_services.py
+59 -61
View File
@@ -19,14 +19,23 @@ from configs import dify_config
from constants.dsl_version import CURRENT_APP_DSL_VERSION
from constants.languages import languages
from core.db.session_factory import get_session_maker
from core.helper.ssrf_proxy import ssrf_proxy
from core.schemas.schema_manager import SchemaManager
from core.tools.tool_file_manager import ToolFileManager
from enums import DeploymentEdition, WebAppAccessMode
from extensions.application_services.agent import AgentAppServices, build_agent_app_services
from extensions.application_services.app import AppServices, build_app_api_key_service, build_app_services
from extensions.application_services.data_sources import (
DataSourceServices,
build_data_source_credentials,
build_data_source_services,
)
from extensions.application_services.datasets import build_dataset_dependencies
from extensions.application_services.installed_app import InstalledAppServices, build_installed_app_services
from extensions.application_services.knowledge import build_dataset_api_key_service
from extensions.application_services.knowledge import (
KnowledgeServices,
build_dataset_api_key_service,
build_knowledge_services,
)
from extensions.application_services.trial_app import TrialAppServices, build_trial_app_services
from extensions.ext_redis import RedisClientWrapper, redis_client
from extensions.ext_storage import storage
@@ -52,8 +61,7 @@ from repositories.app_definition_query_repository import AppDefinitionQueryRepos
from repositories.app_preview_query_repository import AppPreviewQueryRepository
from repositories.app_scoped_end_user_repository import AppScopedEndUserRepo
from repositories.app_statistic_query_repository import AppStatisticQueryRepository
from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
from repositories.data_source_oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository
from repositories.credentials.query_repository import CredentialQueryRepository
from repositories.explore_banner_query_repository import ExploreBannerQueryRepository
from repositories.factory import DifyAPIRepositoryFactory
from repositories.file_grant_repository import FileGrantRepository
@@ -164,15 +172,10 @@ from services.app_statistic_query import AppStatisticQuery
from services.app_task_service import AppTaskControlService
from services.app_tracing_config_gateway import OpsTraceManagerGateway
from services.app_tracing_config_service import AppTracingConfigService
from services.auth.data_source_api_key_auth_gateways import (
ProviderApiKeyAuthCredentialValidator,
TenantApiKeyAuthCredentialEncryptor,
)
from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService
from services.billing_portal_service import BillingPortalService
from services.billing_service import BillingService
from services.compliance_download_service import ComplianceDownloadService
from services.data_source_oauth_service import DataSourceOAuthService, InvalidDataSourceOAuthProviderError
from services.credentials.query import CredentialQuery
from services.enterprise.enterprise_service import EnterpriseService
from services.entities.file_grant_entities import FileGrantLimits
from services.errors.enterprise import EnterpriseServiceError
@@ -191,7 +194,6 @@ from services.message_suggested_questions_adapters import MessageSuggestedQuesti
from services.message_suggested_questions_service import MessageSuggestedQuestions
from services.notification_gateway import BillingNotificationGateway
from services.notification_service import NotificationService
from services.notion_data_source_gateway import NotionDataSourceGateway
from services.oauth_device_adapters import (
DifyConfigOAuthDeviceSettings,
EnterpriseOAuthDeviceSSOGateway,
@@ -337,6 +339,7 @@ class AppScopedEndUserServices:
class ApplicationServices:
agent_apps: AgentAppServices
advanced_prompt_templates: AdvancedPromptTemplateService
credential_queries: CredentialQuery
accounts: AccountServices
app_api_keys: AppApiKeyService
dataset_api_keys: DatasetApiKeyService
@@ -351,8 +354,8 @@ class ApplicationServices:
app_tracing_configs: AppTracingConfigService
billing_portal: BillingPortalService
compliance_downloads: ComplianceDownloadService
data_source_api_key_auth: DataSourceApiKeyAuthService
data_source_oauth: Mapping[str, DataSourceOAuthService]
data_sources: DataSourceServices
knowledge: KnowledgeServices
app_scoped_end_users: AppScopedEndUserServices
webapp_access: WebAppAccessQueryService
web_app_runtime: WebAppRuntimeQueryService
@@ -392,34 +395,6 @@ class ApplicationServices:
tags: TagApplicationService
workflow_statistics: WorkflowStatisticQueryService
def resolve_data_source_oauth(self, provider: str) -> DataSourceOAuthService:
service = self.data_source_oauth.get(provider)
if service is None:
raise InvalidDataSourceOAuthProviderError("Invalid provider")
return service
def _build_data_source_oauth_services(
*,
database_client: sessionmaker[Session],
) -> Mapping[str, DataSourceOAuthService]:
notion_data_source = NotionDataSourceGateway(
client_id=dify_config.NOTION_CLIENT_ID or "",
client_secret=dify_config.NOTION_CLIENT_SECRET or "",
redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/data-source/callback/notion",
http_client=ssrf_proxy,
)
bindings = SQLAlchemyDataSourceOAuthBindingRepository(session_factory=database_client)
return {
"notion": DataSourceOAuthService(
provider_name="notion",
provider_gateway=notion_data_source,
bindings=bindings,
is_internal_provider=dify_config.NOTION_INTEGRATION_TYPE == "internal",
internal_access_token=dify_config.NOTION_INTERNAL_SECRET,
)
}
def _build_oauth_server_service(
*,
@@ -547,7 +522,6 @@ def build_application_services(
redis: RedisClientWrapper,
) -> ApplicationServices:
installation_state = InstallationStateRepository(session_factory=database_client)
data_source_api_key_auth_bindings = SQLAlchemyDataSourceApiKeyAuthBindingRepository(session_factory=database_client)
app_definition_repository = AppDefinitionQueryRepository(session_factory=database_client)
app_definitions = AppDefinitionQueryService(
definitions=app_definition_repository,
@@ -587,6 +561,39 @@ def build_application_services(
sources=database_catalog, exporter=RosterAgentPackageExporter()
)
workspace_query_repository = WorkspaceQueryRepository(session_factory=database_client)
workspace_member_repository = WorkspaceMemberQueryRepository(session_factory=database_client)
dataset_dependencies = build_dataset_dependencies(
database_client=database_client,
workspace_roles=workspace_member_repository,
)
datasource_credentials = build_data_source_credentials(database_client=database_client)
data_sources = build_data_source_services(
database_client=database_client,
dataset_access=dataset_dependencies.access,
datasets=dataset_dependencies.datasets,
documents=dataset_dependencies.documents,
actor_credentials=datasource_credentials.actor,
providers=datasource_credentials.providers,
)
oauth_server = _build_oauth_server_service(database_client=database_client, redis=redis)
apps = build_app_services(
database_client=database_client,
oauth=oauth_server,
recommended_packages=recommended_app_packages,
)
tags = TagApplicationService(tags=TagRepository(session_factory=database_client))
knowledge = build_knowledge_services(
database_client=database_client,
dataset_access=dataset_dependencies.access,
datasets=dataset_dependencies.datasets,
documents=dataset_dependencies.documents,
actor_credentials=datasource_credentials.actor,
stored_credentials=datasource_credentials.stored,
providers=datasource_credentials.providers,
redis=redis,
tags=tags,
app_queries=apps.queries,
)
app_scoped_end_user_repository = AppScopedEndUserRepo(session_factory=database_client)
file_service = FileService(session_factory=database_client)
remote_file_service = RemoteFileService(files=file_service)
@@ -598,8 +605,8 @@ def build_application_services(
workflow_node_execution_repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
session_maker=database_client
)
oauth_server = _build_oauth_server_service(database_client=database_client, redis=redis)
return ApplicationServices(
credential_queries=CredentialQueryRepository(session_factory=database_client),
accounts=AccountServices(
access=AccountAccessService(
accounts=accounts,
@@ -754,11 +761,7 @@ def build_application_services(
enabled=dify_config.RBAC_ENABLED,
),
),
apps=build_app_services(
database_client=database_client,
oauth=oauth_server,
recommended_packages=recommended_app_packages,
),
apps=apps,
agent_apps=build_agent_app_services(database_client=database_client),
advanced_prompt_templates=AdvancedPromptTemplateService(),
app_definitions=app_definitions,
@@ -774,7 +777,10 @@ def build_application_services(
sites=AppSiteCommandRepository(session_factory=database_client),
),
app_api_keys=build_app_api_key_service(database_client=database_client),
dataset_api_keys=build_dataset_api_key_service(database_client=database_client),
dataset_api_keys=build_dataset_api_key_service(
database_client=database_client,
dataset_access=dataset_dependencies.access,
),
app_statistics=AppStatisticQueryRepository(session_factory=database_client),
app_tracing_configs=AppTracingConfigService(
configs=SQLAlchemyAppTracingConfigRepository(session_factory=database_client),
@@ -794,12 +800,8 @@ def build_application_services(
redis_client=redis,
),
),
data_source_api_key_auth=DataSourceApiKeyAuthService(
bindings=data_source_api_key_auth_bindings,
validator=ProviderApiKeyAuthCredentialValidator(),
encryptor=TenantApiKeyAuthCredentialEncryptor(),
),
data_source_oauth=_build_data_source_oauth_services(database_client=database_client),
data_sources=data_sources,
knowledge=knowledge,
app_scoped_end_users=AppScopedEndUserServices(
commands=AppScopedEndUserService(end_users=app_scoped_end_user_repository),
queries=AppScopedEndUserQueryService(end_users=app_scoped_end_user_repository),
@@ -905,9 +907,7 @@ def build_application_services(
plans=DeploymentWorkspacePlanGateway(),
),
workspace_member_queries=WorkspaceMemberQueryService(
members=WorkspaceMemberQueryRepository(
session_factory=database_client,
),
members=workspace_member_repository,
roles=DeploymentWorkspaceMemberRoleResolver(),
),
workflow_app_logs=WorkflowAppLogQueryService(
@@ -927,9 +927,7 @@ def build_application_services(
now=lambda: datetime.now(UTC),
access_token_expire_minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES,
),
tags=TagApplicationService(
tags=TagRepository(session_factory=database_client),
),
tags=tags,
workflow_statistics=WorkflowStatisticQueryService(
workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(
session_maker=database_client,
api/extensions/ext_celery.py
+2
View File
@@ -176,6 +176,8 @@ def init_app(app: DifyApp) -> Celery:
"tasks.trigger_processing_tasks", # async trigger processing
"tasks.generate_summary_index_task", # summary index generation
"tasks.regenerate_summary_index_task", # summary index regeneration
"tasks.rag_pipeline.rag_pipeline_run_task", # regular RAG pipeline execution
"tasks.rag_pipeline.priority_rag_pipeline_run_task", # priority RAG pipeline execution
"tasks.initialize_created_app_rbac_access_task", # app access initialization
"tasks.install_default_plugins_task", # tenant default plugin installation
"tasks.new_agent_beta_task", # New Agent Beta eligibility checks
api/fields/dataset_fields.py
+13 -285
View File
@@ -1,26 +1,21 @@
from collections.abc import Sequence
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
import sqlalchemy as sa
from pydantic import Field, field_validator
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from pydantic import BaseModel, Field, field_validator
from fields.base import ResponseModel
from libs.helper import to_timestamp
from models.account import Account
from models.dataset import (
AppDatasetJoin,
Dataset,
DatasetMetadata,
Document,
ExternalKnowledgeApis,
ExternalKnowledgeBindings,
Pipeline,
)
from models.model import App, Tag, TagBinding
from services.data_source.entities.notion_import import NotionPageType
class NotionEstimatePagePayload(BaseModel):
page_id: str = Field(min_length=1)
page_type: NotionPageType = Field(alias="type")
class NotionEstimateWorkspacePayload(BaseModel):
workspace_id: str = Field(min_length=1)
credential_id: str = Field(min_length=1)
pages: list[NotionEstimatePagePayload] = Field(min_length=1)
class DatasetMetadataResponse(ResponseModel):
@@ -206,270 +201,3 @@ class DatasetDetailResponse(ResponseModel):
@classmethod
def _expand_null_nested(cls, value: object) -> object:
return {} if value is None else value
@dataclass(frozen=True)
class DatasetDetailPrefetch:
"""Batch-loaded values backing the session-scoped fields of a page of datasets.
Building a detail response for a single dataset issues about ten queries for
that row alone, so a list endpoint costs one query per dataset per field.
Loading the same values for every dataset on the page up front keeps the
query count bounded regardless of the page size.
"""
app_counts: dict[str, int] = field(default_factory=dict)
document_counts: dict[str, int] = field(default_factory=dict)
word_counts: dict[str, int] = field(default_factory=dict)
available_document_counts: dict[str, int] = field(default_factory=dict)
author_names: dict[str, str] = field(default_factory=dict)
tags: dict[tuple[str, str], list[Tag]] = field(default_factory=dict)
doc_forms: dict[tuple[str, str], str | None] = field(default_factory=dict)
external_knowledge_infos: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict)
doc_metadata: dict[str, list[DatasetMetadata]] = field(default_factory=dict)
pipeline_published: dict[str, bool] = field(default_factory=dict)
def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Session) -> DatasetDetailPrefetch:
"""Load every session-scoped detail field for the given datasets in a fixed number of queries."""
dataset_ids = [dataset.id for dataset in datasets]
if not dataset_ids:
return DatasetDetailPrefetch()
tenant_ids = {dataset.tenant_id for dataset in datasets}
app_counts = dict(
session.execute(
select(AppDatasetJoin.dataset_id, func.count(AppDatasetJoin.id))
.where(AppDatasetJoin.dataset_id.in_(dataset_ids), App.id == AppDatasetJoin.app_id)
.group_by(AppDatasetJoin.dataset_id)
)
.tuples()
.all()
)
document_counts: dict[str, int] = {}
word_counts: dict[str, int] = {}
available_document_counts: dict[str, int] = {}
available_document_case = sa.case(
(
sa.and_(
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
),
1,
),
else_=0,
)
document_rows = session.execute(
select(
Document.dataset_id,
func.count(Document.id),
func.coalesce(func.sum(Document.word_count), 0),
func.coalesce(func.sum(available_document_case), 0),
)
.where(Document.dataset_id.in_(dataset_ids))
.group_by(Document.dataset_id)
).all()
for dataset_id, document_count, word_count, available_count in document_rows:
document_counts[dataset_id] = document_count
word_counts[dataset_id] = word_count
available_document_counts[dataset_id] = available_count
# doc_form only falls back to the documents table for datasets without a
# chunk structure of their own.
doc_forms: dict[tuple[str, str], str | None] = {}
doc_form_dataset_ids = [dataset.id for dataset in datasets if not dataset.chunk_structure]
if doc_form_dataset_ids:
doc_forms = {
(dataset_id, tenant_id): doc_form
for dataset_id, tenant_id, doc_form in session.execute(
select(Document.dataset_id, Document.tenant_id, func.min(Document.doc_form))
.where(Document.dataset_id.in_(doc_form_dataset_ids), Document.tenant_id.in_(tenant_ids))
.group_by(Document.dataset_id, Document.tenant_id)
).all()
}
author_ids = {dataset.created_by for dataset in datasets if dataset.created_by}
author_names = dict(
session.execute(select(Account.id, Account.name).where(Account.id.in_(author_ids))).tuples().all()
)
tags: dict[tuple[str, str], list[Tag]] = {}
tag_rows = session.execute(
select(TagBinding.target_id, TagBinding.tenant_id, Tag)
.join(TagBinding, Tag.id == TagBinding.tag_id)
.where(
TagBinding.target_id.in_(dataset_ids),
TagBinding.tenant_id.in_(tenant_ids),
Tag.tenant_id == TagBinding.tenant_id,
Tag.type == "knowledge",
)
).all()
for target_id, tenant_id, tag in tag_rows:
tags.setdefault((target_id, tenant_id), []).append(tag)
doc_metadata: dict[str, list[DatasetMetadata]] = {}
for dataset_metadata in session.scalars(
select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids))
).all():
doc_metadata.setdefault(dataset_metadata.dataset_id, []).append(dataset_metadata)
pipeline_ids = {dataset.pipeline_id for dataset in datasets if dataset.pipeline_id}
pipeline_published: dict[str, bool] = {}
if pipeline_ids:
pipeline_published = dict(
session.execute(select(Pipeline.id, Pipeline.is_published).where(Pipeline.id.in_(pipeline_ids)))
.tuples()
.all()
)
return DatasetDetailPrefetch(
app_counts=app_counts,
document_counts=document_counts,
word_counts=word_counts,
available_document_counts=available_document_counts,
author_names=author_names,
tags=tags,
doc_forms=doc_forms,
external_knowledge_infos=_load_external_knowledge_infos(datasets, session=session),
doc_metadata=doc_metadata,
pipeline_published=pipeline_published,
)
def _load_external_knowledge_infos(
datasets: Sequence[Dataset], *, session: Session
) -> dict[tuple[str, str], dict[str, Any]]:
external_datasets = [dataset for dataset in datasets if dataset.provider == "external"]
if not external_datasets:
return {}
external_dataset_ids = [dataset.id for dataset in external_datasets]
external_tenant_ids = {dataset.tenant_id for dataset in external_datasets}
bindings = session.scalars(
select(ExternalKnowledgeBindings).where(
ExternalKnowledgeBindings.dataset_id.in_(external_dataset_ids),
ExternalKnowledgeBindings.tenant_id.in_(external_tenant_ids),
)
).all()
if not bindings:
return {}
apis = {
(api.id, api.tenant_id): api
for api in session.scalars(
select(ExternalKnowledgeApis).where(
ExternalKnowledgeApis.id.in_({binding.external_knowledge_api_id for binding in bindings}),
ExternalKnowledgeApis.tenant_id.in_(external_tenant_ids),
)
).all()
}
infos: dict[tuple[str, str], dict[str, Any]] = {}
for binding in bindings:
info = Dataset.build_external_knowledge_info(
binding, apis.get((binding.external_knowledge_api_id, binding.tenant_id))
)
if info is not None:
infos[(binding.dataset_id, binding.tenant_id)] = info
return infos
@dataclass(frozen=True)
class DatasetDetailResponseSource:
"""Expose session-backed dataset fields during response validation.
With a prefetch the values come from the batch load for the whole page;
without one each field is resolved against the session on access.
"""
dataset: Any
session: Session
prefetch: DatasetDetailPrefetch | None = None
@property
def _prefetch_key(self) -> tuple[str, str]:
return (self.dataset.id, self.dataset.tenant_id)
@property
def app_count(self) -> int:
if self.prefetch is not None:
return self.prefetch.app_counts.get(self.dataset.id, 0)
return self.dataset.get_app_count(session=self.session)
@property
def document_count(self) -> int:
if self.prefetch is not None:
return self.prefetch.document_counts.get(self.dataset.id, 0)
return self.dataset.get_document_count(session=self.session)
@property
def word_count(self) -> int:
if self.prefetch is not None:
return self.prefetch.word_counts.get(self.dataset.id, 0)
return self.dataset.get_word_count(session=self.session)
@property
def author_name(self) -> str | None:
if self.prefetch is not None:
return self.prefetch.author_names.get(self.dataset.created_by)
return self.dataset.get_author_name(session=self.session)
@property
def tags(self) -> Any:
if self.prefetch is not None:
return self.prefetch.tags.get(self._prefetch_key, [])
return self.dataset.get_tags(session=self.session)
@property
def doc_form(self) -> str | None:
if self.prefetch is not None:
if self.dataset.chunk_structure:
return self.dataset.chunk_structure
return self.prefetch.doc_forms.get(self._prefetch_key)
return self.dataset.get_doc_form(session=self.session)
@property
def external_knowledge_info(self) -> Any:
if self.prefetch is not None:
if self.dataset.provider != "external":
return None
return self.prefetch.external_knowledge_infos.get(self._prefetch_key)
return self.dataset.get_external_knowledge_info(session=self.session)
@property
def doc_metadata(self) -> Any:
if self.prefetch is not None:
return self.dataset.build_doc_metadata(self.prefetch.doc_metadata.get(self.dataset.id, []))
return self.dataset.get_doc_metadata(session=self.session)
@property
def is_published(self) -> bool:
if self.prefetch is not None:
if not self.dataset.pipeline_id:
return False
return self.prefetch.pipeline_published.get(self.dataset.pipeline_id, False)
return self.dataset.get_is_published(session=self.session)
@property
def total_documents(self) -> int:
if self.prefetch is not None:
return self.prefetch.document_counts.get(self.dataset.id, 0)
return self.dataset.get_total_documents(session=self.session)
@property
def total_available_documents(self) -> int:
if self.prefetch is not None:
return self.prefetch.available_document_counts.get(self.dataset.id, 0)
return self.dataset.get_total_available_documents(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.dataset, name) # guard-ignore: no-new-getattr -- delegates model fields
def dataset_detail_response_source(
dataset: Any, *, session: Session, prefetch: DatasetDetailPrefetch | None = None
) -> DatasetDetailResponseSource:
return DatasetDetailResponseSource(dataset=dataset, session=session, prefetch=prefetch)
api/fields/document_fields.py
-37
View File
@@ -1,16 +1,12 @@
"""Response schemas for dataset document endpoints."""
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from pydantic import Field, field_validator
from sqlalchemy.orm import Session
from fields.base import ResponseModel
from libs.helper import to_timestamp
from models.dataset import DocMetadataDetailItem, Document
def normalize_enum(value: Any) -> Any:
@@ -70,39 +66,6 @@ class DocumentResponse(ResponseModel):
return to_timestamp(value)
@dataclass(frozen=True)
class DocumentWithSession:
"""Expose session-backed document fields during response validation."""
document: Document
session: Session
@property
def data_source_detail_dict(self) -> dict[str, Any]:
return self.document.get_data_source_detail_dict(session=self.session)
@property
def hit_count(self) -> int:
return self.document.get_hit_count(session=self.session)
@property
def doc_metadata_details(self) -> list[DocMetadataDetailItem] | None:
return self.document.get_doc_metadata_details(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.document, name) # guard-ignore: no-new-getattr -- delegates model fields
def document_response(document: Document, *, session: Session) -> DocumentResponse:
return DocumentResponse.model_validate(
DocumentWithSession(document=document, session=session), from_attributes=True
)
def document_responses(documents: Iterable[Document], *, session: Session) -> list[DocumentResponse]:
return [document_response(document, session=session) for document in documents]
class DocumentListResponse(ResponseModel):
data: list[DocumentResponse]
has_more: bool
api/fields/segment_fields.py
-38
View File
@@ -1,10 +1,6 @@
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from pydantic import field_serializer
from sqlalchemy.orm import Session
from fields.base import ResponseModel
from libs.helper import to_timestamp
@@ -72,40 +68,6 @@ class SegmentResponse(ResponseModel):
return to_timestamp(value)
@dataclass(frozen=True)
class SegmentWithSummary:
"""Expose session-backed segment relations during response validation."""
segment: Any
summary: str | None
session: Session
@property
def child_chunks(self) -> Any:
return self.segment.get_child_chunks(session=self.session, include_full_doc=False)
@property
def attachments(self) -> Any:
return self.segment.get_attachments(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.segment, name)
def segment_response_with_summary(segment: Any, summary: str | None, *, session: Session) -> SegmentResponse:
response_source = SegmentWithSummary(segment=segment, summary=summary, session=session)
return SegmentResponse.model_validate(response_source, from_attributes=True)
def segment_responses_with_summaries(
segments: Iterable[Any],
summaries: Mapping[str, str | None],
*,
session: Session,
) -> list[SegmentResponse]:
return [segment_response_with_summary(segment, summaries.get(segment.id), session=session) for segment in segments]
class SegmentDetailResponse(ResponseModel):
data: SegmentResponse
doc_form: str
api/models/dataset.py
+137 -535
View File
@@ -1,12 +1,5 @@
import base64
import hashlib
import hmac
import json
import logging
import os
import pickle
import re
import time
from collections.abc import Sequence
from datetime import datetime
from json import JSONDecodeError
@@ -14,22 +7,17 @@ from typing import Any, ClassVar, TypedDict, cast, override
from uuid import uuid4
import sqlalchemy as sa
from sqlalchemy import DateTime, String, func, select
from sqlalchemy.orm import Mapped, Session, mapped_column, scoped_session
from sqlalchemy import DateTime, String, func
from sqlalchemy.orm import Mapped, mapped_column
from configs import dify_config
from core.rag.entities import ParentMode, Rule
from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource
from core.rag.entities import PreProcessingRuleKey
from core.rag.index_processor.constant.built_in_field import BuiltInField
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from core.rag.index_processor.constant.query_type import QueryType
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from core.tools.signature import sign_upload_file_preview_url
from extensions.ext_storage import storage
from libs.uuid_utils import uuidv7
from .account import Account
from .base import Base, TypeBase
from .enums import (
from models.base import Base, TypeBase
from models.enums import (
CollectionBindingType,
CreatorUserRole,
DatasetMetadataType,
@@ -45,14 +33,11 @@ from .enums import (
SummaryStatus,
TidbAuthBindingStatus,
)
from .model import App, Tag, TagBinding, UploadFile
from .types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index
logger = logging.getLogger(__name__)
from models.types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index
class PreProcessingRuleItem(TypedDict):
id: str
id: PreProcessingRuleKey
enabled: bool
@@ -95,6 +80,67 @@ class DatasetBindingItem(TypedDict):
name: str
class ExternalKnowledgeApiDict(TypedDict):
id: str
tenant_id: str
name: str
description: str
settings: dict[str, Any] | None
dataset_bindings: list[DatasetBindingItem]
created_by: str
created_at: str
class DocumentDict(TypedDict):
id: str
tenant_id: str
dataset_id: str
position: int
data_source_type: str
data_source_info: str | None
dataset_process_rule_id: str | None
batch: str
name: str
created_from: str
created_by: str
created_api_request_id: str | None
created_at: datetime
processing_started_at: datetime | None
file_id: str | None
word_count: int | None
parsing_completed_at: datetime | None
cleaning_completed_at: datetime | None
splitting_completed_at: datetime | None
tokens: int | None
indexing_latency: float | None
completed_at: datetime | None
is_paused: bool | None
paused_by: str | None
paused_at: datetime | None
error: str | None
stopped_at: datetime | None
indexing_status: str
enabled: bool
disabled_at: datetime | None
disabled_by: str | None
archived: bool
archived_reason: str | None
archived_by: str | None
archived_at: datetime | None
updated_at: datetime
doc_type: str | None
doc_metadata: Any
doc_form: IndexStructureType
doc_language: str | None
display_status: str | None
data_source_info_dict: dict[str, Any]
average_segment_length: int
dataset_process_rule: ProcessRuleDict | None
dataset: None
segment_count: int | None
hit_count: int | None
from models.enums import PermissionEnum
# Backward-compatible alias — new code should import PermissionEnum from models.enums
@@ -149,25 +195,6 @@ class Dataset(Base):
enable_api = mapped_column(sa.Boolean, nullable=False, server_default=sa.true())
is_multimodal = mapped_column(sa.Boolean, default=False, nullable=False)
def get_total_documents(self, *, session: Session) -> int:
return self.get_document_count(session=session)
def get_total_available_documents(self, *, session: Session) -> int:
return (
session.scalar(
select(func.count(Document.id)).where(
Document.dataset_id == self.id,
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
)
)
or 0
)
def get_dataset_keyword_table(self, *, session: Session) -> "DatasetKeywordTable | None":
return session.scalar(select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.id))
@property
def index_struct_dict(self):
return json.loads(self.index_struct) if self.index_struct else None
@@ -180,78 +207,6 @@ class Dataset(Base):
}
return self.retrieval_model or default_retrieval_model
def get_created_by_account(self, *, session: Session) -> Account | None:
return session.get(Account, self.created_by)
def get_author_name(self, *, session: Session) -> str | None:
account = self.get_created_by_account(session=session)
if account:
return account.name
return None
def get_latest_process_rule(self, *, session: Session) -> "DatasetProcessRule | None":
return session.scalar(
select(DatasetProcessRule)
.where(DatasetProcessRule.dataset_id == self.id)
.order_by(DatasetProcessRule.created_at.desc())
.limit(1)
)
def get_app_count(self, *, session: Session) -> int:
return (
session.scalar(
select(func.count(AppDatasetJoin.id)).where(
AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id
)
)
or 0
)
def get_document_count(self, *, session: Session) -> int:
return session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0
def get_available_document_count(self, *, session: Session) -> int:
return (
session.scalar(
select(func.count(Document.id)).where(
Document.dataset_id == self.id,
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
)
)
or 0
)
def get_available_segment_count(self, *, session: Session) -> int:
return (
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.dataset_id == self.id,
DocumentSegment.status == "completed",
DocumentSegment.enabled == True,
)
)
or 0
)
def get_word_count(self, *, session: Session) -> int:
return (
session.scalar(
select(func.coalesce(func.sum(Document.word_count), 0)).where(Document.dataset_id == self.id)
)
or 0
)
def get_doc_form(self, *, session: Session) -> str | None:
if self.chunk_structure:
return self.chunk_structure
return session.scalar(
select(Document.doc_form)
.where(Document.dataset_id == self.id, Document.tenant_id == self.tenant_id)
.limit(1)
)
@property
def retrieval_model_dict(self):
"""Return a normalized retrieval model payload for API responses.
@@ -272,39 +227,6 @@ class Dataset(Base):
return {**default_retrieval_model, **self.retrieval_model}
def get_tags(self, *, session: Session) -> Sequence[Tag]:
tags = session.scalars(
select(Tag)
.join(TagBinding, Tag.id == TagBinding.tag_id)
.where(
TagBinding.target_id == self.id,
TagBinding.tenant_id == self.tenant_id,
Tag.tenant_id == self.tenant_id,
Tag.type == "knowledge",
)
).all()
return tags or []
def get_external_knowledge_info(self, *, session: Session) -> dict[str, Any] | None:
if self.provider != "external":
return None
external_knowledge_binding = session.scalar(
select(ExternalKnowledgeBindings).where(
ExternalKnowledgeBindings.dataset_id == self.id,
ExternalKnowledgeBindings.tenant_id == self.tenant_id,
)
)
if not external_knowledge_binding:
return None
external_knowledge_api = session.scalar(
select(ExternalKnowledgeApis).where(
ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id,
ExternalKnowledgeApis.tenant_id == self.tenant_id,
)
)
return self.build_external_knowledge_info(external_knowledge_binding, external_knowledge_api)
@staticmethod
def build_external_knowledge_info(
external_knowledge_binding: "ExternalKnowledgeBindings | None",
@@ -322,18 +244,6 @@ class Dataset(Base):
"external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""),
}
def get_is_published(self, *, session: Session) -> bool:
if self.pipeline_id:
pipeline = session.scalar(select(Pipeline).where(Pipeline.id == self.pipeline_id))
if pipeline:
return pipeline.is_published
return False
def get_doc_metadata(self, *, session: Session) -> list[dict[str, str]]:
dataset_metadatas = session.scalars(select(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id)).all()
return self.build_doc_metadata(dataset_metadatas)
def build_doc_metadata(self, dataset_metadatas: "Sequence[DatasetMetadata]") -> list[dict[str, str]]:
"""Format metadata rows, appending the built-in fields when they are enabled."""
doc_metadata = [
@@ -407,7 +317,6 @@ class DatasetProcessRule(TypeBase):
)
MODES = ["automatic", "custom", "hierarchical"]
PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
AUTOMATIC_RULES: ClassVar[AutomaticRulesConfig] = {
"pre_processing_rules": [
{"id": "remove_extra_spaces", "enabled": True},
@@ -482,11 +391,11 @@ class Document(Base):
paused_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
# error
error = mapped_column(LongText, nullable=True)
error: Mapped[str | None] = mapped_column(LongText, nullable=True)
stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
# basic fields
indexing_status = mapped_column(
indexing_status: Mapped[str] = mapped_column(
EnumText(IndexingStatus, length=255), nullable=False, server_default=sa.text("'waiting'")
)
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.true())
@@ -512,10 +421,10 @@ class Document(Base):
@property
def display_status(self):
status = None
if self.indexing_status == "waiting":
status = "queuing"
elif self.indexing_status not in {"completed", "error", "waiting"} and self.is_paused:
if self.indexing_status not in {"completed", "error"} and self.is_paused:
status = "paused"
elif self.indexing_status == "waiting":
status = "queuing"
elif self.indexing_status in {"parsing", "cleaning", "splitting", "indexing"}:
status = "indexing"
elif self.indexing_status == "error":
@@ -539,56 +448,6 @@ class Document(Base):
return data_source_info_dict
return {}
def get_data_source_detail_dict(self, *, session: Session) -> dict[str, Any]:
if self.data_source_info:
if self.data_source_type == "upload_file":
data_source_info_dict: dict[str, Any] = json.loads(self.data_source_info)
file_detail = session.scalar(
select(UploadFile).where(UploadFile.id == data_source_info_dict["upload_file_id"])
)
if file_detail:
return {
"upload_file": {
"id": file_detail.id,
"name": file_detail.name,
"size": file_detail.size,
"extension": file_detail.extension,
"mime_type": file_detail.mime_type,
"created_by": file_detail.created_by,
"created_at": file_detail.created_at.timestamp(),
}
}
elif self.data_source_type in {"notion_import", "website_crawl"}:
result: dict[str, Any] = json.loads(self.data_source_info)
return result
return {}
def get_dataset_process_rule(self, *, session: Session) -> "DatasetProcessRule | None":
if self.dataset_process_rule_id:
return session.get(DatasetProcessRule, self.dataset_process_rule_id)
return None
def get_dataset(self, *, session: Session) -> Dataset | None:
"""Load the owning dataset with the caller-owned database session."""
return session.get(Dataset, self.dataset_id)
def get_segment_count(self, *, session: Session) -> int:
return session.scalar(select(func.count(DocumentSegment.id)).where(DocumentSegment.document_id == self.id)) or 0
def get_hit_count(self, *, session: Session) -> int:
return (
session.scalar(
select(func.coalesce(func.sum(DocumentSegment.hit_count), 0)).where(
DocumentSegment.document_id == self.id
)
)
or 0
)
def get_uploader(self, *, session: Session) -> str | None:
user = session.scalar(select(Account).where(Account.id == self.created_by))
return user.name if user else None
@property
def upload_date(self):
return self.created_at
@@ -597,77 +456,57 @@ class Document(Base):
def last_update_date(self):
return self.updated_at
def get_doc_metadata_details(self, *, session: Session) -> list[DocMetadataDetailItem] | None:
if self.doc_metadata:
document_metadatas = session.scalars(
select(DatasetMetadata)
.join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id)
.where(
DatasetMetadata.tenant_id == self.tenant_id,
DatasetMetadata.dataset_id == self.dataset_id,
DatasetMetadataBinding.tenant_id == self.tenant_id,
DatasetMetadataBinding.dataset_id == self.dataset_id,
DatasetMetadataBinding.document_id == self.id,
)
).all()
metadata_list: list[DocMetadataDetailItem] = []
for metadata in document_metadatas:
metadata_dict: DocMetadataDetailItem = {
"id": metadata.id,
"name": metadata.name,
"type": metadata.type,
"value": self.doc_metadata.get(metadata.name),
}
metadata_list.append(metadata_dict)
# deal built-in fields
metadata_list.extend(self.get_built_in_fields(session=session))
return metadata_list
return None
def get_built_in_fields(self, *, session: Session) -> list[DocMetadataDetailItem]:
built_in_fields: list[DocMetadataDetailItem] = []
built_in_fields.append(
{
"id": "built-in",
"name": BuiltInField.document_name,
"type": "string",
"value": self.name,
}
)
built_in_fields.append(
{
"id": "built-in",
"name": BuiltInField.uploader,
"type": "string",
"value": self.get_uploader(session=session),
}
)
built_in_fields.append(
{
"id": "built-in",
"name": BuiltInField.upload_date,
"type": "time",
"value": str(self.created_at.timestamp()),
}
)
built_in_fields.append(
{
"id": "built-in",
"name": BuiltInField.last_update_date,
"type": "time",
"value": str(self.updated_at.timestamp()),
}
)
built_in_fields.append(
{
"id": "built-in",
"name": BuiltInField.source,
"type": "string",
"value": MetadataDataSource[self.data_source_type],
}
)
return built_in_fields
def to_dict(self, *, segment_count: int, hit_count: int, process_rule: DatasetProcessRule | None) -> DocumentDict:
result: DocumentDict = {
"id": self.id,
"tenant_id": self.tenant_id,
"dataset_id": self.dataset_id,
"position": self.position,
"data_source_type": self.data_source_type,
"data_source_info": self.data_source_info,
"dataset_process_rule_id": self.dataset_process_rule_id,
"batch": self.batch,
"name": self.name,
"created_from": self.created_from,
"created_by": self.created_by,
"created_api_request_id": self.created_api_request_id,
"created_at": self.created_at,
"processing_started_at": self.processing_started_at,
"file_id": self.file_id,
"word_count": self.word_count,
"parsing_completed_at": self.parsing_completed_at,
"cleaning_completed_at": self.cleaning_completed_at,
"splitting_completed_at": self.splitting_completed_at,
"tokens": self.tokens,
"indexing_latency": self.indexing_latency,
"completed_at": self.completed_at,
"is_paused": self.is_paused,
"paused_by": self.paused_by,
"paused_at": self.paused_at,
"error": self.error,
"stopped_at": self.stopped_at,
"indexing_status": self.indexing_status,
"enabled": self.enabled,
"disabled_at": self.disabled_at,
"disabled_by": self.disabled_by,
"archived": self.archived,
"archived_reason": self.archived_reason,
"archived_by": self.archived_by,
"archived_at": self.archived_at,
"updated_at": self.updated_at,
"doc_type": self.doc_type,
"doc_metadata": self.doc_metadata,
"doc_form": self.doc_form,
"doc_language": self.doc_language,
"display_status": self.display_status,
"data_source_info_dict": self.data_source_info_dict,
"average_segment_length": (self.word_count or 0) // segment_count if segment_count else 0,
"dataset_process_rule": process_rule.to_dict() if process_rule else None,
"dataset": None,
"segment_count": segment_count,
"hit_count": hit_count,
}
return result
@classmethod
def from_dict(cls, data: dict[str, Any]):
@@ -761,155 +600,6 @@ class DocumentSegment(TypeBase):
stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)
def get_dataset(self, *, session: Session) -> Dataset | None:
"""Load the owning dataset with the caller-owned database session."""
return session.get(Dataset, self.dataset_id)
def get_document(self, *, session: Session) -> Document | None:
"""Load the owning document with the caller-owned database session."""
return session.get(Document, self.document_id)
def previous_segment(self, session: Session) -> "DocumentSegment | None":
return session.scalar(
select(DocumentSegment).where(
DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1
)
)
def next_segment(self, session: Session) -> "DocumentSegment | None":
return session.scalar(
select(DocumentSegment).where(
DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1
)
)
def get_child_chunks(self, *, session: Session, include_full_doc: bool = True) -> Sequence["ChildChunk"]:
"""Load hierarchical child chunks with the caller-owned database session."""
document = session.get(Document, self.document_id)
if not document:
return []
process_rule = document.get_dataset_process_rule(session=session)
if process_rule and process_rule.mode == "hierarchical":
rules_dict = process_rule.rules_dict
if rules_dict:
rules = Rule.model_validate(rules_dict)
if rules.parent_mode and (include_full_doc or rules.parent_mode != ParentMode.FULL_DOC):
child_chunks = session.scalars(
select(ChildChunk).where(ChildChunk.segment_id == self.id).order_by(ChildChunk.position.asc())
).all()
return child_chunks or []
return []
@property
def sign_content(self) -> str:
return self.get_sign_content()
def get_sign_content(self) -> str:
signed_urls: list[tuple[int, int, str]] = []
text = self.content
# For data before v0.10.0
pattern = r"/files/([a-f0-9\-]+)/image-preview(?:\?.*?)?"
matches = re.finditer(pattern, text)
for match in matches:
upload_file_id = match.group(1)
nonce = os.urandom(16).hex()
timestamp = str(int(time.time()))
data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
secret_key = dify_config.SECRET_KEY.encode()
sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
encoded_sign = base64.urlsafe_b64encode(sign).decode()
params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
base_url = f"/files/{upload_file_id}/image-preview"
signed_url = f"{base_url}?{params}"
signed_urls.append((match.start(), match.end(), signed_url))
# For data after v0.10.0
pattern = r"/files/([a-f0-9\-]+)/file-preview(?:\?.*?)?"
matches = re.finditer(pattern, text)
for match in matches:
upload_file_id = match.group(1)
nonce = os.urandom(16).hex()
timestamp = str(int(time.time()))
data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
secret_key = dify_config.SECRET_KEY.encode()
sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
encoded_sign = base64.urlsafe_b64encode(sign).decode()
params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
base_url = f"/files/{upload_file_id}/file-preview"
signed_url = f"{base_url}?{params}"
signed_urls.append((match.start(), match.end(), signed_url))
# For tools directory - direct file formats (e.g., .png, .jpg, etc.)
# Match URL including any query parameters up to common URL boundaries (space, parenthesis, quotes)
pattern = r"/files/tools/([a-f0-9\-]+)\.([a-zA-Z0-9]+)(?:\?[^\s\)\"\']*)?"
matches = re.finditer(pattern, text)
for match in matches:
upload_file_id = match.group(1)
file_extension = match.group(2)
nonce = os.urandom(16).hex()
timestamp = str(int(time.time()))
data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
secret_key = dify_config.SECRET_KEY.encode()
sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
encoded_sign = base64.urlsafe_b64encode(sign).decode()
params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
base_url = f"/files/tools/{upload_file_id}.{file_extension}"
signed_url = f"{base_url}?{params}"
signed_urls.append((match.start(), match.end(), signed_url))
# Reconstruct the text with signed URLs
offset = 0
for start, end, signed_url in signed_urls:
text = text[: start + offset] + signed_url + text[end + offset :]
offset += len(signed_url) - (end - start)
return text
def get_attachments(self, *, session: Session) -> list[AttachmentItem]:
"""Load attachment metadata with the caller-owned database session."""
# Use JOIN to fetch attachments in a single query instead of two separate queries
attachments_with_bindings = session.execute(
select(SegmentAttachmentBinding, UploadFile)
.join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
.where(
SegmentAttachmentBinding.tenant_id == self.tenant_id,
SegmentAttachmentBinding.dataset_id == self.dataset_id,
SegmentAttachmentBinding.document_id == self.document_id,
SegmentAttachmentBinding.segment_id == self.id,
)
).all()
if not attachments_with_bindings:
return []
attachment_list: list[AttachmentItem] = []
for _, attachment in attachments_with_bindings:
upload_file_id = attachment.id
nonce = os.urandom(16).hex()
timestamp = str(int(time.time()))
data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
secret_key = dify_config.SECRET_KEY.encode()
sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
encoded_sign = base64.urlsafe_b64encode(sign).decode()
params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
reference_url = dify_config.FILES_URL or dify_config.CONSOLE_API_URL or ""
base_url = f"{reference_url}/files/{upload_file_id}/image-preview"
source_url = f"{base_url}?{params}"
attachment_list.append(
{
"id": attachment.id,
"name": attachment.name,
"size": attachment.size,
"extension": attachment.extension,
"mime_type": attachment.mime_type,
"source_url": source_url,
}
)
return attachment_list
class ChildChunk(TypeBase):
__tablename__ = "child_chunks"
@@ -957,15 +647,6 @@ class ChildChunk(TypeBase):
)
error: Mapped[str | None] = mapped_column(LongText, nullable=True, init=False)
def dataset(self, session: Session) -> Dataset | None:
return session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))
def document(self, session: Session) -> Document | None:
return session.scalar(select(Document).where(Document.id == self.document_id))
def segment(self, session: Session) -> DocumentSegment | None:
return session.scalar(select(DocumentSegment).where(DocumentSegment.id == self.segment_id))
class AppDatasetJoin(TypeBase):
__tablename__ = "app_dataset_joins"
@@ -988,9 +669,6 @@ class AppDatasetJoin(TypeBase):
DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
)
def app(self, session: Session) -> App | None:
return session.get(App, self.app_id)
class DatasetQuery(TypeBase):
__tablename__ = "dataset_queries"
@@ -1017,37 +695,6 @@ class DatasetQuery(TypeBase):
DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
)
def get_queries(self, *, session: Session) -> list[dict[str, Any]]:
try:
queries = json.loads(self.content)
if isinstance(queries, list):
for query in queries:
if query["content_type"] == QueryType.IMAGE_QUERY:
file_info = session.scalar(select(UploadFile).where(UploadFile.id == query["content"]))
if file_info:
query["file_info"] = {
"id": file_info.id,
"name": file_info.name,
"size": file_info.size,
"extension": file_info.extension,
"mime_type": file_info.mime_type,
"source_url": sign_upload_file_preview_url(file_info.id, file_info.extension),
}
else:
query["file_info"] = None
return queries
else:
return [queries]
except JSONDecodeError:
return [
{
"content_type": QueryType.TEXT_QUERY,
"content": self.content,
"file_info": None,
}
]
class DatasetKeywordTable(TypeBase):
__tablename__ = "dataset_keyword_tables"
@@ -1067,40 +714,6 @@ class DatasetKeywordTable(TypeBase):
keyword_table: Mapped[str] = mapped_column(LongText, nullable=False)
data_source_type: Mapped[str] = mapped_column(String(255), nullable=False, default="database")
def get_keyword_table_dict(self, *, session: Session) -> dict[str, set[Any]] | None:
class SetDecoder(json.JSONDecoder):
def __init__(self, *args: Any, **kwargs: Any) -> None:
def object_hook(dct: Any) -> Any:
if isinstance(dct, dict):
result: dict[str, Any] = {}
items = cast(dict[str, Any], dct).items()
for keyword, node_idxs in items:
if isinstance(node_idxs, list):
result[keyword] = set(node_idxs)
else:
result[keyword] = node_idxs
return result
return dct
super().__init__(object_hook=object_hook, *args, **kwargs)
# get dataset
dataset = session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))
if not dataset:
return None
if self.data_source_type == "database":
return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None
else:
file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt"
try:
keyword_table_text = storage.load_once(file_key)
if keyword_table_text:
return json.loads(keyword_table_text.decode("utf-8"), cls=SetDecoder)
return None
except Exception:
logger.exception("Failed to load keyword table from file: %s", file_key)
return None
class Embedding(TypeBase):
__tablename__ = "embeddings"
@@ -1263,6 +876,18 @@ class ExternalKnowledgeApis(TypeBase):
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
def to_dict(self, *, dataset_bindings: list[DatasetBindingItem]) -> ExternalKnowledgeApiDict:
return {
"id": self.id,
"tenant_id": self.tenant_id,
"name": self.name,
"description": self.description,
"settings": self.settings_dict,
"dataset_bindings": dataset_bindings,
"created_by": self.created_by,
"created_at": self.created_at.isoformat(),
}
@property
def settings_dict(self) -> dict[str, Any] | None:
try:
@@ -1270,18 +895,6 @@ class ExternalKnowledgeApis(TypeBase):
except JSONDecodeError:
return None
def get_dataset_bindings(self, *, session: Session) -> list[DatasetBindingItem]:
external_knowledge_bindings = session.scalars(
select(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.external_knowledge_api_id == self.id)
).all()
dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings]
datasets = session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all()
dataset_bindings: list[DatasetBindingItem] = []
for dataset in datasets:
dataset_bindings.append({"id": dataset.id, "name": dataset.name})
return dataset_bindings
class ExternalKnowledgeBindings(TypeBase):
__tablename__ = "external_knowledge_bindings"
@@ -1467,12 +1080,6 @@ class PipelineCustomizedTemplate(TypeBase):
init=False,
)
def created_user_name(self, session: Session) -> str:
account = session.scalar(select(Account).where(Account.id == self.created_by))
if account:
return account.name
return ""
class Pipeline(TypeBase):
__tablename__ = "pipelines"
@@ -1500,11 +1107,6 @@ class Pipeline(TypeBase):
init=False,
)
def retrieve_dataset(self, session: Session | scoped_session):
return session.scalar(
select(Dataset).where(Dataset.pipeline_id == self.id, Dataset.tenant_id == self.tenant_id)
)
class DocumentPipelineExecutionLog(TypeBase):
__tablename__ = "document_pipeline_execution_logs"
api/openapi/markdown/console-openapi.md
+38 -65
View File
@@ -5083,30 +5083,9 @@ Get compliance document download link
### [GET] /data-source/integrates
#### Responses
| Code | Description |
| ---- | ----------- |
| 200 | Success |
### [PATCH] /data-source/integrates
#### Responses
| Code | Description |
| ---- | ----------- |
| 200 | Success |
### [GET] /data-source/integrates/{binding_id}/{action}
#### Parameters
| Name | Located in | Description | Required | Schema |
| ---- | ---------- | ----------- | -------- | ------ |
| action | path | | Yes | string |
| binding_id | path | | Yes | string (uuid) |
#### Responses
| Code | Description |
| ---- | ----------- |
| 200 | Success |
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Success | **application/json**: [DataSourceIntegrateListResponse](#datasourceintegratelistresponse)<br> |
### [PATCH] /data-source/integrates/{binding_id}/{action}
#### Parameters
@@ -5118,9 +5097,9 @@ Get compliance document download link
#### Responses
| Code | Description |
| ---- | ----------- |
| 200 | Success |
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)<br> |
### [GET] /datasets
Get list of datasets
@@ -5155,7 +5134,7 @@ Create a new dataset
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 201 | Dataset created successfully | **application/json**: [DatasetDetailResponse](#datasetdetailresponse)<br> |
| 201 | Dataset created successfully | **application/json**: [DatasetDetailWithPartialMembersResponse](#datasetdetailwithpartialmembersresponse)<br> |
| 400 | Invalid request parameters | |
### [GET] /datasets/api-base-info
@@ -5218,25 +5197,6 @@ Delete dataset API key
| ---- | ----------- | ------ |
| 200 | Batch import status | **application/json**: [SegmentBatchImportStatusResponse](#segmentbatchimportstatusresponse)<br> |
### [POST] /datasets/batch_import_status/{job_id}
#### Parameters
| Name | Located in | Description | Required | Schema |
| ---- | ---------- | ----------- | -------- | ------ |
| job_id | path | | Yes | string (uuid) |
#### Request Body
| Required | Schema |
| -------- | ------ |
| Yes | **application/json**: [BatchImportPayload](#batchimportpayload)<br> |
#### Responses
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Batch import started | **application/json**: [SegmentBatchImportStatusResponse](#segmentbatchimportstatusresponse)<br> |
### [POST] /datasets/external
Create external knowledge dataset
@@ -5970,20 +5930,6 @@ Update document processing status (pause/resume)
| ---- | ----------- | ------ |
| 200 | Segments retrieved successfully | **application/json**: [ConsoleSegmentListResponse](#consolesegmentlistresponse)<br> |
### [GET] /datasets/{dataset_id}/documents/{document_id}/segments/batch_import
#### Parameters
| Name | Located in | Description | Required | Schema |
| ---- | ---------- | ----------- | -------- | ------ |
| dataset_id | path | | Yes | string (uuid) |
| document_id | path | | Yes | string (uuid) |
#### Responses
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Batch import status | **application/json**: [SegmentBatchImportStatusResponse](#segmentbatchimportstatusresponse)<br> |
### [POST] /datasets/{dataset_id}/documents/{document_id}/segments/batch_import
#### Parameters
@@ -18102,7 +18048,7 @@ Site token customization strategy
| page_id | string | | Yes |
| page_name | string | | Yes |
| parent_id | string | | Yes |
| type | string | | Yes |
| type | [NotionPageType](#notionpagetype) | | Yes |
#### DataSourceIntegrateResponse
@@ -21037,15 +20983,30 @@ Coarse node-level status used by Inspector to pick a banner.
| notifications | [ [NotificationItemResponse](#notificationitemresponse) ] | | Yes |
| should_show | boolean | | Yes |
#### NotionEstimatePagePayload
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| page_id | string | | Yes |
| type | [NotionPageType](#notionpagetype) | | Yes |
#### NotionEstimatePayload
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| doc_form | string, <br>**Default:** text_model | | No |
| doc_language | string, <br>**Default:** English | | No |
| notion_info_list | [ object ] | | Yes |
| notion_info_list | [ [NotionEstimateWorkspacePayload](#notionestimateworkspacepayload) ] | | Yes |
| process_rule | object | | Yes |
#### NotionEstimateWorkspacePayload
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| credential_id | string | | Yes |
| pages | [ [NotionEstimatePagePayload](#notionestimatepagepayload) ] | | Yes |
| workspace_id | string | | Yes |
#### NotionIcon
| Name | Type | Description | Required |
@@ -21077,7 +21038,7 @@ Coarse node-level status used by Inspector to pick a banner.
| page_id | string | | Yes |
| page_name | string | | Yes |
| parent_id | string | | Yes |
| type | string | | Yes |
| type | [NotionPageType](#notionpagetype) | | Yes |
#### NotionIntegrateWorkspaceResponse
@@ -21097,6 +21058,12 @@ Coarse node-level status used by Inspector to pick a banner.
| page_name | string | | Yes |
| type | string | | Yes |
#### NotionPageType
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| NotionPageType | string | | |
#### OAuthCallbackQuery
| Name | Type | Description | Required |
@@ -22232,7 +22199,13 @@ Verification of the plugin.
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| enabled | boolean | Whether this preprocessing rule is enabled. | Yes |
| id | string, <br>**Available values:** "remove_extra_spaces", "remove_stopwords", "remove_urls_emails" | Rule identifier.<br>*Enum:* `"remove_extra_spaces"`, `"remove_stopwords"`, `"remove_urls_emails"` | Yes |
| id | [PreProcessingRuleKey](#preprocessingrulekey) | Rule identifier. | Yes |
#### PreProcessingRuleKey
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| PreProcessingRuleKey | string | | |
#### PreviewDetail
api/openapi/markdown/service-openapi.md
+7 -1
View File
@@ -3836,7 +3836,13 @@ JSON result for published runs and draft runs using `response_mode: blocking`.
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| enabled | boolean | Whether this preprocessing rule is enabled. | Yes |
| id | string, <br>**Available values:** "remove_extra_spaces", "remove_stopwords", "remove_urls_emails" | Rule identifier.<br>*Enum:* `"remove_extra_spaces"`, `"remove_stopwords"`, `"remove_urls_emails"` | Yes |
| id | [PreProcessingRuleKey](#preprocessingrulekey) | Rule identifier. | Yes |
#### PreProcessingRuleKey
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| PreProcessingRuleKey | string | | |
#### ProcessRule
api/repositories/credentials/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Persistence adapters for the credentials domain."""
api/repositories/credentials/query_repository.py
+250
View File
@@ -0,0 +1,250 @@
"""Persistence-owned credential queries and visibility filtering shared across domains."""
from collections.abc import Sequence
from typing import override
from sqlalchemy import and_, func, or_, select
from sqlalchemy.orm import InstrumentedAttribute, Session, sessionmaker
from models.credential_permission import CredentialPermission, CredentialType
from models.enums import PermissionEnum
from models.oauth import DatasourceProvider
from models.provider import ProviderCredential
from models.provider_ids import ModelProviderID
from models.tools import BuiltinToolProvider
from models.trigger import TriggerSubscription, WorkflowPluginTrigger
from services.credentials.query import (
CredentialQuery,
DatasourceCredentialListItem,
ModelCredentialRecord,
ToolCredentialRecord,
TriggerSubscriptionRecord,
)
def apply_credential_visibility_filter_for_actor(
query,
*,
tenant_id: str,
model_id_column: InstrumentedAttribute,
model_user_id_column: InstrumentedAttribute,
model_visibility_column: InstrumentedAttribute,
credential_type: str,
actor_id: str,
):
"""Apply the canonical tenant- and actor-scoped credential visibility policy.
Team credentials and legacy rows without an owner are visible to all tenant
members. Personal credentials are visible to their owner, while partially
shared credentials additionally require an enabled permission row for the
same tenant and credential type. Administrator roles do not bypass this
policy.
"""
partial_member_ids = (
select(CredentialPermission.credential_id)
.where(
CredentialPermission.tenant_id == tenant_id,
CredentialPermission.credential_type == credential_type,
CredentialPermission.account_id == actor_id,
CredentialPermission.has_permission.is_(True),
)
.correlate_except(CredentialPermission)
)
return query.where(
or_(
model_visibility_column == PermissionEnum.ALL_TEAM,
model_user_id_column.is_(None),
model_user_id_column == actor_id,
and_(model_visibility_column == PermissionEnum.PARTIAL_TEAM, model_id_column.in_(partial_member_ids)),
)
)
class CredentialQueryRepository(CredentialQuery):
"""Load complete credential lists without exposing ORM objects or SQL to callers."""
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
@override
def list_models(self, *, workspace_id: str, provider: str, actor_id: str | None) -> Sequence[ModelCredentialRecord]:
names = [provider]
provider_id = ModelProviderID(provider)
if provider_id.is_langgenius():
names.append(provider_id.provider_name if "/" in provider else str(provider_id))
query = (
select(ProviderCredential)
.where(ProviderCredential.tenant_id == workspace_id, ProviderCredential.provider_name.in_(names))
.order_by(ProviderCredential.created_at.desc())
)
if actor_id is not None:
query = apply_credential_visibility_filter_for_actor(
query,
tenant_id=workspace_id,
model_id_column=ProviderCredential.id,
model_user_id_column=ProviderCredential.user_id,
model_visibility_column=ProviderCredential.visibility,
credential_type=CredentialType.PROVIDER_CREDENTIAL,
actor_id=actor_id,
)
with self._session_factory() as session:
return [ModelCredentialRecord(row.id, row.credential_name) for row in session.scalars(query)]
@override
def list_datasources(
self, *, workspace_id: str, provider: str, plugin_id: str, actor_id: str | None
) -> Sequence[DatasourceCredentialListItem]:
owner = (
DatasourceProvider.tenant_id == workspace_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
query = select(DatasourceProvider).where(*owner)
if actor_id is not None:
query = apply_credential_visibility_filter_for_actor(
query,
tenant_id=workspace_id,
model_id_column=DatasourceProvider.id,
model_user_id_column=DatasourceProvider.user_id,
model_visibility_column=DatasourceProvider.visibility,
credential_type=CredentialType.DATASOURCE_PROVIDER,
actor_id=actor_id,
)
with self._session_factory() as session:
rows = session.scalars(query).all()
if not rows:
return []
# Preserve the configured default even when that credential is hidden from this actor.
default_id = session.scalar(
select(DatasourceProvider.id)
.where(*owner)
.order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc())
.limit(1)
)
return [
DatasourceCredentialListItem(
id=row.id,
name=row.name,
auth_type=row.auth_type,
encrypted_credentials=dict(row.encrypted_credentials),
avatar_url=row.avatar_url,
is_default=row.id == default_id,
)
for row in rows
]
@override
def list_tools(
self,
*,
workspace_id: str,
provider: str,
actor_id: str | None,
include_credential_ids: Sequence[str] = (),
) -> Sequence[ToolCredentialRecord]:
owner = (BuiltinToolProvider.tenant_id == workspace_id, BuiltinToolProvider.provider == provider)
order = (BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc())
query = select(BuiltinToolProvider).where(*owner).order_by(*order)
if actor_id is not None:
query = apply_credential_visibility_filter_for_actor(
query,
tenant_id=workspace_id,
model_id_column=BuiltinToolProvider.id,
model_user_id_column=BuiltinToolProvider.user_id,
model_visibility_column=BuiltinToolProvider.visibility,
credential_type=CredentialType.BUILTIN_TOOL_PROVIDER,
actor_id=actor_id,
)
with self._session_factory() as session:
visible = list(session.scalars(query))
wanted = set(include_credential_ids) - {row.id for row in visible} - {""}
borrowed: list[BuiltinToolProvider] = (
list(
session.scalars(
select(BuiltinToolProvider).where(*owner, BuiltinToolProvider.id.in_(wanted)).order_by(*order)
)
)
if wanted
else []
)
rows = visible + borrowed
partial_ids = [row.id for row in rows if row.visibility == PermissionEnum.PARTIAL_TEAM]
members: dict[str, list[str]] = {}
if partial_ids:
for credential_id, account_id in session.execute(
select(CredentialPermission.credential_id, CredentialPermission.account_id).where(
CredentialPermission.tenant_id == workspace_id,
CredentialPermission.credential_type == CredentialType.BUILTIN_TOOL_PROVIDER,
CredentialPermission.credential_id.in_(partial_ids),
CredentialPermission.has_permission.is_(True),
)
):
members.setdefault(credential_id, []).append(account_id)
borrowed_ids = {row.id for row in borrowed}
return [
ToolCredentialRecord(
id=row.id,
name=row.name,
provider=row.provider,
credential_type=row.credential_type,
credentials=dict(row.credentials),
is_default=(bool(visible) and row.id == visible[0].id) or row.is_default,
visibility=str(row.visibility),
created_by=row.user_id or "",
partial_member_ids=tuple(members.get(row.id, [])),
from_other_member=row.id in borrowed_ids,
)
for row in rows
]
@override
def list_trigger_subscriptions(
self, *, workspace_id: str, provider: str, actor_id: str | None
) -> Sequence[TriggerSubscriptionRecord]:
query = (
select(TriggerSubscription)
.where(TriggerSubscription.tenant_id == workspace_id, TriggerSubscription.provider_id == provider)
.order_by(TriggerSubscription.created_at.desc())
)
if actor_id is not None:
query = apply_credential_visibility_filter_for_actor(
query,
tenant_id=workspace_id,
model_id_column=TriggerSubscription.id,
model_user_id_column=TriggerSubscription.user_id,
model_visibility_column=TriggerSubscription.visibility,
credential_type=CredentialType.TRIGGER_SUBSCRIPTION,
actor_id=actor_id,
)
with self._session_factory() as session:
rows = session.scalars(query).all()
if not rows:
return []
counts = dict(
session.execute(
select(
WorkflowPluginTrigger.subscription_id, func.count(func.distinct(WorkflowPluginTrigger.app_id))
)
.where(
WorkflowPluginTrigger.tenant_id == workspace_id,
WorkflowPluginTrigger.subscription_id.in_([row.id for row in rows]),
)
.group_by(WorkflowPluginTrigger.subscription_id)
)
.tuples()
.all()
)
return [
TriggerSubscriptionRecord(
id=row.id,
name=row.name,
provider=row.provider_id,
credential_type=row.credential_type,
credentials=dict(row.credentials),
endpoint_id=row.endpoint_id,
parameters=dict(row.parameters),
properties=dict(row.properties),
workflows_in_use=counts.get(row.id, 0),
)
for row in rows
]
api/repositories/data_source/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Persistence adapters for the data-source domain."""
api/repositories/data_source_api_key_auth_repository.py → api/repositories/data_source/api_key_auth_repository.py
+2 -2
View File
@@ -8,8 +8,8 @@ from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from models.source import DataSourceApiKeyAuthBinding
from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthBindingRepository
from services.entities.data_source_api_key_auth_entities import (
from services.data_source.auth.api_key_service import DataSourceApiKeyAuthBindingRepository
from services.data_source.entities.api_key_auth import (
DataSourceApiKeyAuthBindingRecord,
DataSourceApiKeyAuthCredentials,
)
api/repositories/data_source/credential_repository.py
+252
View File
@@ -0,0 +1,252 @@
"""Actor-aware and trusted-source SQLAlchemy repository for datasource credentials."""
import json
from collections.abc import Mapping
from typing import cast
from pydantic import TypeAdapter, ValidationError
from sqlalchemy import select, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session, sessionmaker
from libs.datetime_utils import naive_utc_now
from models.credential_permission import CredentialType
from models.dataset import Dataset, Document
from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider
from repositories.credentials.query_repository import apply_credential_visibility_filter_for_actor
from services.data_source.entities.credential import DatasourceCredentialRecord, DatasourceOAuthClientConfigRecord
_CREDENTIALS_ADAPTER = TypeAdapter(dict[str, object])
def _source_mapping(value: object) -> Mapping[str, object]:
try:
if isinstance(value, str):
return _CREDENTIALS_ADAPTER.validate_python(json.loads(value))
return _CREDENTIALS_ADAPTER.validate_python(value)
except (TypeError, ValueError, ValidationError):
return {}
def _datasource_credential_record(provider: DatasourceProvider) -> DatasourceCredentialRecord:
return DatasourceCredentialRecord(
id=provider.id,
workspace_id=provider.tenant_id,
owner_id=provider.user_id,
name=provider.name,
provider=provider.provider,
plugin_id=provider.plugin_id,
auth_type=provider.auth_type,
encrypted_credentials=_CREDENTIALS_ADAPTER.validate_python(provider.encrypted_credentials),
expires_at=provider.expires_at,
updated_at=provider.updated_at,
)
class SQLAlchemyDatasourceCredentialRepository:
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
def get_provider_credential(
self, *, workspace_id: str, provider: str, plugin_id: str, credential_id: str | None
) -> DatasourceCredentialRecord | None:
with self._session_factory() as session:
rows = _read_provider_credentials(
session,
workspace_id=workspace_id,
provider=provider,
plugin_id=plugin_id,
credential_id=credential_id,
first_only=True,
)
return _datasource_credential_record(rows[0]) if rows else None
def list_provider_credentials(
self, *, workspace_id: str, provider: str, plugin_id: str
) -> list[DatasourceCredentialRecord]:
with self._session_factory() as session:
return [
_datasource_credential_record(row)
for row in _read_provider_credentials(
session,
workspace_id=workspace_id,
provider=provider,
plugin_id=plugin_id,
)
]
def get_visible(
self,
*,
workspace_id: str,
actor_id: str,
credential_id: str,
provider: str,
plugin_id: str,
) -> DatasourceCredentialRecord | None:
query = select(DatasourceProvider).where(
DatasourceProvider.tenant_id == workspace_id,
DatasourceProvider.id == credential_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
query = apply_credential_visibility_filter_for_actor(
query,
tenant_id=workspace_id,
model_id_column=DatasourceProvider.id,
model_user_id_column=DatasourceProvider.user_id,
model_visibility_column=DatasourceProvider.visibility,
credential_type=CredentialType.DATASOURCE_PROVIDER,
actor_id=actor_id,
)
with self._session_factory() as session:
provider_record = session.scalar(query.limit(1))
return _datasource_credential_record(provider_record) if provider_record is not None else None
def get_oauth_client_config(
self,
*,
workspace_id: str,
provider: str,
plugin_id: str,
) -> DatasourceOAuthClientConfigRecord:
"""Load OAuth client rows in one bounded read transaction."""
with self._session_factory() as session:
return _read_oauth_client_config(session, workspace_id=workspace_id, provider=provider, plugin_id=plugin_id)
def update_if_unchanged(
self,
*,
record: DatasourceCredentialRecord,
encrypted_credentials: Mapping[str, object],
expires_at: int,
) -> bool:
with self._session_factory.begin() as session:
return _update_datasource_credentials_if_unchanged(
session, record=record, encrypted_credentials=encrypted_credentials, expires_at=expires_at
)
def get_for_stored_document(
self,
*,
workspace_id: str,
dataset_id: str,
document_id: str,
credential_id: str | None,
provider: str,
plugin_id: str,
) -> DatasourceCredentialRecord | None:
with self._session_factory() as session:
query = (
select(DatasourceProvider, Document.data_source_info)
.select_from(Document)
.join(Dataset, Dataset.id == Document.dataset_id)
.join(DatasourceProvider, DatasourceProvider.tenant_id == Document.tenant_id)
.where(
Dataset.id == dataset_id,
Dataset.tenant_id == workspace_id,
Document.id == document_id,
Document.dataset_id == dataset_id,
Document.tenant_id == workspace_id,
Document.data_source_type == "notion_import",
DatasourceProvider.tenant_id == workspace_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
)
if credential_id is not None:
query = query.where(DatasourceProvider.id == credential_id)
else:
query = query.where(DatasourceProvider.is_default.is_(True)).order_by(
DatasourceProvider.created_at.asc()
)
row = session.execute(query.limit(1)).one_or_none()
if row is None:
return None
datasource_provider, raw_source_info = row
source_info = _source_mapping(raw_source_info)
if credential_id is not None and source_info.get("credential_id") != credential_id:
return None
return _datasource_credential_record(datasource_provider)
def _read_oauth_client_config(
session: Session, *, workspace_id: str, provider: str, plugin_id: str
) -> DatasourceOAuthClientConfigRecord:
tenant_config = session.scalar(
select(DatasourceOauthTenantParamConfig)
.where(
DatasourceOauthTenantParamConfig.tenant_id == workspace_id,
DatasourceOauthTenantParamConfig.provider == provider,
DatasourceOauthTenantParamConfig.plugin_id == plugin_id,
DatasourceOauthTenantParamConfig.enabled.is_(True),
)
.limit(1)
)
system_config = session.scalar(
select(DatasourceOauthParamConfig)
.where(
DatasourceOauthParamConfig.provider == provider,
DatasourceOauthParamConfig.plugin_id == plugin_id,
)
.limit(1)
)
return DatasourceOAuthClientConfigRecord(
encrypted_tenant_params=(dict(tenant_config.client_params) if tenant_config is not None else None),
system_credentials=(dict(system_config.system_credentials) if system_config is not None else None),
)
def _update_datasource_credentials_if_unchanged(
session: Session,
*,
record: DatasourceCredentialRecord,
encrypted_credentials: Mapping[str, object],
expires_at: int,
) -> bool:
credentials = _CREDENTIALS_ADAPTER.validate_python(encrypted_credentials)
result = session.execute(
update(DatasourceProvider)
.where(
DatasourceProvider.id == record.id,
DatasourceProvider.tenant_id == record.workspace_id,
DatasourceProvider.provider == record.provider,
DatasourceProvider.plugin_id == record.plugin_id,
DatasourceProvider.user_id == record.owner_id,
DatasourceProvider.encrypted_credentials
== _CREDENTIALS_ADAPTER.validate_python(record.encrypted_credentials),
DatasourceProvider.expires_at == record.expires_at,
)
.values(
encrypted_credentials=credentials,
expires_at=expires_at,
updated_at=naive_utc_now(),
)
)
return cast(CursorResult[object], result).rowcount == 1
def _read_provider_credentials(
session: Session,
*,
workspace_id: str,
provider: str,
plugin_id: str,
credential_id: str | None = None,
first_only: bool = False,
) -> list[DatasourceProvider]:
statement = (
select(DatasourceProvider)
.where(
DatasourceProvider.tenant_id == workspace_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
.order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc())
)
if credential_id is not None:
statement = statement.where(DatasourceProvider.id == credential_id)
if first_only:
statement = statement.limit(1)
return list(session.scalars(statement))
api/repositories/data_source_oauth_binding_repository.py → api/repositories/data_source/oauth_binding_repository.py
+66 -7
View File
@@ -1,24 +1,27 @@
"""SQLAlchemy repository for OAuth data-source bindings."""
from collections.abc import Mapping
from typing import override
from typing import cast, override
from pydantic import TypeAdapter
from sqlalchemy import select
from sqlalchemy import select, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session, sessionmaker
from libs.datetime_utils import naive_utc_now
from models.source import DataSourceOauthBinding
from services.data_source_oauth_service import DataSourceOAuthBindingRepository
from services.entities.data_source_oauth_entities import (
from services.data_source.binding_application_service import BindingMutationResult, DataSourceBindingStore
from services.data_source.entities.oauth import (
DataSourceBindingSummary,
DataSourceOAuthAuthorization,
DataSourceOAuthBindingRecord,
)
from services.data_source.oauth_service import DataSourceOAuthBindingRepository
_SOURCE_INFO_ADAPTER = TypeAdapter(dict[str, object])
class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepository):
class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepository, DataSourceBindingStore):
def __init__(self, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
@@ -68,7 +71,7 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor
DataSourceOauthBinding.tenant_id == workspace_id,
DataSourceOauthBinding.provider == provider,
DataSourceOauthBinding.id == binding_id,
DataSourceOauthBinding.disabled.is_(False),
DataSourceOauthBinding.disabled.is_not(True),
)
)
if binding is None:
@@ -96,7 +99,7 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor
DataSourceOauthBinding.tenant_id == workspace_id,
DataSourceOauthBinding.provider == provider,
DataSourceOauthBinding.id == binding_id,
DataSourceOauthBinding.disabled.is_(False),
DataSourceOauthBinding.disabled.is_not(True),
)
)
if binding is None:
@@ -105,3 +108,59 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor
binding.source_info = validated_source_info
binding.updated_at = naive_utc_now()
return True
@override
def list_enabled_bindings(self, *, workspace_id: str) -> tuple[DataSourceBindingSummary, ...]:
with self._session_factory() as session:
bindings = session.scalars(
select(DataSourceOauthBinding).where(
DataSourceOauthBinding.tenant_id == workspace_id,
DataSourceOauthBinding.disabled.is_not(True),
)
).all()
return tuple(
DataSourceBindingSummary(
id=binding.id,
provider=binding.provider,
created_at=binding.created_at,
disabled=False,
source_info=_SOURCE_INFO_ADAPTER.validate_python(binding.source_info),
)
for binding in bindings
)
@override
def change_disabled_state(
self,
*,
workspace_id: str,
binding_id: str,
disabled: bool,
) -> BindingMutationResult:
state_predicate = (
DataSourceOauthBinding.disabled.is_not(True) if disabled else DataSourceOauthBinding.disabled.is_(True)
)
with self._session_factory.begin() as session:
result = session.execute(
update(DataSourceOauthBinding)
.where(
DataSourceOauthBinding.tenant_id == workspace_id,
DataSourceOauthBinding.id == binding_id,
state_predicate,
)
.values(disabled=disabled, updated_at=naive_utc_now())
)
if cast(CursorResult[object], result).rowcount == 1:
return "updated"
binding_exists = session.scalar(
select(DataSourceOauthBinding.id)
.where(
DataSourceOauthBinding.tenant_id == workspace_id,
DataSourceOauthBinding.id == binding_id,
)
.limit(1)
)
if binding_exists is None:
return "not_found"
return "already_disabled" if disabled else "already_enabled"
api/repositories/knowledge/__init__.py
+1 -1
View File
@@ -1 +1 @@
"""Knowledge base persistence and queries."""
"""Knowledge-domain persistence adapters and queries."""
api/repositories/knowledge/dataset_api_key_repository.py
+2 -27
View File
@@ -5,8 +5,8 @@ from typing import override
from sqlalchemy import Select, func, select
from sqlalchemy.orm import Session, sessionmaker
from models.account import Tenant, TenantAccountJoin
from models.dataset import Dataset, DatasetPermission
from models.account import Tenant
from models.dataset import Dataset
from models.enums import ApiTokenType
from models.model import ApiToken
from repositories.knowledge import dataset_api_key_bindings
@@ -17,37 +17,12 @@ from services.auth.api_key_contracts import (
ApiKeyResourceNotFoundError,
)
from services.knowledge.api_key_service import DatasetApiKeyStore, UnknownDatasetIdsError
from services.knowledge.dataset_access import DatasetAccess
class DatasetApiKeyRepository(DatasetApiKeyStore):
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
@override
def get_dataset_access(self, workspace_id: str, dataset_id: str, account_id: str) -> DatasetAccess:
with self._session_factory() as session:
dataset = self._get_dataset(session, workspace_id, dataset_id)
role = session.scalar(
select(TenantAccountJoin.role).where(
TenantAccountJoin.tenant_id == workspace_id, TenantAccountJoin.account_id == account_id
)
)
has_permission = (
session.scalar(
select(DatasetPermission.id)
.where(
DatasetPermission.tenant_id == workspace_id,
DatasetPermission.dataset_id == dataset_id,
DatasetPermission.account_id == account_id,
DatasetPermission.has_permission.is_(True),
)
.limit(1)
)
is not None
)
return DatasetAccess(dataset.permission, dataset.maintainer, role, has_permission)
@override
def list_keys(self, workspace_id: str, dataset_id: str) -> tuple[ApiKeyRecord, ...]:
with self._session_factory() as session:
api/repositories/knowledge/dataset_read_repository.py
+755
View File
@@ -0,0 +1,755 @@
"""Explicit reads for knowledge models within the caller's transaction.
These query functions are shared by repositories and existing transaction owners;
they never create, commit, or close a session and perform no external I/O.
"""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any
import sqlalchemy as sa
from sqlalchemy import func, select, tuple_
from sqlalchemy.orm import Session, scoped_session
from core.app.file_access import DatabaseFileAccessController
from core.rag.entities import ParentMode, Rule
from models.account import Account
from models.dataset import (
AppDatasetJoin,
ChildChunk,
Dataset,
DatasetBindingItem,
DatasetKeywordTable,
DatasetMetadata,
DatasetMetadataBinding,
DatasetProcessRule,
DatasetQuery,
Document,
DocumentSegment,
ExternalKnowledgeApis,
ExternalKnowledgeBindings,
Pipeline,
PipelineCustomizedTemplate,
SegmentAttachmentBinding,
)
from models.model import App, Tag, TagBinding, UploadFile
from models.tools import ToolFile
@dataclass(frozen=True)
class DocumentReadBatch:
uploads: dict[tuple[str, str], UploadFile] = field(default_factory=dict)
metadatas: dict[str, list[DatasetMetadata]] = field(default_factory=dict)
built_in_enabled: dict[tuple[str, str], bool] = field(default_factory=dict)
uploader_names: dict[str, str] = field(default_factory=dict)
segment_counts: dict[str, int] = field(default_factory=dict)
hit_counts: dict[str, int] = field(default_factory=dict)
process_rules: dict[str, DatasetProcessRule] = field(default_factory=dict)
def load_document_read_batch(documents: Sequence[Document], *, session: Session) -> DocumentReadBatch:
"""Load a page's document relations with a bounded number of queries."""
if not documents:
return DocumentReadBatch()
owners = [(doc.tenant_id, doc.dataset_id, doc.id) for doc in documents]
dataset_owners = list({(doc.tenant_id, doc.dataset_id) for doc in documents})
upload_owners = list(
{
(doc.tenant_id, file_id)
for doc in documents
if doc.data_source_type == "upload_file" and (file_id := doc.data_source_info_dict.get("upload_file_id"))
}
)
uploads = (
{
(file.tenant_id, file.id): file
for file in session.scalars(
select(UploadFile).where(tuple_(UploadFile.tenant_id, UploadFile.id).in_(upload_owners))
)
}
if upload_owners
else {}
)
metadatas: dict[str, list[DatasetMetadata]] = {}
for doc_id, metadata in session.execute(
select(DatasetMetadataBinding.document_id, DatasetMetadata)
.join(DatasetMetadata, DatasetMetadata.id == DatasetMetadataBinding.metadata_id)
.where(
tuple_(
DatasetMetadataBinding.tenant_id, DatasetMetadataBinding.dataset_id, DatasetMetadataBinding.document_id
).in_(owners),
DatasetMetadata.tenant_id == DatasetMetadataBinding.tenant_id,
DatasetMetadata.dataset_id == DatasetMetadataBinding.dataset_id,
)
):
metadatas.setdefault(doc_id, []).append(metadata)
built_in_enabled = {
(tenant_id, dataset_id): enabled
for tenant_id, dataset_id, enabled in session.execute(
select(Dataset.tenant_id, Dataset.id, Dataset.built_in_field_enabled).where(
tuple_(Dataset.tenant_id, Dataset.id).in_(dataset_owners)
)
)
}
uploader_names = dict(
session.execute(select(Account.id, Account.name).where(Account.id.in_({doc.created_by for doc in documents})))
.tuples()
.all()
)
counts = session.execute(
select(
DocumentSegment.document_id,
func.count(DocumentSegment.id),
func.coalesce(func.sum(DocumentSegment.hit_count), 0),
)
.where(tuple_(DocumentSegment.tenant_id, DocumentSegment.dataset_id, DocumentSegment.document_id).in_(owners))
.group_by(DocumentSegment.document_id)
).all()
rule_ids = {doc.dataset_process_rule_id for doc in documents if doc.dataset_process_rule_id}
process_rules = (
{
rule.id: rule
for rule in session.scalars(
select(DatasetProcessRule)
.join(Dataset, Dataset.id == DatasetProcessRule.dataset_id)
.where(DatasetProcessRule.id.in_(rule_ids), tuple_(Dataset.tenant_id, Dataset.id).in_(dataset_owners))
)
}
if rule_ids
else {}
)
return DocumentReadBatch(
uploads=uploads,
metadatas=metadatas,
built_in_enabled=built_in_enabled,
uploader_names=uploader_names,
segment_counts={doc_id: count for doc_id, count, _ in counts},
hit_counts={doc_id: hits for doc_id, _, hits in counts},
process_rules=process_rules,
)
def get_segment_content_file_ids(
segment: DocumentSegment, *, upload_ids: set[str], tool_ids: set[str], session: Session
) -> tuple[set[str], set[str]]:
access = DatabaseFileAccessController()
uploads = (
session.scalars(
access.apply_upload_file_filters(
select(UploadFile).where(UploadFile.id.in_(upload_ids), UploadFile.tenant_id == segment.tenant_id)
)
).all()
if upload_ids
else []
)
tools = (
session.scalars(
access.apply_tool_file_filters(
select(ToolFile).where(ToolFile.id.in_(tool_ids), ToolFile.tenant_id == segment.tenant_id)
)
).all()
if tool_ids
else []
)
return {file.id for file in uploads}, {file.id for file in tools}
def get_segment_attachment_files(segment: DocumentSegment, *, session: Session) -> Sequence[UploadFile]:
return session.scalars(
select(UploadFile)
.join(SegmentAttachmentBinding, UploadFile.id == SegmentAttachmentBinding.attachment_id)
.where(
SegmentAttachmentBinding.tenant_id == segment.tenant_id,
SegmentAttachmentBinding.dataset_id == segment.dataset_id,
SegmentAttachmentBinding.document_id == segment.document_id,
SegmentAttachmentBinding.segment_id == segment.id,
UploadFile.tenant_id == segment.tenant_id,
)
).all()
def is_retrieved_segment_owned(
segment: DocumentSegment, *, tenant_id: str, dataset_ids: Sequence[str], session: Session
) -> bool:
"""Validate retrieved identifiers against the complete persisted owner chain."""
if segment.tenant_id != tenant_id or segment.dataset_id not in dataset_ids:
return False
return (
session.scalar(
select(DocumentSegment.id)
.join(Document, Document.id == DocumentSegment.document_id)
.join(Dataset, Dataset.id == Document.dataset_id)
.where(
Dataset.tenant_id == tenant_id,
Dataset.id == segment.dataset_id,
Document.tenant_id == tenant_id,
Document.dataset_id == segment.dataset_id,
Document.id == segment.document_id,
Document.enabled.is_(True),
Document.archived.is_(False),
DocumentSegment.tenant_id == tenant_id,
DocumentSegment.dataset_id == segment.dataset_id,
DocumentSegment.id == segment.id,
DocumentSegment.enabled.is_(True),
)
)
is not None
)
def get_document_upload_file(document: Document, *, session: Session) -> UploadFile | None:
file_id = document.data_source_info_dict.get("upload_file_id")
if not file_id:
return None
return session.scalar(
select(UploadFile).where(UploadFile.id == file_id, UploadFile.tenant_id == document.tenant_id)
)
def get_document_metadata_rows(document: Document, *, session: Session) -> Sequence[DatasetMetadata]:
return session.scalars(
select(DatasetMetadata)
.join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id)
.where(
DatasetMetadata.tenant_id == document.tenant_id,
DatasetMetadata.dataset_id == document.dataset_id,
DatasetMetadataBinding.tenant_id == document.tenant_id,
DatasetMetadataBinding.dataset_id == document.dataset_id,
DatasetMetadataBinding.document_id == document.id,
)
).all()
def get_query_upload_files(query: DatasetQuery, file_ids: set[str], *, session: Session) -> dict[str, UploadFile]:
if not file_ids:
return {}
rows = session.scalars(
select(UploadFile)
.join(Dataset, Dataset.tenant_id == UploadFile.tenant_id)
.where(Dataset.id == query.dataset_id, UploadFile.id.in_(file_ids))
).all()
return {file.id: file for file in rows}
def get_external_api_bindings_batch(
apis: Sequence[ExternalKnowledgeApis], *, session: Session
) -> dict[str, list[DatasetBindingItem]]:
if not apis:
return {}
allowed = {(api.id, api.tenant_id) for api in apis}
rows = session.execute(
select(
ExternalKnowledgeBindings.external_knowledge_api_id,
ExternalKnowledgeBindings.tenant_id,
Dataset.id,
Dataset.name,
)
.join(Dataset, Dataset.id == ExternalKnowledgeBindings.dataset_id)
.where(
tuple_(ExternalKnowledgeBindings.external_knowledge_api_id, ExternalKnowledgeBindings.tenant_id).in_(
allowed
),
Dataset.tenant_id == ExternalKnowledgeBindings.tenant_id,
)
.distinct()
).all()
bindings: dict[str, list[DatasetBindingItem]] = {}
for api_id, tenant_id, dataset_id, name in rows:
if (api_id, tenant_id) in allowed:
bindings.setdefault(api_id, []).append({"id": dataset_id, "name": name})
return bindings
def get_external_api_dataset_bindings(api: ExternalKnowledgeApis, *, session: Session) -> list[DatasetBindingItem]:
return get_external_api_bindings_batch([api], session=session).get(api.id, [])
def get_dataset_keyword_table(dataset: Dataset, *, session: Session) -> DatasetKeywordTable | None:
return session.scalar(select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == dataset.id))
def get_dataset_creator(dataset: Dataset, *, session: Session) -> Account | None:
return session.get(Account, dataset.created_by)
def get_dataset_author_name(dataset: Dataset, *, session: Session) -> str | None:
account = get_dataset_creator(dataset, session=session)
if account:
return account.name
return None
def get_latest_dataset_process_rule(dataset: Dataset, *, session: Session) -> DatasetProcessRule | None:
return session.scalar(
select(DatasetProcessRule)
.where(DatasetProcessRule.dataset_id == dataset.id)
.order_by(DatasetProcessRule.created_at.desc())
.limit(1)
)
def get_dataset_app_count(dataset: Dataset, *, session: Session) -> int:
return (
session.scalar(
select(func.count(AppDatasetJoin.id)).where(
AppDatasetJoin.dataset_id == dataset.id,
App.id == AppDatasetJoin.app_id,
App.tenant_id == dataset.tenant_id,
)
)
or 0
)
def get_dataset_document_count(dataset: Dataset, *, session: Session) -> int:
return (
session.scalar(
select(func.count(Document.id)).where(
Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id
)
)
or 0
)
def get_dataset_available_document_count(dataset: Dataset, *, session: Session) -> int:
return (
session.scalar(
select(func.count(Document.id)).where(
Document.dataset_id == dataset.id,
Document.tenant_id == dataset.tenant_id,
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
)
)
or 0
)
def get_dataset_available_segment_count(dataset: Dataset, *, session: Session) -> int:
return (
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.tenant_id == dataset.tenant_id,
DocumentSegment.status == "completed",
DocumentSegment.enabled == True,
)
)
or 0
)
def get_dataset_word_count(dataset: Dataset, *, session: Session) -> int:
return (
session.scalar(
select(func.coalesce(func.sum(Document.word_count), 0)).where(
Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id
)
)
or 0
)
def get_dataset_doc_form(dataset: Dataset, *, session: Session) -> str | None:
if dataset.chunk_structure:
return dataset.chunk_structure
return session.scalar(
select(Document.doc_form)
.where(Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id)
.limit(1)
)
def get_dataset_tags(dataset: Dataset, *, session: Session) -> Sequence[Tag]:
tags = session.scalars(
select(Tag)
.join(TagBinding, Tag.id == TagBinding.tag_id)
.where(
TagBinding.target_id == dataset.id,
TagBinding.tenant_id == dataset.tenant_id,
Tag.tenant_id == dataset.tenant_id,
Tag.type == "knowledge",
)
).all()
return tags or []
def get_external_knowledge_info(dataset: Dataset, *, session: Session) -> dict[str, Any] | None:
if dataset.provider != "external":
return None
external_knowledge_binding = session.scalar(
select(ExternalKnowledgeBindings).where(
ExternalKnowledgeBindings.dataset_id == dataset.id, ExternalKnowledgeBindings.tenant_id == dataset.tenant_id
)
)
if not external_knowledge_binding:
return None
external_knowledge_api = session.scalar(
select(ExternalKnowledgeApis).where(
ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id,
ExternalKnowledgeApis.tenant_id == dataset.tenant_id,
)
)
return dataset.build_external_knowledge_info(external_knowledge_binding, external_knowledge_api)
def get_dataset_is_published(dataset: Dataset, *, session: Session) -> bool:
if dataset.pipeline_id:
pipeline = session.scalar(
select(Pipeline).where(Pipeline.id == dataset.pipeline_id, Pipeline.tenant_id == dataset.tenant_id)
)
if pipeline:
return pipeline.is_published
return False
def get_dataset_doc_metadata(dataset: Dataset, *, session: Session) -> list[dict[str, str]]:
dataset_metadatas = session.scalars(
select(DatasetMetadata).where(
DatasetMetadata.dataset_id == dataset.id, DatasetMetadata.tenant_id == dataset.tenant_id
)
).all()
return dataset.build_doc_metadata(dataset_metadatas)
def get_document_process_rule(document: Document, *, session: Session) -> DatasetProcessRule | None:
if document.dataset_process_rule_id:
return session.scalar(
select(DatasetProcessRule).where(
DatasetProcessRule.id == document.dataset_process_rule_id,
DatasetProcessRule.dataset_id == document.dataset_id,
)
)
return None
def get_document_dataset(document: Document, *, session: Session) -> Dataset | None:
"""Load the owning dataset with the caller-owned database session."""
return session.scalar(
select(Dataset).where(Dataset.id == document.dataset_id, Dataset.tenant_id == document.tenant_id)
)
def get_document_segment_count(document: Document, *, session: Session) -> int:
return (
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.document_id == document.id,
DocumentSegment.dataset_id == document.dataset_id,
DocumentSegment.tenant_id == document.tenant_id,
)
)
or 0
)
def get_document_hit_count(document: Document, *, session: Session) -> int:
return (
session.scalar(
select(func.coalesce(func.sum(DocumentSegment.hit_count), 0)).where(
DocumentSegment.document_id == document.id,
DocumentSegment.dataset_id == document.dataset_id,
DocumentSegment.tenant_id == document.tenant_id,
)
)
or 0
)
def get_document_uploader(document: Document, *, session: Session) -> str | None:
user = session.scalar(select(Account).where(Account.id == document.created_by))
return user.name if user else None
def get_segment_dataset(segment: DocumentSegment, *, session: Session) -> Dataset | None:
"""Load the owning dataset with the caller-owned database session."""
return session.scalar(
select(Dataset).where(Dataset.id == segment.dataset_id, Dataset.tenant_id == segment.tenant_id)
)
def get_segment_document(segment: DocumentSegment, *, session: Session) -> Document | None:
"""Load the owning document with the caller-owned database session."""
return session.scalar(
select(Document).where(
Document.id == segment.document_id,
Document.dataset_id == segment.dataset_id,
Document.tenant_id == segment.tenant_id,
)
)
def get_previous_segment(segment: DocumentSegment, session: Session) -> DocumentSegment | None:
return session.scalar(
select(DocumentSegment).where(
DocumentSegment.document_id == segment.document_id,
DocumentSegment.dataset_id == segment.dataset_id,
DocumentSegment.tenant_id == segment.tenant_id,
DocumentSegment.position == segment.position - 1,
)
)
def get_next_segment(segment: DocumentSegment, session: Session) -> DocumentSegment | None:
return session.scalar(
select(DocumentSegment).where(
DocumentSegment.document_id == segment.document_id,
DocumentSegment.dataset_id == segment.dataset_id,
DocumentSegment.tenant_id == segment.tenant_id,
DocumentSegment.position == segment.position + 1,
)
)
def get_segment_child_chunks(
segment: DocumentSegment, *, session: Session, include_full_doc: bool = True
) -> Sequence[ChildChunk]:
"""Load hierarchical child chunks with the caller-owned database session."""
document = get_segment_document(segment, session=session)
if not document:
return []
process_rule = get_document_process_rule(document, session=session)
if process_rule and process_rule.mode == "hierarchical":
rules_dict = process_rule.rules_dict
if rules_dict:
rules = Rule.model_validate(rules_dict)
if rules.parent_mode and (include_full_doc or rules.parent_mode != ParentMode.FULL_DOC):
child_chunks = session.scalars(
select(ChildChunk)
.where(
ChildChunk.segment_id == segment.id,
ChildChunk.tenant_id == segment.tenant_id,
ChildChunk.dataset_id == segment.dataset_id,
ChildChunk.document_id == segment.document_id,
)
.order_by(ChildChunk.position.asc())
).all()
return child_chunks or []
return []
def get_child_dataset(chunk: ChildChunk, session: Session) -> Dataset | None:
return session.scalar(select(Dataset).where(Dataset.id == chunk.dataset_id, Dataset.tenant_id == chunk.tenant_id))
def get_child_document(chunk: ChildChunk, session: Session) -> Document | None:
return session.scalar(
select(Document).where(
Document.id == chunk.document_id,
Document.dataset_id == chunk.dataset_id,
Document.tenant_id == chunk.tenant_id,
)
)
def get_child_segment(chunk: ChildChunk, session: Session) -> DocumentSegment | None:
return session.scalar(
select(DocumentSegment).where(
DocumentSegment.id == chunk.segment_id,
DocumentSegment.document_id == chunk.document_id,
DocumentSegment.dataset_id == chunk.dataset_id,
DocumentSegment.tenant_id == chunk.tenant_id,
)
)
def get_joined_app(binding: AppDatasetJoin, session: Session) -> App | None:
return session.get(App, binding.app_id)
def get_pipeline_template_creator_name(template: PipelineCustomizedTemplate, session: Session) -> str:
account = session.scalar(select(Account).where(Account.id == template.created_by))
if account:
return account.name
return ""
def get_pipeline_dataset(pipeline: Pipeline, session: Session | scoped_session) -> Dataset | None:
return session.scalar(
select(Dataset).where(Dataset.pipeline_id == pipeline.id, Dataset.tenant_id == pipeline.tenant_id)
)
@dataclass(frozen=True)
class DatasetDetailPrefetch:
"""Batch-loaded values backing the session-scoped fields of a page of datasets.
Building a detail response for a single dataset issues about ten queries for
that row alone, so a list endpoint costs one query per dataset per field.
Loading the same values for every dataset on the page up front keeps the
query count bounded regardless of the page size.
"""
app_counts: dict[str, int] = field(default_factory=dict)
document_counts: dict[str, int] = field(default_factory=dict)
word_counts: dict[str, int] = field(default_factory=dict)
available_document_counts: dict[str, int] = field(default_factory=dict)
author_names: dict[str, str] = field(default_factory=dict)
tags: dict[tuple[str, str], list[Tag]] = field(default_factory=dict)
doc_forms: dict[tuple[str, str], str | None] = field(default_factory=dict)
external_knowledge_infos: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict)
doc_metadata: dict[str, list[DatasetMetadata]] = field(default_factory=dict)
pipeline_published: dict[str, bool] = field(default_factory=dict)
def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Session) -> DatasetDetailPrefetch:
"""Load every session-scoped detail field for the given datasets in a fixed number of queries."""
dataset_ids = [dataset.id for dataset in datasets]
if not dataset_ids:
return DatasetDetailPrefetch()
tenant_ids = {dataset.tenant_id for dataset in datasets}
dataset_owners = [(dataset.tenant_id, dataset.id) for dataset in datasets]
app_counts = dict(
session.execute(
select(AppDatasetJoin.dataset_id, func.count(AppDatasetJoin.id))
.where(
AppDatasetJoin.dataset_id.in_(dataset_ids),
App.id == AppDatasetJoin.app_id,
tuple_(App.tenant_id, AppDatasetJoin.dataset_id).in_(dataset_owners),
)
.group_by(AppDatasetJoin.dataset_id)
)
.tuples()
.all()
)
document_counts: dict[str, int] = {}
word_counts: dict[str, int] = {}
available_document_counts: dict[str, int] = {}
available_document_case = sa.case(
(
sa.and_(
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
),
1,
),
else_=0,
)
document_rows = session.execute(
select(
Document.dataset_id,
func.count(Document.id),
func.coalesce(func.sum(Document.word_count), 0),
func.coalesce(func.sum(available_document_case), 0),
)
.where(tuple_(Document.tenant_id, Document.dataset_id).in_(dataset_owners))
.group_by(Document.dataset_id)
).all()
for dataset_id, document_count, word_count, available_count in document_rows:
document_counts[dataset_id] = document_count
word_counts[dataset_id] = word_count
available_document_counts[dataset_id] = available_count
# doc_form only falls back to the documents table for datasets without a
# chunk structure of their own.
doc_forms: dict[tuple[str, str], str | None] = {}
doc_form_dataset_ids = [dataset.id for dataset in datasets if not dataset.chunk_structure]
if doc_form_dataset_ids:
doc_forms = {
(dataset_id, tenant_id): doc_form
for dataset_id, tenant_id, doc_form in session.execute(
select(Document.dataset_id, Document.tenant_id, func.min(Document.doc_form))
.where(
Document.dataset_id.in_(doc_form_dataset_ids),
tuple_(Document.tenant_id, Document.dataset_id).in_(dataset_owners),
)
.group_by(Document.dataset_id, Document.tenant_id)
).all()
}
author_ids = {dataset.created_by for dataset in datasets if dataset.created_by}
author_names = dict(
session.execute(select(Account.id, Account.name).where(Account.id.in_(author_ids))).tuples().all()
)
tags: dict[tuple[str, str], list[Tag]] = {}
tag_rows = session.execute(
select(TagBinding.target_id, TagBinding.tenant_id, Tag)
.join(TagBinding, Tag.id == TagBinding.tag_id)
.where(
tuple_(TagBinding.tenant_id, TagBinding.target_id).in_(dataset_owners),
Tag.tenant_id == TagBinding.tenant_id,
Tag.type == "knowledge",
)
).all()
for target_id, tenant_id, tag in tag_rows:
tags.setdefault((target_id, tenant_id), []).append(tag)
doc_metadata: dict[str, list[DatasetMetadata]] = {}
for dataset_metadata in session.scalars(
select(DatasetMetadata).where(tuple_(DatasetMetadata.tenant_id, DatasetMetadata.dataset_id).in_(dataset_owners))
).all():
doc_metadata.setdefault(dataset_metadata.dataset_id, []).append(dataset_metadata)
pipeline_ids = {dataset.pipeline_id for dataset in datasets if dataset.pipeline_id}
pipeline_published: dict[str, bool] = {}
if pipeline_ids:
pipeline_owners = [(dataset.tenant_id, dataset.pipeline_id) for dataset in datasets if dataset.pipeline_id]
pipeline_published = dict(
session.execute(
select(Pipeline.id, Pipeline.is_published).where(
tuple_(Pipeline.tenant_id, Pipeline.id).in_(pipeline_owners)
)
)
.tuples()
.all()
)
return DatasetDetailPrefetch(
app_counts=app_counts,
document_counts=document_counts,
word_counts=word_counts,
available_document_counts=available_document_counts,
author_names=author_names,
tags=tags,
doc_forms=doc_forms,
external_knowledge_infos=_load_external_knowledge_infos(datasets, session=session),
doc_metadata=doc_metadata,
pipeline_published=pipeline_published,
)
def _load_external_knowledge_infos(
datasets: Sequence[Dataset], *, session: Session
) -> dict[tuple[str, str], dict[str, Any]]:
external_datasets = [dataset for dataset in datasets if dataset.provider == "external"]
if not external_datasets:
return {}
external_owners = [(dataset.tenant_id, dataset.id) for dataset in external_datasets]
external_tenant_ids = {dataset.tenant_id for dataset in external_datasets}
bindings = session.scalars(
select(ExternalKnowledgeBindings).where(
tuple_(ExternalKnowledgeBindings.tenant_id, ExternalKnowledgeBindings.dataset_id).in_(external_owners),
)
).all()
if not bindings:
return {}
apis = {
(api.id, api.tenant_id): api
for api in session.scalars(
select(ExternalKnowledgeApis).where(
ExternalKnowledgeApis.id.in_({binding.external_knowledge_api_id for binding in bindings}),
ExternalKnowledgeApis.tenant_id.in_(external_tenant_ids),
)
).all()
}
infos: dict[tuple[str, str], dict[str, Any]] = {}
for binding in bindings:
info = Dataset.build_external_knowledge_info(
binding, apis.get((binding.external_knowledge_api_id, binding.tenant_id))
)
if info is not None:
infos[(binding.dataset_id, binding.tenant_id)] = info
return infos
api/repositories/knowledge/dataset_repository.py
+153
View File
@@ -0,0 +1,153 @@
"""SQLAlchemy repository for tenant-owned dataset state."""
from sqlalchemy import and_, exists, literal, select
from sqlalchemy.orm import Session, sessionmaker
from models.dataset import Dataset, DatasetPermission, Document
from services.knowledge.dataset_access import DatasetAccessRecord, DatasetAccessSnapshot
from services.knowledge.indexing.estimate import DatasetEstimateRecord
from services.knowledge.resource_scope import DatasetRef
from services.knowledge.segments.application import SegmentDatasetRecord, SegmentDocumentRecord, SegmentScope
def _dataset_access_record(dataset: Dataset) -> DatasetAccessRecord:
return DatasetAccessRecord(
id=dataset.id,
workspace_id=dataset.tenant_id,
maintainer_id=dataset.maintainer,
permission=str(dataset.permission),
)
def _dataset_estimate_record(dataset: Dataset) -> DatasetEstimateRecord:
return DatasetEstimateRecord(
id=dataset.id,
workspace_id=dataset.tenant_id,
indexing_technique=str(dataset.indexing_technique) if dataset.indexing_technique is not None else None,
)
def _get_dataset(session: Session, dataset_ref: DatasetRef) -> Dataset | None:
"""Shared SQL for implementation adapters that already own a transaction."""
return session.scalar(
select(Dataset)
.where(
Dataset.id == dataset_ref.dataset_id,
Dataset.tenant_id == dataset_ref.tenant_id,
)
.limit(1)
)
class SQLAlchemyDatasetRepository:
"""Own SQL access to tenant-scoped dataset state."""
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
def get_estimate_record(self, dataset_ref: DatasetRef) -> DatasetEstimateRecord | None:
with self._session_factory() as session:
dataset = _get_dataset(session, dataset_ref)
return _dataset_estimate_record(dataset) if dataset is not None else None
def is_notion_dataset(self, dataset_ref: DatasetRef) -> bool:
"""Return whether the owned dataset is configured for Notion imports."""
with self._session_factory() as session:
return bool(
session.scalar(
select(
exists().where(
Dataset.id == dataset_ref.dataset_id,
Dataset.tenant_id == dataset_ref.tenant_id,
Dataset.data_source_type == "notion_import",
)
)
)
)
def get_access_snapshot(
self,
*,
workspace_id: str,
dataset_id: str,
actor_id: str,
) -> DatasetAccessSnapshot | None:
with self._session_factory() as session:
dataset = _get_dataset(session, DatasetRef(workspace_id, dataset_id))
if dataset is None:
return None
actor_has_partial_access = bool(
session.scalar(
select(
exists().where(
DatasetPermission.tenant_id == workspace_id,
DatasetPermission.dataset_id == dataset_id,
DatasetPermission.account_id == actor_id,
DatasetPermission.has_permission.is_(True),
)
)
)
)
return DatasetAccessSnapshot(
dataset=_dataset_access_record(dataset),
actor_has_partial_access=actor_has_partial_access,
)
def get_segment_scope(
self,
*,
dataset_ref: DatasetRef,
document_id: str,
actor_id: str | None,
) -> SegmentScope:
"""Read ownership, model configuration and the actor grant in one snapshot."""
grant = (
exists().where(
DatasetPermission.tenant_id == dataset_ref.tenant_id,
DatasetPermission.dataset_id == dataset_ref.dataset_id,
DatasetPermission.account_id == actor_id,
DatasetPermission.has_permission.is_(True),
)
if actor_id is not None
else literal(False)
)
with self._session_factory() as session:
row = session.execute(
select(Dataset, Document, grant)
.outerjoin(
Document,
and_(
Document.id == document_id,
Document.dataset_id == Dataset.id,
Document.tenant_id == Dataset.tenant_id,
),
)
.where(Dataset.id == dataset_ref.dataset_id, Dataset.tenant_id == dataset_ref.tenant_id)
).one_or_none()
if row is None:
return SegmentScope(None, None)
dataset, document, actor_has_access = row
return SegmentScope(
dataset=SegmentDatasetRecord(
id=dataset.id,
workspace_id=dataset.tenant_id,
indexing_technique=str(dataset.indexing_technique),
embedding_model_provider=dataset.embedding_model_provider,
embedding_model=dataset.embedding_model,
),
document=SegmentDocumentRecord(
id=document.id,
dataset_id=document.dataset_id,
workspace_id=document.tenant_id,
doc_form=str(document.doc_form),
)
if document is not None
else None,
access_snapshot=DatasetAccessSnapshot(
dataset=_dataset_access_record(dataset),
actor_has_partial_access=bool(actor_has_access),
)
if actor_id is not None
else None,
)
api/repositories/knowledge/document_repository.py
+241
View File
@@ -0,0 +1,241 @@
"""SQLAlchemy repository for tenant-owned document state."""
import json
from collections.abc import Mapping
from pydantic import TypeAdapter, ValidationError
from sqlalchemy import exists, select, update
from sqlalchemy.orm import Session, sessionmaker
from libs.datetime_utils import naive_utc_now
from models import Account
from models.dataset import Dataset, DatasetProcessRule, Document, Pipeline
from models.enums import IndexingStatus
from repositories.knowledge.dataset_read_repository import get_pipeline_dataset
from services.knowledge.document_sync import SyncDocumentRecord
from services.knowledge.indexing.errors import DocumentIsDeletedPausedError, DocumentIsPausedError
from services.knowledge.indexing.estimate import EstimateDocumentRecord, StoredSource
from services.knowledge.indexing.execution import IndexingDocument
from services.knowledge.resource_scope import DatasetRef, DocumentRef
_MAPPING_ADAPTER = TypeAdapter(dict[str, object])
def _mapping(value: object) -> Mapping[str, object] | None:
if value is None:
return None
try:
if isinstance(value, str):
return _MAPPING_ADAPTER.validate_python(json.loads(value))
return _MAPPING_ADAPTER.validate_python(value)
except (TypeError, ValueError, ValidationError):
return {}
def _estimate_document_record(document: Document) -> EstimateDocumentRecord:
return EstimateDocumentRecord(
id=document.id,
workspace_id=document.tenant_id,
dataset_id=document.dataset_id,
data_source_type=str(document.data_source_type),
data_source_info=_mapping(document.data_source_info),
indexing_status=str(document.indexing_status),
doc_form=str(document.doc_form),
doc_language=document.doc_language,
dataset_process_rule_id=document.dataset_process_rule_id,
)
def _get_document(session: Session, document_ref: DocumentRef) -> Document | None:
"""Shared SQL for implementation adapters that already own a transaction."""
return session.scalar(
select(Document)
.where(
Document.id == document_ref.document_id,
Document.dataset_id == document_ref.dataset.dataset_id,
Document.tenant_id == document_ref.dataset.tenant_id,
)
.limit(1)
)
def require_indexing_document(session: Session, ref: DocumentRef, *, lock: bool = False) -> Document:
"""Validate the persisted owner chain before each indexing phase."""
statement = (
select(Document)
.join(Dataset, Dataset.id == Document.dataset_id)
.where(
Document.id == ref.document_id,
Document.dataset_id == ref.dataset.dataset_id,
Document.tenant_id == ref.dataset.tenant_id,
Dataset.tenant_id == ref.dataset.tenant_id,
)
)
if lock:
statement = statement.with_for_update(of=Document)
document = session.scalar(statement)
if document is None:
raise DocumentIsDeletedPausedError()
if document.is_paused:
raise DocumentIsPausedError()
return document
class SQLAlchemyDocumentRepository:
"""Own SQL access to documents through their complete owner chain."""
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
def get_pipeline_dataset(self, pipeline: Pipeline, *, session: Session) -> Dataset | None:
"""Read within the pipeline caller's transaction, including uncommitted changes."""
return get_pipeline_dataset(pipeline, session=session)
def get_indexing_document(self, ref: DocumentRef) -> IndexingDocument:
with self._session_factory() as session:
document = require_indexing_document(session, ref)
rule = session.scalar(
select(DatasetProcessRule).where(
DatasetProcessRule.id == document.dataset_process_rule_id,
DatasetProcessRule.dataset_id == ref.dataset.dataset_id,
)
)
return IndexingDocument(
source=StoredSource.from_document(_estimate_document_record(document)),
processing_rule=rule.to_dict() if rule else {},
doc_language=document.doc_language or "English",
need_summary=bool(document.need_summary),
)
def get_indexing_models(self, ref: DocumentRef) -> tuple[Dataset, Document]:
"""Detached legacy processor inputs; never pass them to another thread."""
with self._session_factory() as session:
document = require_indexing_document(session, ref)
dataset = session.scalar(
select(Dataset).where(
Dataset.id == ref.dataset.dataset_id,
Dataset.tenant_id == ref.dataset.tenant_id,
)
)
assert dataset is not None
return dataset, document
def get_indexing_user(self, ref: DocumentRef) -> Account:
with self._session_factory() as session:
document = require_indexing_document(session, ref)
account = session.get(Account, document.created_by)
if account is None:
raise ValueError("no current user found")
account.set_tenant_id_with_session(ref.dataset.tenant_id, session=session)
return account
def mark_splitting(self, ref: DocumentRef) -> None:
with self._session_factory.begin() as session:
document = require_indexing_document(session, ref, lock=True)
document.indexing_status = IndexingStatus.SPLITTING
document.parsing_completed_at = naive_utc_now()
def complete_indexing(self, ref: DocumentRef, *, tokens: int, latency: float) -> None:
with self._session_factory.begin() as session:
document = require_indexing_document(session, ref, lock=True)
document.indexing_status = IndexingStatus.COMPLETED
document.tokens = tokens
document.indexing_latency = latency
document.completed_at = naive_utc_now()
document.error = None
def fail_indexing(self, ref: DocumentRef, error: str) -> None:
try:
with self._session_factory.begin() as session:
document = require_indexing_document(session, ref, lock=True)
document.indexing_status = IndexingStatus.ERROR
document.error = error
document.stopped_at = naive_utc_now()
except DocumentIsDeletedPausedError:
return
def exists(self, *, workspace_id: str, dataset_id: str, document_id: str) -> bool:
with self._session_factory() as session:
return bool(
session.scalar(
select(
exists().where(
Document.tenant_id == workspace_id,
Document.dataset_id == dataset_id,
Document.id == document_id,
)
)
)
)
def mark_failed(self, *, workspace_id: str, dataset_id: str, document_id: str, error: str) -> None:
with self._session_factory.begin() as session:
session.execute(
update(Document)
.where(
Document.tenant_id == workspace_id,
Document.dataset_id == dataset_id,
Document.id == document_id,
)
.values(indexing_status=IndexingStatus.ERROR, error=error)
)
def get_estimate_document(self, document_ref: DocumentRef) -> EstimateDocumentRecord | None:
with self._session_factory() as session:
document = _get_document(session, document_ref)
return _estimate_document_record(document) if document is not None else None
def list_estimate_documents_by_batch(
self, dataset_ref: DatasetRef, batch: str
) -> tuple[EstimateDocumentRecord, ...]:
with self._session_factory() as session:
documents = session.scalars(
select(Document).where(
Document.tenant_id == dataset_ref.tenant_id,
Document.dataset_id == dataset_ref.dataset_id,
Document.batch == batch,
)
).all()
return tuple(_estimate_document_record(document) for document in documents)
def get_sync_document(self, document_ref: DocumentRef) -> SyncDocumentRecord | None:
with self._session_factory() as session:
document = _get_document(session, document_ref)
if document is None:
return None
return SyncDocumentRecord(id=document.id, data_source_type=str(document.data_source_type))
def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]:
with self._session_factory() as session:
return tuple(
dataset_ref.document(document_id)
for document_id in session.scalars(
select(Document.id).where(
Document.tenant_id == dataset_ref.tenant_id,
Document.dataset_id == dataset_ref.dataset_id,
Document.data_source_type == "notion_import",
Document.enabled.is_(True),
Document.archived.is_(False),
)
).all()
)
def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]:
page_ids: set[str] = set()
with self._session_factory() as session:
documents = session.scalars(
select(Document).where(
Document.tenant_id == dataset_ref.tenant_id,
Document.dataset_id == dataset_ref.dataset_id,
Document.data_source_type == "notion_import",
Document.enabled.is_(True),
)
).all()
for document in documents:
source_info = _mapping(document.data_source_info)
if source_info is None:
continue
page_id = source_info.get("notion_page_id")
if isinstance(page_id, str) and page_id:
page_ids.add(page_id)
return frozenset(page_ids)
api/repositories/knowledge/keyword_table_repository.py
+121
View File
@@ -0,0 +1,121 @@
"""Keyword table serialization, object storage, and tenant-scoped SQL persistence."""
import json
from collections.abc import Mapping, Sequence
from typing import Any, Protocol
import orjson
from sqlalchemy import select
from sqlalchemy.orm import Session
from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment
from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table
class _KeywordStorage(Protocol):
def load_once(self, key: str, /) -> bytes: ...
def exists(self, key: str, /) -> bool: ...
def save(self, key: str, data: bytes, /) -> None: ...
def delete(self, key: str, /) -> None: ...
def keyword_file_key(tenant_id: str, dataset_id: str) -> str:
return f"keyword_files/{tenant_id}/{dataset_id}.txt"
def _decode_sets(values: dict[str, Any]) -> dict[str, Any]:
return {key: set(value) if isinstance(value, list) else value for key, value in values.items()}
def _encode_sets(value: object) -> list[object]:
if isinstance(value, set):
return list(value)
raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
def load_keyword_table(
*, storage: _KeywordStorage, tenant_id: str, dataset_id: str, storage_type: str, data: str | None
) -> dict[str, Any] | None:
if storage_type == "database":
return json.loads(data, object_hook=_decode_sets) if data else None
file_key = keyword_file_key(tenant_id, dataset_id)
try:
payload = storage.load_once(file_key)
return json.loads(payload.decode("utf-8"), object_hook=_decode_sets) if payload else None
except FileNotFoundError:
return None
def load_dataset_keyword_table(
dataset: Dataset, *, storage: _KeywordStorage, session: Session
) -> dict[str, Any] | None:
"""Resolve a dataset's stored keyword payload using the caller's session."""
row = get_dataset_keyword_table(dataset, session=session)
if row is None:
return None
return load_keyword_table(
storage=storage,
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
storage_type=row.data_source_type,
data=row.keyword_table,
)
def save_keyword_table(
*, storage: _KeywordStorage, tenant_id: str, dataset_id: str, storage_type: str, table: dict[str, set[str]]
) -> str:
data = orjson.dumps(
{"__type__": "keyword_table", "__data__": {"index_id": dataset_id, "summary": None, "table": table}},
default=_encode_sets,
)
if storage_type != "database":
file_key = keyword_file_key(tenant_id, dataset_id)
if storage.exists(file_key):
storage.delete(file_key)
storage.save(file_key, data)
return data.decode("utf-8")
def persist_keyword_table(
session: Session,
*,
tenant_id: str,
dataset_id: str,
storage_type: str,
data: str,
keywords: Mapping[str, Sequence[str]],
) -> None:
"""Write a tenant-owned table and segment keywords in the caller's transaction."""
result = session.execute(
select(Dataset.id, DatasetKeywordTable)
.outerjoin(DatasetKeywordTable, DatasetKeywordTable.dataset_id == Dataset.id)
.where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id)
).one_or_none()
if result is None:
raise LookupError("Dataset no longer exists")
_, row = result
if row is None:
row = DatasetKeywordTable(
dataset_id=dataset_id,
data_source_type=storage_type,
keyword_table=data if storage_type == "database" else "",
)
session.add(row)
elif storage_type == "database":
row.keyword_table = data
if keywords:
segments = session.scalars(
select(DocumentSegment).where(
DocumentSegment.tenant_id == tenant_id,
DocumentSegment.dataset_id == dataset_id,
DocumentSegment.index_node_id.in_(keywords),
)
)
for segment in segments:
assert segment.index_node_id is not None
segment.keywords = list(keywords[segment.index_node_id])
session.flush()
api/repositories/knowledge/metadata_repository.py
+212
View File
@@ -0,0 +1,212 @@
"""Tenant-scoped metadata operations; each mutation owns one atomic transaction."""
from collections.abc import Sequence
from sqlalchemy import func, select
from sqlalchemy.orm import Session, sessionmaker
from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource
from libs.datetime_utils import naive_utc_now
from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document
from repositories.knowledge.dataset_read_repository import get_document_uploader
from services.errors.metadata import MetadataResourceNotFoundError
from services.knowledge.metadata.application import MetadataList, MetadataOperation, MetadataRecord
from services.knowledge.resource_scope import DatasetRef
class SQLAlchemyMetadataRepository:
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._sessions = session_factory
@staticmethod
def _dataset(session: Session, ref: DatasetRef, *, write: bool = False) -> Dataset:
query = select(Dataset).where(Dataset.id == ref.dataset_id, Dataset.tenant_id == ref.tenant_id)
if write:
query = query.with_for_update()
dataset = session.scalar(query)
if dataset is None:
raise MetadataResourceNotFoundError("Dataset not found.")
return dataset
@staticmethod
def _fields(ref: DatasetRef):
return select(DatasetMetadata).where(
DatasetMetadata.tenant_id == ref.tenant_id, DatasetMetadata.dataset_id == ref.dataset_id
)
@staticmethod
def _documents(ref: DatasetRef):
return select(Document).where(Document.tenant_id == ref.tenant_id, Document.dataset_id == ref.dataset_id)
@staticmethod
def _bindings(ref: DatasetRef):
return select(DatasetMetadataBinding).where(
DatasetMetadataBinding.tenant_id == ref.tenant_id, DatasetMetadataBinding.dataset_id == ref.dataset_id
)
@staticmethod
def _record(row: DatasetMetadata) -> MetadataRecord:
return MetadataRecord(
id=row.id,
name=row.name,
type=row.type,
tenant_id=row.tenant_id,
dataset_id=row.dataset_id,
created_by=row.created_by,
created_at=row.created_at,
updated_by=row.updated_by,
updated_at=row.updated_at,
)
@classmethod
def _check_name(cls, session: Session, ref: DatasetRef, name: str) -> None:
if session.scalar(cls._fields(ref).where(DatasetMetadata.name == name).limit(1)) is not None:
raise ValueError("Metadata name already exists.")
def create(self, ref: DatasetRef, *, name: str, field_type: str, actor_id: str) -> MetadataRecord:
with self._sessions.begin() as session:
self._dataset(session, ref, write=True)
self._check_name(session, ref, name)
row = DatasetMetadata(
tenant_id=ref.tenant_id, dataset_id=ref.dataset_id, name=name, type=field_type, created_by=actor_id
)
session.add(row)
session.flush()
return self._record(row)
def rename(self, ref: DatasetRef, metadata_id: str, *, name: str, actor_id: str) -> MetadataRecord:
with self._sessions.begin() as session:
self._dataset(session, ref, write=True)
self._check_name(session, ref, name)
row = session.scalar(self._fields(ref).where(DatasetMetadata.id == metadata_id))
if row is None:
raise MetadataResourceNotFoundError("Metadata not found.")
old_name = row.name
row.name, row.updated_by, row.updated_at = name, actor_id, naive_utc_now()
bindings = session.scalars(
self._bindings(ref).where(DatasetMetadataBinding.metadata_id == metadata_id)
).all()
for document in session.scalars(
self._documents(ref).where(Document.id.in_([b.document_id for b in bindings]))
):
values = dict(document.doc_metadata or {})
values[name] = values.pop(old_name, None)
document.doc_metadata = values
return self._record(row)
def delete(self, ref: DatasetRef, metadata_id: str) -> MetadataRecord:
with self._sessions.begin() as session:
self._dataset(session, ref, write=True)
row = session.scalar(self._fields(ref).where(DatasetMetadata.id == metadata_id))
if row is None:
raise MetadataResourceNotFoundError("Metadata not found.")
bindings = session.scalars(
self._bindings(ref).where(DatasetMetadataBinding.metadata_id == metadata_id)
).all()
for document in session.scalars(
self._documents(ref).where(Document.id.in_([b.document_id for b in bindings]))
):
values = dict(document.doc_metadata or {})
values.pop(row.name, None)
document.doc_metadata = values
for binding in bindings:
session.delete(binding)
session.delete(row)
return self._record(row)
@staticmethod
def _built_in(session: Session, document: Document) -> dict[str, str | int | float | None]:
return {
BuiltInField.document_name: document.name,
BuiltInField.uploader: get_document_uploader(document, session=session),
BuiltInField.upload_date: document.upload_date.timestamp(),
BuiltInField.last_update_date: document.last_update_date.timestamp(),
BuiltInField.source: MetadataDataSource[document.data_source_type],
}
def set_built_in(self, ref: DatasetRef, *, enabled: bool) -> None:
with self._sessions.begin() as session:
dataset = self._dataset(session, ref, write=True)
if dataset.built_in_field_enabled == enabled:
return
for document in session.scalars(
self._documents(ref).where(
Document.enabled.is_(True), Document.archived.is_(False), Document.indexing_status == "completed"
)
):
values = dict(document.doc_metadata or {})
if enabled:
values.update(self._built_in(session, document))
else:
for key in BuiltInField:
values.pop(key, None)
document.doc_metadata = values
dataset.built_in_field_enabled = enabled
def update_documents(self, ref: DatasetRef, operations: Sequence[MetadataOperation], *, actor_id: str) -> None:
with self._sessions.begin() as session:
dataset = self._dataset(session, ref, write=True)
metadata_ids = {value.id for operation in operations for value in operation.metadata_list}
fields = {
row.id: row for row in session.scalars(self._fields(ref).where(DatasetMetadata.id.in_(metadata_ids)))
}
if metadata_ids != set(fields):
raise MetadataResourceNotFoundError("Metadata not found.")
document_ids = {operation.document_id for operation in operations}
documents = {
row.id: row
for row in session.scalars(self._documents(ref).where(Document.id.in_(document_ids)).with_for_update())
}
if document_ids != set(documents):
raise MetadataResourceNotFoundError("Document not found.")
for operation in operations:
document = documents[operation.document_id]
values = dict(document.doc_metadata or {}) if operation.partial_update else {}
values.update({fields[value.id].name: value.value for value in operation.metadata_list})
if dataset.built_in_field_enabled:
values.update(self._built_in(session, document))
document.doc_metadata = values
bindings = session.scalars(
self._bindings(ref).where(DatasetMetadataBinding.document_id == document.id)
).all()
previous = {binding.metadata_id for binding in bindings}
if not operation.partial_update:
for binding in bindings:
session.delete(binding)
previous.clear()
for value in operation.metadata_list:
if value.id in previous:
continue
session.add(
DatasetMetadataBinding(
tenant_id=ref.tenant_id,
dataset_id=ref.dataset_id,
document_id=document.id,
metadata_id=value.id,
created_by=actor_id,
)
)
previous.add(value.id)
def list_fields(self, ref: DatasetRef) -> MetadataList:
with self._sessions() as session:
dataset = self._dataset(session, ref)
counts = dict(
session.execute(
select(DatasetMetadataBinding.metadata_id, func.count())
.where(
DatasetMetadataBinding.tenant_id == ref.tenant_id,
DatasetMetadataBinding.dataset_id == ref.dataset_id,
)
.group_by(DatasetMetadataBinding.metadata_id)
)
.tuples()
.all()
)
return {
"doc_metadata": [
{"id": row.id, "name": row.name, "type": row.type, "count": counts.get(row.id, 0)}
for row in session.scalars(self._fields(ref))
],
"built_in_field_enabled": dataset.built_in_field_enabled,
}
api/repositories/knowledge/segment_read_adapter.py
+87
View File
@@ -0,0 +1,87 @@
"""Read adapter for authorized segment content and attachment preview values."""
import base64
import hashlib
import hmac
import os
import re
import time
from collections.abc import Sequence
from typing import Literal
from sqlalchemy.orm import Session
from configs import dify_config
from core.tools.signature import bind_file_uri, sign_tool_file_uri
from models.dataset import AttachmentItem, DocumentSegment
from models.model import UploadFile
from repositories.knowledge.dataset_read_repository import get_segment_attachment_files, get_segment_content_file_ids
# Consume the entire old query string and replace matches in source order.
_FILE_ID = r"[a-fA-F0-9]{8}(?:-[a-fA-F0-9]{4}){3}-[a-fA-F0-9]{12}"
_FILE_LINK = re.compile(
rf"/files/(?:(?P<upload>{_FILE_ID})/(?P<preview>image-preview|file-preview)"
rf"|tools/(?P<tool>{_FILE_ID})\.(?P<extension>[a-zA-Z0-9]+))"
r"(?:\?[^\s)\"'<>]*)?"
)
def _sign_upload_file_uri(upload_file_id: str, *, preview_kind: Literal["image", "file"] = "file") -> str:
"""Sign an upload ID whose ownership was checked by this read adapter."""
timestamp = str(int(time.time()))
nonce = os.urandom(16).hex()
payload = f"{preview_kind}-preview|{upload_file_id}|{timestamp}|{nonce}"
signature = hmac.new(dify_config.SECRET_KEY.encode(), payload.encode(), hashlib.sha256).digest()
encoded_sign = base64.urlsafe_b64encode(signature).decode()
return f"/files/{upload_file_id}/{preview_kind}-preview?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
def sign_segment_content(segment: DocumentSegment, *, session: Session) -> str:
"""Refresh links only after checking persisted file ownership.
The segment must come from an authorized knowledge read. File queries also
enforce its tenant and the current execution's file-access scope. Inaccessible
references lose their old signature and never receive a new one.
"""
matches = list(_FILE_LINK.finditer(segment.content))
if not matches:
return segment.content
upload_ids = {match["upload"].lower() for match in matches if match["upload"]}
tool_ids = {match["tool"].lower() for match in matches if match["tool"]}
allowed_uploads, allowed_tools = get_segment_content_file_ids(
segment, upload_ids=upload_ids, tool_ids=tool_ids, session=session
)
def replace(match: re.Match[str]) -> str:
if upload_id := match["upload"]:
upload_id = upload_id.lower()
if upload_id not in allowed_uploads:
return match.group().split("?", 1)[0]
return _sign_upload_file_uri(
upload_id, preview_kind="image" if match["preview"] == "image-preview" else "file"
)
tool_id = match["tool"].lower()
if tool_id not in allowed_tools:
return match.group().split("?", 1)[0]
return sign_tool_file_uri(tool_id, "." + match["extension"])
return _FILE_LINK.sub(replace, segment.content)
def format_segment_attachments(files: Sequence[UploadFile]) -> list[AttachmentItem]:
base_url = dify_config.FILES_URL or dify_config.CONSOLE_API_URL or ""
return [
{
"id": file.id,
"name": file.name,
"size": file.size,
"extension": file.extension,
"mime_type": file.mime_type,
"source_url": bind_file_uri(_sign_upload_file_uri(file.id, preview_kind="image"), base_url),
}
for file in files
]
def get_segment_attachments(segment: DocumentSegment, *, session: Session) -> list[AttachmentItem]:
return format_segment_attachments(get_segment_attachment_files(segment, session=session))
api/repositories/knowledge/segment_repository.py
+716
View File
@@ -0,0 +1,716 @@
"""Infrastructure adapters for dataset segment use cases."""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import cast
from sqlalchemy import String, case, delete, func, literal, or_, select, update
from sqlalchemy import cast as sql_cast
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.docstore.dataset_docstore import DatasetDocumentStore
from core.rag.index_processor.constant.index_type import IndexStructureType
from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor
from core.rag.models.document import ChildDocument
from core.rag.models.document import Document as IndexDocument
from graphon.file import File
from libs.datetime_utils import naive_utc_now
from libs.helper import escape_like_pattern
from libs.pagination import PaginatedResult, paginate_query
from models.dataset import (
ChildChunk,
Dataset,
DatasetProcessRule,
Document,
DocumentSegment,
DocumentSegmentSummary,
SegmentAttachmentBinding,
)
from models.enums import IndexingStatus, SegmentStatus, SegmentType, SummaryStatus
from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table, get_segment_child_chunks
from repositories.knowledge.dataset_repository import _get_dataset
from repositories.knowledge.document_repository import _get_document, require_indexing_document
from repositories.knowledge.keyword_table_repository import persist_keyword_table
from repositories.knowledge.segment_read_adapter import get_segment_attachments, sign_segment_content
from services.knowledge.entities.segments import ChildChunkRecord, SegmentRecord
from services.knowledge.indexing.execution import IndexingDocument
from services.knowledge.resource_scope import DatasetRef, DocumentRef, SegmentRef
from services.knowledge.segments.application import (
ChildChunkListFilter,
ChildChunkPage,
ChildChunkState,
SegmentDetail,
SegmentIndexTarget,
SegmentListFilter,
SegmentPage,
SegmentUpdateState,
)
class SQLAlchemySegmentRepository:
"""Read and write scoped segment rows in bounded database transactions."""
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
def clear_for_indexing(self, document: IndexingDocument) -> None:
ref = document.ref
with self._session_factory.begin() as session:
require_indexing_document(session, ref, lock=True)
if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
session.execute(
delete(ChildChunk).where(
ChildChunk.tenant_id == ref.dataset.tenant_id,
ChildChunk.dataset_id == ref.dataset.dataset_id,
ChildChunk.document_id == ref.document_id,
)
)
session.execute(delete(DocumentSegment).where(*self._indexing_scope(ref)))
def save_for_indexing(
self, document: IndexingDocument, chunks: list[IndexDocument], token_counts: list[int]
) -> None:
ref = document.ref
with self._session_factory.begin() as session:
row = require_indexing_document(session, ref, lock=True)
dataset = _get_dataset(session, ref.dataset)
assert dataset is not None
DatasetDocumentStore(dataset=dataset, user_id=row.created_by, document_id=row.id).add_documents(
session=session,
docs=chunks,
token_counts=token_counts,
save_child=document.doc_form == IndexStructureType.PARENT_CHILD_INDEX,
)
now = naive_utc_now()
row.indexing_status = IndexingStatus.INDEXING
row.cleaning_completed_at = now
row.splitting_completed_at = now
row.word_count = sum(len(chunk.page_content) for chunk in chunks)
session.execute(
update(DocumentSegment)
.where(*self._indexing_scope(ref))
.values(
status=SegmentStatus.INDEXING,
indexing_at=now,
)
)
def resume_indexing(self, document: IndexingDocument) -> tuple[list[IndexDocument], int]:
ref = document.ref
with self._session_factory() as session:
require_indexing_document(session, ref)
rows = session.scalars(
select(DocumentSegment).where(*self._indexing_scope(ref)).order_by(DocumentSegment.position)
).all()
chunks = []
for row in rows:
if row.status == SegmentStatus.COMPLETED:
continue
chunk = IndexDocument(
page_content=row.content,
metadata={
"doc_id": row.index_node_id,
"doc_hash": row.index_node_hash,
"document_id": ref.document_id,
"dataset_id": ref.dataset.dataset_id,
},
)
if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
chunk.children = [
ChildDocument(
page_content=child.content,
metadata={
"doc_id": child.index_node_id,
"doc_hash": child.index_node_hash,
"document_id": ref.document_id,
"dataset_id": ref.dataset.dataset_id,
},
)
for child in get_segment_child_chunks(row, session=session)
]
chunks.append(chunk)
return chunks, sum(row.tokens for row in rows)
def complete_indexing_segments(self, ref: DocumentRef, node_ids: Sequence[str]) -> None:
with self._session_factory.begin() as session:
require_indexing_document(session, ref, lock=True)
session.execute(
update(DocumentSegment)
.where(
*self._indexing_scope(ref),
DocumentSegment.index_node_id.in_(node_ids),
DocumentSegment.status == SegmentStatus.INDEXING,
)
.values(status=SegmentStatus.COMPLETED, enabled=True, completed_at=naive_utc_now())
)
@staticmethod
def _indexing_scope(ref: DocumentRef):
return (
DocumentSegment.tenant_id == ref.dataset.tenant_id,
DocumentSegment.dataset_id == ref.dataset.dataset_id,
DocumentSegment.document_id == ref.document_id,
)
def list_segments(self, document_ref: DocumentRef, query: SegmentListFilter) -> SegmentPage:
with self._session_factory() as session:
statement = (
select(DocumentSegment)
.where(
DocumentSegment.tenant_id == document_ref.dataset.tenant_id,
DocumentSegment.dataset_id == document_ref.dataset.dataset_id,
DocumentSegment.document_id == document_ref.document_id,
)
.order_by(DocumentSegment.position.asc())
)
if query.statuses:
statement = statement.where(DocumentSegment.status.in_(query.statuses))
if query.hit_count_gte is not None:
statement = statement.where(DocumentSegment.hit_count >= query.hit_count_gte)
if query.keyword:
escaped_keyword = escape_like_pattern(query.keyword)
if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql":
keywords_jsonb = sql_cast(DocumentSegment.keywords, JSONB)
keywords_array = case(
(func.jsonb_typeof(keywords_jsonb) == "array", keywords_jsonb),
else_=sql_cast(literal("[]"), JSONB),
)
keywords_condition = func.array_to_string(
func.array(
select(func.jsonb_array_elements_text(keywords_array))
.correlate(DocumentSegment)
.scalar_subquery()
),
",",
).ilike(f"%{escaped_keyword}%", escape="\\")
else:
keywords_condition = sql_cast(DocumentSegment.keywords, String).ilike(
f"%{escaped_keyword}%", escape="\\"
)
statement = statement.where(
or_(
DocumentSegment.content.ilike(f"%{escaped_keyword}%", escape="\\"),
keywords_condition,
)
)
if query.enabled.lower() == "true":
statement = statement.where(DocumentSegment.enabled.is_(True))
elif query.enabled.lower() == "false":
statement = statement.where(DocumentSegment.enabled.is_(False))
limit = min(query.limit, 100)
segments = paginate_query(statement, session=session, page=query.page, per_page=limit, max_per_page=100)
items = list(segments.items)
summaries = (
{
summary.chunk_id: summary
for summary in session.scalars(
select(DocumentSegmentSummary).where(
DocumentSegmentSummary.chunk_id.in_([segment.id for segment in items]),
DocumentSegmentSummary.dataset_id == document_ref.dataset.dataset_id,
DocumentSegmentSummary.document_id == document_ref.document_id,
)
)
}
if items
else {}
)
return SegmentPage(
items=tuple(
_segment_data(
segment,
summaries[segment.id].summary_content if segment.id in summaries else None,
session,
)
for segment in items
),
total=segments.total,
total_pages=segments.pages,
page=query.page,
limit=limit,
)
def list_child_chunks(
self,
segment_ref: SegmentRef,
query: ChildChunkListFilter,
) -> ChildChunkPage | None:
with self._session_factory() as session:
if _get_segment(session, segment_ref) is None:
return None
limit = min(query.limit, 100)
child_chunks = query_child_chunks(
session,
segment_ref,
page=query.page,
limit=limit,
keyword=query.keyword,
)
return ChildChunkPage(
items=tuple(_child_chunk_data(chunk) for chunk in child_chunks.items),
total=child_chunks.total,
total_pages=child_chunks.pages,
page=query.page,
limit=limit,
)
def get_segment(self, segment_ref: SegmentRef) -> SegmentDetail | None:
with self._session_factory() as session:
segment = _get_segment(session, segment_ref)
if segment is None:
return None
_, document = _require_scope_models(session, segment_ref.document)
return _detail(session, segment, document)
def get_segment_update_state(self, segment_ref: SegmentRef) -> SegmentUpdateState | None:
with self._session_factory() as session:
segment = _get_segment(session, segment_ref)
if segment is None:
return None
# Legacy rows may contain scalar or malformed JSON keywords. Treat
# those as absent so a write can repair them without building a DTO.
keywords = segment.keywords
return SegmentUpdateState(
content=segment.content,
enabled=segment.enabled,
keywords=tuple(keywords)
if isinstance(keywords, list) and all(isinstance(keyword, str) for keyword in keywords)
else None,
)
def get_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str]) -> tuple[SegmentIndexTarget, ...]:
if not segment_ids:
return ()
with self._session_factory() as session:
segments = session.scalars(
select(DocumentSegment).where(
DocumentSegment.tenant_id == document_ref.dataset.tenant_id,
DocumentSegment.dataset_id == document_ref.dataset.dataset_id,
DocumentSegment.document_id == document_ref.document_id,
DocumentSegment.id.in_(segment_ids),
)
).all()
return tuple(SegmentIndexTarget(segment.id, segment.index_node_id, segment.enabled) for segment in segments)
def get_children_for_segments(
self, document_ref: DocumentRef, segment_ids: Sequence[str]
) -> tuple[ChildChunkState, ...]:
if not segment_ids:
return ()
with self._session_factory() as session:
return tuple(_child_state(child) for child in session.scalars(_children_query(document_ref, segment_ids)))
def set_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str], values: Mapping[str, object]) -> None:
if not segment_ids:
return
with self._session_factory.begin() as session:
session.execute(
update(DocumentSegment)
.where(
DocumentSegment.tenant_id == document_ref.dataset.tenant_id,
DocumentSegment.dataset_id == document_ref.dataset.dataset_id,
DocumentSegment.document_id == document_ref.document_id,
DocumentSegment.id.in_(segment_ids),
)
.values(**values)
)
def save_segment(self, segment_ref: SegmentRef, values: Mapping[str, object], *, create: bool = False) -> None:
with self._session_factory.begin() as session:
_, document = _require_scope_models(session, segment_ref.document)
segment = _get_segment(session, segment_ref)
previous_words = segment.word_count if segment is not None else 0
if create:
if segment is not None:
raise ValueError("Segment already exists")
position = (
session.scalar(
select(func.max(DocumentSegment.position)).where(
DocumentSegment.tenant_id == segment_ref.document.dataset.tenant_id,
DocumentSegment.dataset_id == segment_ref.document.dataset.dataset_id,
DocumentSegment.document_id == segment_ref.document.document_id,
)
)
or 0
)
segment = DocumentSegment(
tenant_id=segment_ref.document.dataset.tenant_id,
dataset_id=segment_ref.document.dataset.dataset_id,
document_id=segment_ref.document.document_id,
position=position + 1,
content=cast(str, values["content"]),
word_count=cast(int, values["word_count"]),
tokens=cast(int, values["tokens"]),
created_by=cast(str, values["created_by"]),
)
segment.id = segment_ref.segment_id
session.add(segment)
if segment is None:
raise LookupError("Segment no longer exists")
for key, value in values.items():
setattr(segment, key, value)
if "word_count" in values:
document.word_count = max(0, (document.word_count or 0) + segment.word_count - previous_words)
def delete_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str]) -> None:
if not segment_ids:
return
with self._session_factory.begin() as session:
_, document = _require_scope_models(session, document_ref)
rows = session.scalars(
select(DocumentSegment).where(
DocumentSegment.tenant_id == document_ref.dataset.tenant_id,
DocumentSegment.dataset_id == document_ref.dataset.dataset_id,
DocumentSegment.document_id == document_ref.document_id,
DocumentSegment.id.in_(segment_ids),
)
).all()
document.word_count = max(0, (document.word_count or 0) - sum(row.word_count for row in rows))
for row in rows:
session.delete(row)
def get_children(self, segment_ref: SegmentRef) -> tuple[ChildChunkState, ...] | None:
with self._session_factory() as session:
if _get_segment(session, segment_ref) is None:
return None
children = session.scalars(
_children_query(segment_ref.document, (segment_ref.segment_id,)).order_by(ChildChunk.position.asc())
).all()
return tuple(_child_state(child) for child in children)
def save_children(
self,
segment_ref: SegmentRef,
*,
added: Sequence[ChildChunkState] = (),
updated: Sequence[ChildChunkState] = (),
deleted: Sequence[ChildChunkState] = (),
) -> None:
with self._session_factory.begin() as session:
if _get_segment(session, segment_ref) is None:
raise LookupError("Segment no longer exists")
for child in deleted:
row = _get_child_chunk(session, segment_ref, child.data.id)
if row is not None:
session.delete(row)
for child in updated:
row = _get_child_chunk(session, segment_ref, child.data.id)
if row is None:
raise LookupError("Child chunk no longer exists")
row.content = child.data.content
row.word_count = child.data.word_count
row.type = SegmentType(child.data.type)
row.updated_by = child.updated_by
row.updated_at = child.data.updated_at
for child in added:
row = ChildChunk(
**child.data.model_dump(exclude={"id", "created_at", "updated_at"}),
tenant_id=segment_ref.document.dataset.tenant_id,
dataset_id=segment_ref.document.dataset.dataset_id,
document_id=segment_ref.document.document_id,
index_node_id=child.index_node_id,
index_node_hash=child.index_node_hash,
created_by=child.created_by,
)
row.id = child.data.id
row.created_at = child.data.created_at
row.updated_at = child.data.updated_at
row.updated_by = child.updated_by
session.add(row)
def get_indexing_snapshot(self, segment_ref: SegmentRef) -> SegmentIndexingSnapshot:
return self.get_indexing_snapshots(segment_ref.document, (segment_ref.segment_id,))[0]
def get_indexing_snapshots(
self, document_ref: DocumentRef, segment_ids: Sequence[str]
) -> tuple[SegmentIndexingSnapshot, ...]:
"""Load an ordered batch and its owner chain in one bounded read session."""
if not segment_ids:
return ()
with self._session_factory() as session:
dataset, document = _require_scope_models(session, document_ref)
segments = {
segment.id: segment
for segment in session.scalars(
select(DocumentSegment).where(
DocumentSegment.tenant_id == dataset.tenant_id,
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.document_id == document.id,
DocumentSegment.id.in_(segment_ids),
)
)
}
if set(segments) != set(segment_ids):
raise LookupError("Segment no longer exists")
rule = (
session.scalar(
select(DatasetProcessRule).where(
DatasetProcessRule.id == document.dataset_process_rule_id,
DatasetProcessRule.dataset_id == dataset.id,
)
)
if document.dataset_process_rule_id
else None
)
vector_type = (
Vector.resolve_vector_type(dataset, session=session)
if dataset.indexing_technique == "high_quality"
else None
)
attachments: dict[str, list[str]] = {segment_id: [] for segment_id in segment_ids}
for binding in session.scalars(
select(SegmentAttachmentBinding).where(
SegmentAttachmentBinding.tenant_id == dataset.tenant_id,
SegmentAttachmentBinding.dataset_id == dataset.id,
SegmentAttachmentBinding.document_id == document.id,
SegmentAttachmentBinding.segment_id.in_(segment_ids),
)
):
attachments[binding.segment_id].append(binding.attachment_id)
summaries = {
summary.chunk_id: summary
for summary in session.scalars(
select(DocumentSegmentSummary).where(
DocumentSegmentSummary.dataset_id == dataset.id,
DocumentSegmentSummary.document_id == document.id,
DocumentSegmentSummary.chunk_id.in_(segment_ids),
)
)
}
return tuple(
SegmentIndexingSnapshot(
dataset,
document,
segments[segment_id],
rule,
vector_type,
summaries.get(segment_id),
tuple(attachments[segment_id]),
)
for segment_id in segment_ids
)
def replace_attachments(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> None:
with self._session_factory.begin() as session:
if _get_segment(session, segment_ref) is None:
raise LookupError("Segment no longer exists")
session.execute(
delete(SegmentAttachmentBinding).where(
SegmentAttachmentBinding.tenant_id == segment_ref.document.dataset.tenant_id,
SegmentAttachmentBinding.dataset_id == segment_ref.document.dataset.dataset_id,
SegmentAttachmentBinding.document_id == segment_ref.document.document_id,
SegmentAttachmentBinding.segment_id == segment_ref.segment_id,
)
)
session.add_all(
[
SegmentAttachmentBinding(
tenant_id=segment_ref.document.dataset.tenant_id,
dataset_id=segment_ref.document.dataset.dataset_id,
document_id=segment_ref.document.document_id,
segment_id=segment_ref.segment_id,
attachment_id=file_id,
)
for file_id in attachment_ids
]
)
def save_summary(
self, segment_ref: SegmentRef, content: str | None, *, error: str | None = None
) -> DocumentSegmentSummary | None:
with self._session_factory.begin() as session:
segment = _get_segment(session, segment_ref)
if segment is None:
raise LookupError("Segment no longer exists")
summary = _summary(session, segment)
if content is None:
if summary is not None:
session.delete(summary)
return None
if summary is None:
summary = DocumentSegmentSummary(
dataset_id=segment.dataset_id, document_id=segment.document_id, chunk_id=segment.id
)
session.add(summary)
summary.summary_content = content
summary.status = SummaryStatus.ERROR if error else SummaryStatus.GENERATING
summary.error = error
session.flush()
session.expunge(summary)
return summary
def get_summary_images(self, segment_ref: SegmentRef) -> list[File]:
with self._session_factory() as session:
segment = _get_segment(session, segment_ref)
if segment is None:
raise LookupError("Segment no longer exists")
images = ParagraphIndexProcessor._extract_images_from_segment_attachments(
segment.tenant_id, segment.id, session
)
return images or ParagraphIndexProcessor._extract_images_from_text(
segment.tenant_id, segment.content, session
)
def get_keyword_table(self, dataset_ref: DatasetRef) -> tuple[str, str | None]:
with self._session_factory() as session:
dataset = _get_dataset(session, dataset_ref)
if dataset is None:
raise LookupError("Dataset no longer exists")
row = get_dataset_keyword_table(dataset, session=session)
return (row.data_source_type, row.keyword_table) if row else (dify_config.KEYWORD_DATA_SOURCE_TYPE, None)
def save_keyword_table(
self, dataset_ref: DatasetRef, *, storage_type: str, data: str, keywords: Mapping[str, Sequence[str]]
) -> None:
with self._session_factory.begin() as session:
persist_keyword_table(
session,
tenant_id=dataset_ref.tenant_id,
dataset_id=dataset_ref.dataset_id,
storage_type=storage_type,
data=data,
keywords=keywords,
)
@dataclass(frozen=True)
class SegmentIndexingSnapshot:
dataset: Dataset
document: Document
segment: DocumentSegment
process_rule: DatasetProcessRule | None
vector_type: str | None
summary: DocumentSegmentSummary | None
attachment_ids: tuple[str, ...]
def _summary(session: Session, segment: DocumentSegment) -> DocumentSegmentSummary | None:
return session.scalar(
select(DocumentSegmentSummary)
.where(
DocumentSegmentSummary.dataset_id == segment.dataset_id,
DocumentSegmentSummary.document_id == segment.document_id,
DocumentSegmentSummary.chunk_id == segment.id,
)
.limit(1)
)
def _children_query(document_ref: DocumentRef, segment_ids: Sequence[str]):
return select(ChildChunk).where(
ChildChunk.tenant_id == document_ref.dataset.tenant_id,
ChildChunk.dataset_id == document_ref.dataset.dataset_id,
ChildChunk.document_id == document_ref.document_id,
ChildChunk.segment_id.in_(segment_ids),
)
def _child_state(child: ChildChunk) -> ChildChunkState:
return ChildChunkState(
data=_child_chunk_data(child),
index_node_id=child.index_node_id,
index_node_hash=child.index_node_hash,
created_by=child.created_by,
updated_by=child.updated_by,
)
def _detail(session: Session, segment: DocumentSegment, document: Document) -> SegmentDetail:
summary = _summary(session, segment)
return SegmentDetail(
data=_segment_data(segment, summary.summary_content if summary else None, session),
doc_form=str(document.doc_form),
)
def _require_scope_models(session: Session, document_ref: DocumentRef) -> tuple[Dataset, Document]:
dataset = _get_dataset(session, document_ref.dataset)
document = _get_document(session, document_ref)
if dataset is None or document is None:
raise LookupError("Dataset document scope no longer exists")
return dataset, document
def _get_segment(session: Session, segment_ref: SegmentRef) -> DocumentSegment | None:
return session.scalar(
select(DocumentSegment)
.where(
DocumentSegment.id == segment_ref.segment_id,
DocumentSegment.document_id == segment_ref.document.document_id,
DocumentSegment.dataset_id == segment_ref.document.dataset.dataset_id,
DocumentSegment.tenant_id == segment_ref.document.dataset.tenant_id,
)
.limit(1)
)
def _get_child_chunk(session: Session, segment_ref: SegmentRef, child_chunk_id: str) -> ChildChunk | None:
return session.scalar(
select(ChildChunk)
.where(
ChildChunk.id == child_chunk_id,
ChildChunk.segment_id == segment_ref.segment_id,
ChildChunk.document_id == segment_ref.document.document_id,
ChildChunk.dataset_id == segment_ref.document.dataset.dataset_id,
ChildChunk.tenant_id == segment_ref.document.dataset.tenant_id,
)
.limit(1)
)
def _segment_data(segment: DocumentSegment, summary: str | None, session: Session) -> SegmentRecord:
return SegmentRecord.model_validate(
{
"id": segment.id,
"position": segment.position,
"document_id": segment.document_id,
"content": segment.content,
"answer": segment.answer,
"word_count": segment.word_count,
"tokens": segment.tokens,
"keywords": segment.keywords,
"index_node_id": segment.index_node_id,
"index_node_hash": segment.index_node_hash,
"hit_count": segment.hit_count,
"enabled": segment.enabled,
"disabled_at": segment.disabled_at,
"disabled_by": segment.disabled_by,
"status": segment.status,
"created_by": segment.created_by,
"created_at": segment.created_at,
"updated_at": segment.updated_at,
"updated_by": segment.updated_by,
"indexing_at": segment.indexing_at,
"completed_at": segment.completed_at,
"error": segment.error,
"stopped_at": segment.stopped_at,
"sign_content": sign_segment_content(segment, session=session),
"summary": summary,
"child_chunks": get_segment_child_chunks(segment, session=session, include_full_doc=False),
"attachments": get_segment_attachments(segment, session=session),
}
)
def _child_chunk_data(child_chunk: ChildChunk) -> ChildChunkRecord:
return ChildChunkRecord.model_validate(child_chunk)
def query_child_chunks(
session: Session,
segment_ref: SegmentRef,
*,
page: int,
limit: int,
keyword: str | None,
) -> PaginatedResult[ChildChunk]:
"""Share the query with callers that already own an explicit session."""
query = _children_query(segment_ref.document, (segment_ref.segment_id,)).order_by(ChildChunk.position.asc())
if keyword:
query = query.where(ChildChunk.content.ilike(f"%{escape_like_pattern(keyword)}%", escape="\\"))
return paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100)
api/repositories/knowledge/summary_repository.py
+74
View File
@@ -0,0 +1,74 @@
"""Persist summary indexing outcomes using one owned session per attempt."""
from collections.abc import Callable
from sqlalchemy import select
from sqlalchemy.orm import Session
from libs.datetime_utils import naive_utc_now
from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary
from models.enums import SummaryStatus
from services.knowledge.summaries.application import SummaryVectorInput, SummaryVectorResult
class SQLAlchemySummaryRepository:
def __init__(self, *, new_session: Callable[[], Session]) -> None:
self._new_session = new_session
@staticmethod
def _record(session: Session, command: SummaryVectorInput) -> DocumentSegmentSummary | None:
ref = command.ref
segment = session.scalar(
select(DocumentSegment)
.join(Document, Document.id == DocumentSegment.document_id)
.join(Dataset, Dataset.id == Document.dataset_id)
.where(
Dataset.id == ref.document.dataset.dataset_id,
Dataset.tenant_id == ref.document.dataset.tenant_id,
Document.id == ref.document.document_id,
Document.tenant_id == ref.document.dataset.tenant_id,
DocumentSegment.id == ref.segment_id,
DocumentSegment.dataset_id == ref.document.dataset.dataset_id,
DocumentSegment.tenant_id == ref.document.dataset.tenant_id,
)
)
if segment is None:
raise LookupError("Summary segment no longer exists in this dataset")
return session.scalar(
select(DocumentSegmentSummary)
.where(
DocumentSegmentSummary.dataset_id == ref.document.dataset.dataset_id,
DocumentSegmentSummary.document_id == ref.document.document_id,
DocumentSegmentSummary.chunk_id == ref.segment_id,
)
.limit(1)
)
def completed(self, command: SummaryVectorInput, result: SummaryVectorResult) -> None:
with self._new_session() as session, session.begin():
row = self._record(session, command)
if row is None:
row = DocumentSegmentSummary(
dataset_id=command.ref.document.dataset.dataset_id,
document_id=command.ref.document.document_id,
chunk_id=command.ref.segment_id,
summary_content=command.content,
enabled=True,
)
row.id = command.summary_id
session.add(row)
row.summary_content = command.content
row.summary_index_node_id = result.node_id
row.summary_index_node_hash = result.content_hash
row.tokens = result.tokens
row.status = SummaryStatus.COMPLETED
row.error = None
row.updated_at = naive_utc_now()
def failed(self, command: SummaryVectorInput, error: str) -> None:
with self._new_session() as session, session.begin():
row = self._record(session, command)
if row is not None:
row.status = SummaryStatus.ERROR
row.error = error
row.updated_at = naive_utc_now()
api/repositories/knowledge/upload_file_repository.py
+55
View File
@@ -0,0 +1,55 @@
"""SQLAlchemy repository for upload files used by knowledge extraction."""
from collections.abc import Sequence
from sqlalchemy import delete, select
from sqlalchemy.orm import Session, sessionmaker
from core.rag.extractor.entity.extract_setting import UploadFileExtractionInput
from models.model import UploadFile
def _query_files(session: Session, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFile]:
if not file_ids:
return {}
uploads = session.scalars(
select(UploadFile).where(UploadFile.id.in_(file_ids), UploadFile.tenant_id == workspace_id)
)
return {upload.id: upload for upload in uploads}
def query_upload_extraction_inputs(
session: Session, *, workspace_id: str, file_ids: Sequence[str]
) -> dict[str, UploadFileExtractionInput]:
"""Read detached extraction inputs without changing the caller's session lifecycle."""
return {
upload.id: UploadFileExtractionInput.model_validate(upload)
for upload in _query_files(session, workspace_id=workspace_id, file_ids=file_ids).values()
}
class SQLAlchemyKnowledgeUploadRepository:
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
def get_by_id(self, *, workspace_id: str, file_id: str) -> UploadFileExtractionInput | None:
return self.get_by_ids(workspace_id=workspace_id, file_ids=(file_id,)).get(file_id)
def get_by_ids(self, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFileExtractionInput]:
with self._session_factory() as session:
return query_upload_extraction_inputs(session, workspace_id=workspace_id, file_ids=file_ids)
def get_files(self, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFile]:
with self._session_factory() as session:
return _query_files(session, workspace_id=workspace_id, file_ids=file_ids)
def get_file_name(self, *, workspace_id: str, upload_file_id: str) -> str | None:
upload = self.get_files(workspace_id=workspace_id, file_ids=(upload_file_id,)).get(upload_file_id)
return upload.name if upload is not None else None
def delete_files(self, *, workspace_id: str, file_ids: Sequence[str]) -> None:
"""Remove preview records after their storage cleanup has finished."""
if not file_ids:
return
with self._session_factory.begin() as session:
session.execute(delete(UploadFile).where(UploadFile.tenant_id == workspace_id, UploadFile.id.in_(file_ids)))
api/repositories/workspace_member_query_repository.py
+14
View File
@@ -60,3 +60,17 @@ class WorkspaceMemberQueryRepository(WorkspaceMemberQuery):
)
return records
def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None:
"""Return one member's role without loading the whole workspace roster."""
with self._session_factory() as session:
role = session.scalar(
select(TenantAccountJoin.role)
.where(
TenantAccountJoin.tenant_id == workspace_id,
TenantAccountJoin.account_id == account_id,
)
.limit(1)
)
return role.value if role is not None else None
api/schedule/clean_unused_datasets_task.py
+2 -1
View File
@@ -14,6 +14,7 @@ from enums import CloudPlan
from extensions.ext_redis import redis_client
from libs.pagination import paginate_query
from models.dataset import Dataset, DatasetAutoDisableLog, DatasetQuery, Document
from repositories.knowledge.dataset_read_repository import get_dataset_doc_form
from services.feature_service import FeatureService
@@ -140,7 +141,7 @@ def clean_unused_datasets_task():
# Remove index
index_processor = IndexProcessorFactory(
dataset.get_doc_form(session=session)
get_dataset_doc_form(dataset, session=session)
).init_index_processor()
index_processor.clean(dataset, None, session=session)
api/services/agent/knowledge_datasets.py
+1 -1
View File
@@ -35,7 +35,7 @@ def get_tenant_knowledge_dataset_rows(*, session: Session, tenant_id: str, datas
a dataset row, so they are treated as missing instead of breaking the
UUID-typed dataset lookup.
"""
from services.dataset_service import DatasetService
from services.knowledge.dataset_service import DatasetService
valid_ids: list[str] = []
for dataset_id in dataset_ids:
api/services/app_service.py
+1 -1
View File
@@ -60,7 +60,7 @@ from services.entities.app_entities import (
AppListSortBy,
CreateAppParams,
)
from services.model_provider_service import ModelProviderService
from services.model_provider.service import ModelProviderService
from services.openapi.visibility import apply_openapi_gate, is_openapi_visible
from services.rbac_agent_access_service import initialize_agent_rbac_access
from services.system_feature_service import SystemFeatureService
Loaded 100 of 417 files, more files were not shown because too many files have changed in this diff. Show more