opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-21 18:58:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: explicit DB session propagation across backend paths (#38559)

Browse Source 
Co-authored-by: WH-2099 <wh2099@pm.me>
This commit is contained in:
Byron.wang 2026-07-15 14:48:28 +08:00 committed by GitHub
parent 7d5835fbc0
commit ab3e4daa95
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
397 changed files with 14206 additions and 9968 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
api/commands/plugin.py
View File

@ -9,6 +9,7 @@ from sqlalchemy import delete, func, select
from sqlalchemy.engine import CursorResult
from configs import dify_config
from core.db.session_factory import session_factory
from core.helper import encrypter
from core.plugin.entities.plugin_daemon import CredentialType
from core.plugin.impl.plugin import PluginInstaller
@ -578,9 +579,6 @@ def install_rag_pipeline_plugins(input_file, output_file, workers):
"""
click.echo(click.style("Installing rag pipeline plugins", fg="yellow"))
plugin_migration = PluginMigration()
plugin_migration.install_rag_pipeline_plugins(
input_file,
output_file,
workers,
)
with session_factory.create_session() as session:
plugin_migration.install_rag_pipeline_plugins(input_file, output_file, workers, session=session)
click.echo(click.style("Installing rag pipeline plugins successfully", fg="green"))

11
api/commands/system.py
View File

@ -188,23 +188,26 @@ where sites.id is null limit 1000"""
if app_id in failed_app_ids:
continue
session = db.session()
try:
app = db.session.scalar(select(App).where(App.id == app_id))
app = session.scalar(select(App).where(App.id == app_id))
if not app:
logger.info("App %s not found", app_id)
continue
tenant = app.tenant
tenant = session.get(Tenant, app.tenant_id)
if tenant:
accounts = tenant.get_accounts()
accounts = tenant.get_accounts(session=session)
if not accounts:
logger.info("Fix failed for app %s", app.id)
continue
account = accounts[0]
logger.info("Fixing missing site for app %s", app.id)
app_was_created.send(app, account=account)
app_was_created.send(app, account=account, session=session)
session.commit()
except Exception:
session.rollback()
failed_app_ids.append(app_id)
click.echo(click.style(f"Failed to fix missing site for app {app_id}", fg="red"))
logger.exception("Failed to fix app related site missing issue, app_id: %s", app_id)

18
api/commands/vector.py
View File

@ -1,10 +1,11 @@
import json
from typing import cast
import click
from flask import current_app
from sqlalchemy import select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config
from core.rag.datasource.vdb.vector_factory import Vector
@ -101,7 +102,8 @@ def migrate_annotation_vector_database():
)
documents.append(document)
vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"])
with Session(db.engine) as session:
vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"], session=session)
click.echo(f"Migrating annotations for app: {app.id}.")
try:
@ -176,6 +178,7 @@ def migrate_knowledge_vector_database():
VectorType.OCEANBASE,
}
page = 1
db_session = db.session()
while True:
try:
stmt = (
@ -184,7 +187,7 @@ def migrate_knowledge_vector_database():
.order_by(Dataset.created_at.desc())
)
datasets = paginate_query(stmt, page=page, per_page=50, max_per_page=50)
datasets = paginate_query(stmt, page=page, per_page=50, max_per_page=50, session=db_session)
if not datasets.items:
break
except SQLAlchemyError:
@ -227,7 +230,8 @@ def migrate_knowledge_vector_database():
index_struct_dict = {"type": vector_type, "vector_store": {"class_prefix": collection_name}}
dataset.index_struct = json.dumps(index_struct_dict)
vector = Vector(dataset)
with Session(db.engine) as session:
vector = Vector(dataset, session=session)
click.echo(f"Migrating dataset {dataset.id}.")
try:
@ -274,7 +278,7 @@ def migrate_knowledge_vector_database():
},
)
if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
child_chunks = segment.get_child_chunks()
child_chunks = segment.get_child_chunks(session=db_session)
if child_chunks:
child_documents = []
for child_chunk in child_chunks:
@ -410,7 +414,9 @@ def old_metadata_migration():
.where(DatasetDocument.doc_metadata.is_not(None))
.order_by(DatasetDocument.created_at.desc())
)
documents = paginate_query(stmt, page=page, per_page=50, max_per_page=50)
documents = paginate_query(
stmt, page=page, per_page=50, max_per_page=50, session=cast(Session, db.session())
)
except SQLAlchemyError:
raise
if not documents:

19
api/controllers/common/agent_app_parameters.py
View File

@ -1,27 +1,29 @@
from typing import Any
from sqlalchemy import select
from sqlalchemy.orm import Session
from core.app.apps.agent_app.app_feature_projection import merge_agent_app_features
from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form
from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError
from extensions.ext_database import db
from models.agent import Agent, AgentConfigSnapshot, AgentStatus
from models.agent_config_entities import AgentSoulConfig
from models.model import App
from models.model import App, load_annotation_reply_config
def get_published_agent_app_feature_dict_and_user_input_form(
app_model: App,
*,
session: Session,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
"""Return public Agent App parameters backed by the published Agent Soul."""
app_model_config = app_model.app_model_config
app_model_config = app_model.app_model_config_with_session(session=session)
agent_id = app_model.bound_agent_id
if not agent_id:
raise AgentAppGeneratorError("Agent App has no bound Agent")
agent = db.session.scalar(
agent = session.scalar(
select(Agent)
.where(
Agent.tenant_id == app_model.tenant_id,
@ -37,7 +39,7 @@ def get_published_agent_app_feature_dict_and_user_input_form(
if not agent.active_config_snapshot_id:
raise AgentAppNotPublishedError("Agent has not been published")
snapshot = db.session.scalar(
snapshot = session.scalar(
select(AgentConfigSnapshot)
.where(
AgentConfigSnapshot.tenant_id == app_model.tenant_id,
@ -50,5 +52,10 @@ def get_published_agent_app_feature_dict_and_user_input_form(
raise AgentAppGeneratorError("Agent published version not found")
agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict)
features_dict = merge_agent_app_features(agent_soul=agent_soul, app_model_config=app_model_config)
annotation_reply = load_annotation_reply_config(session, app_model.id) if app_model_config else None
features_dict = merge_agent_app_features(
agent_soul=agent_soul,
app_model_config=app_model_config,
annotation_reply=annotation_reply,
)
return features_dict, agent_app_variables_to_user_input_form(agent_soul.app_variables)

6
api/controllers/common/app_access.py
View File

@ -4,7 +4,8 @@ from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING
from extensions.ext_database import db
from sqlalchemy.orm import Session
from services.enterprise import rbac_service as enterprise_rbac_service
if TYPE_CHECKING:
@ -68,6 +69,7 @@ def resolve_app_access_filter(
tenant_id: str,
account_id: str,
*,
session: Session,
permissions: MyPermissionsResponse | None = None,
) -> AppAccessFilter:
"""Compute the RBAC app-access filter for ``account_id`` in ``tenant_id``.
@ -77,7 +79,7 @@ def resolve_app_access_filter(
inner-API round trip; otherwise it is fetched here.
"""
if permissions is None:
permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id, session=db.session())
permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id, session=session)
whitelist_scope = enterprise_rbac_service.RBACService.AppAccess.whitelist_resources(tenant_id, account_id)
can_manage_own_apps = _MANAGE_OWN_APPS_PERMISSION_KEY in permissions.workspace.permission_keys

18
api/controllers/common/session.py
View File

@ -2,8 +2,10 @@
`with_session` is an HTTP controller helper: it opens one SQLAlchemy session
for a Resource handler and injects it as the first argument after `self`.
Handlers use a transaction by default so migrated write paths keep
commit/rollback handling; pure read handlers may opt out with `write=False`.
Write handlers commit on success and roll back on failure. They use a regular
Session context so existing services may commit an intermediate unit and keep
using the same Session through SQLAlchemy's autobegin behavior. Pure read
handlers may opt out with `write=False`.
"""
from collections.abc import Callable
@ -38,14 +40,20 @@ def with_session[T, **P, R](
) -> (
Callable[Concatenate[T, P], R] | Callable[[Callable[Concatenate[T, Session, P], R]], Callable[Concatenate[T, P], R]]
):
"""Inject a request-scoped session, using a transaction only for write handlers."""
"""Inject a request-scoped session and finalize write handlers."""
def decorator(view: Callable[Concatenate[T, Session, P], R]) -> Callable[Concatenate[T, P], R]:
@wraps(view)
def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R:
if write:
with session_factory.get_session_maker().begin() as session:
return view(self, session, *args, **kwargs)
with session_factory.create_session() as session:
try:
result = view(self, session, *args, **kwargs)
session.commit()
return result
except Exception:
session.rollback() # noqa: no-new-controller-sqlalchemy decorator owns transaction rollback
raise
with session_factory.create_session() as session:
return view(self, session, *args, **kwargs)

11
api/controllers/console/agent/app_helpers.py
View File

@ -1,20 +1,21 @@
from uuid import UUID
from extensions.ext_database import db
from sqlalchemy.orm import Session
from models.model import App
from services.agent.roster_service import AgentRosterService
def resolve_agent_app_model(*, tenant_id: str, agent_id: UUID) -> App:
def resolve_agent_app_model(*, session: Session, tenant_id: str, agent_id: UUID) -> App:
"""Resolve a roster Agent's public Agent App."""
return AgentRosterService(db.session).get_agent_app_model(tenant_id=tenant_id, agent_id=str(agent_id))
return AgentRosterService(session).get_agent_app_model(tenant_id=tenant_id, agent_id=str(agent_id))
def resolve_agent_runtime_app_model(*, tenant_id: str, agent_id: UUID) -> App:
def resolve_agent_runtime_app_model(*, session: Session, tenant_id: str, agent_id: UUID) -> App:
"""Resolve the App that backs an Agent runtime surface.
This accepts both roster Agent Apps and workflow-only inline Agents with a
hidden backing App.
"""
return AgentRosterService(db.session).get_agent_runtime_app_model(tenant_id=tenant_id, agent_id=str(agent_id))
return AgentRosterService(session).get_agent_runtime_app_model(tenant_id=tenant_id, agent_id=str(agent_id))

143
api/controllers/console/agent/composer.py
View File

@ -2,9 +2,11 @@ from uuid import UUID
from flask import request
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import (
@ -17,7 +19,6 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user_id,
)
from extensions.ext_database import db
from fields.agent_fields import (
AgentAppComposerResponse,
AgentComposerCandidatesResponse,
@ -59,20 +60,21 @@ class WorkflowAgentComposerApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_user_id
@with_current_tenant_id
def get(self, tenant_id: str, account_id: str, app_model: App, node_id: str):
@with_session
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def get(self, session: Session, tenant_id: str, account_id: str, app_model: App, node_id: str):
query = WorkflowAgentComposerQuery.model_validate(request.args.to_dict(flat=True))
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.load_workflow_composer(
session=session,
tenant_id=tenant_id,
app_id=app_model.id,
node_id=node_id,
account_id=account_id,
snapshot_id=query.snapshot_id,
session=db.session(),
),
)
@ -85,20 +87,21 @@ class WorkflowAgentComposerApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_user_id
@with_current_tenant_id
def put(self, tenant_id: str, account_id: str, app_model: App, node_id: str):
@with_session
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def put(self, session: Session, tenant_id: str, account_id: str, app_model: App, node_id: str):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.save_workflow_composer(
session=session,
tenant_id=tenant_id,
app_id=app_model.id,
node_id=node_id,
account_id=account_id,
payload=payload,
session=db.session(),
),
)
@ -116,14 +119,16 @@ class WorkflowAgentComposerCopyFromRosterApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_user_id
@with_current_tenant_id
def post(self, tenant_id: str, account_id: str, app_model: App, node_id: str):
@with_session
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def post(self, session: Session, tenant_id: str, account_id: str, app_model: App, node_id: str):
payload = WorkflowComposerCopyFromRosterPayload.model_validate(console_ns.payload or {})
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.copy_workflow_composer_from_roster(
session=session,
tenant_id=tenant_id,
app_id=app_model.id,
node_id=node_id,
@ -131,7 +136,6 @@ class WorkflowAgentComposerCopyFromRosterApi(Resource):
source_agent_id=payload.source_agent_id,
source_snapshot_id=payload.source_snapshot_id,
idempotency_key=payload.idempotency_key,
session=db.session(),
),
)
@ -145,19 +149,22 @@ class WorkflowAgentComposerValidateApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_tenant_id
def post(self, tenant_id: str, app_model: App, node_id: str):
@with_session(write=False)
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def post(self, session: Session, tenant_id: str, app_model: App, node_id: str):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
ComposerConfigValidator.validate_publish_payload(payload)
AgentComposerService.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul)
AgentComposerService.validate_knowledge_datasets(
session=session, tenant_id=tenant_id, agent_soul=payload.agent_soul
)
findings = AgentComposerService.collect_validation_findings(
session=session,
tenant_id=tenant_id,
payload=payload,
agent_id=AgentComposerService.resolve_workflow_node_agent_id(
tenant_id=tenant_id, app_id=app_model.id, node_id=node_id, session=db.session()
session=session, tenant_id=tenant_id, app_id=app_model.id, node_id=node_id
),
session=db.session(),
)
return dump_response(AgentComposerValidateResponse, {"result": "success", "errors": [], **findings})
@ -170,18 +177,19 @@ class WorkflowAgentComposerCandidatesApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_user_id
@with_current_tenant_id
def get(self, tenant_id: str, current_user_id: str, app_model: App, node_id: str):
@with_session(write=False)
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def get(self, session: Session, tenant_id: str, current_user_id: str, app_model: App, node_id: str):
return dump_response(
AgentComposerCandidatesResponse,
AgentComposerService.get_workflow_candidates(
session=session,
tenant_id=tenant_id,
app_id=app_model.id,
node_id=node_id,
user_id=current_user_id,
session=db.session(),
),
)
@ -193,9 +201,10 @@ class WorkflowAgentComposerImpactApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_tenant_id
def post(self, tenant_id: str, app_model: App, node_id: str):
@with_session(write=False)
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def post(self, session: Session, tenant_id: str, app_model: App, node_id: str):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
current_snapshot_id = payload.binding.current_snapshot_id if payload.binding else None
if not current_snapshot_id:
@ -205,7 +214,7 @@ class WorkflowAgentComposerImpactApi(Resource):
return dump_response(
AgentComposerImpactResponse,
AgentComposerService.calculate_impact(
tenant_id=tenant_id, current_snapshot_id=current_snapshot_id, session=db.session()
session=session, tenant_id=tenant_id, current_snapshot_id=current_snapshot_id
),
)
@ -221,26 +230,27 @@ class WorkflowAgentComposerSaveToRosterApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
@with_current_user_id
@with_current_tenant_id
def post(self, tenant_id: str, account_id: str, app_model: App, node_id: str):
@with_session
@get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
def post(self, session: Session, tenant_id: str, account_id: str, app_model: App, node_id: str):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.save_workflow_composer(
session=session,
tenant_id=tenant_id,
app_id=app_model.id,
node_id=node_id,
account_id=account_id,
payload=payload,
session=db.session(),
),
)
def _require_snippet_app_id(*, tenant_id: str, snippet_id: UUID) -> str:
snippet = SnippetService(session=db.session()).get_snippet_by_id(
def _require_snippet_app_id(*, session: Session, tenant_id: str, snippet_id: UUID) -> str:
snippet = SnippetService(session=session).get_snippet_by_id(
snippet_id=str(snippet_id),
tenant_id=tenant_id,
)
@ -258,17 +268,18 @@ class SnippetAgentComposerApi(Resource):
@account_initialization_required
@with_current_user_id
@with_current_tenant_id
def get(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
@with_session
def get(self, session: Session, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
query = WorkflowAgentComposerQuery.model_validate(request.args.to_dict(flat=True))
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.load_workflow_composer(
session=session,
tenant_id=tenant_id,
app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
app_id=_require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id),
node_id=node_id,
account_id=account_id,
snapshot_id=query.snapshot_id,
session=db.session(),
),
)
@ -283,17 +294,18 @@ class SnippetAgentComposerApi(Resource):
)
@with_current_user_id
@with_current_tenant_id
def put(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
@with_session
def put(self, session: Session, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.save_workflow_composer(
session=session,
tenant_id=tenant_id,
app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
app_id=_require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id),
node_id=node_id,
account_id=account_id,
payload=payload,
session=db.session(),
),
)
@ -313,19 +325,20 @@ class SnippetAgentComposerCopyFromRosterApi(Resource):
)
@with_current_user_id
@with_current_tenant_id
def post(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
@with_session
def post(self, session: Session, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
payload = WorkflowComposerCopyFromRosterPayload.model_validate(console_ns.payload or {})
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.copy_workflow_composer_from_roster(
session=session,
tenant_id=tenant_id,
app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
app_id=_require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id),
node_id=node_id,
account_id=account_id,
source_agent_id=payload.source_agent_id,
source_snapshot_id=payload.source_snapshot_id,
idempotency_key=payload.idempotency_key,
session=db.session(),
),
)
@ -340,21 +353,24 @@ class SnippetAgentComposerValidateApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def post(self, tenant_id: str, snippet_id: UUID, node_id: str):
app_id = _require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id)
@with_session(write=False)
def post(self, session: Session, tenant_id: str, snippet_id: UUID, node_id: str):
app_id = _require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id)
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
ComposerConfigValidator.validate_publish_payload(payload)
AgentComposerService.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul)
AgentComposerService.validate_knowledge_datasets(
session=session, tenant_id=tenant_id, agent_soul=payload.agent_soul
)
findings = AgentComposerService.collect_validation_findings(
session=session,
tenant_id=tenant_id,
payload=payload,
agent_id=AgentComposerService.resolve_workflow_node_agent_id(
session=session,
tenant_id=tenant_id,
app_id=app_id,
node_id=node_id,
session=db.session(),
),
session=db.session(),
)
return dump_response(AgentComposerValidateResponse, {"result": "success", "errors": [], **findings})
@ -369,15 +385,16 @@ class SnippetAgentComposerCandidatesApi(Resource):
@account_initialization_required
@with_current_user_id
@with_current_tenant_id
def get(self, tenant_id: str, current_user_id: str, snippet_id: UUID, node_id: str):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user_id: str, snippet_id: UUID, node_id: str):
return dump_response(
AgentComposerCandidatesResponse,
AgentComposerService.get_workflow_candidates(
session=session,
tenant_id=tenant_id,
app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
app_id=_require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id),
node_id=node_id,
user_id=current_user_id,
session=db.session(),
),
)
@ -390,8 +407,9 @@ class SnippetAgentComposerImpactApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def post(self, tenant_id: str, snippet_id: UUID, node_id: str):
_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id)
@with_session(write=False)
def post(self, session: Session, tenant_id: str, snippet_id: UUID, node_id: str):
_require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id)
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
current_snapshot_id = payload.binding.current_snapshot_id if payload.binding else None
if not current_snapshot_id:
@ -401,9 +419,9 @@ class SnippetAgentComposerImpactApi(Resource):
return dump_response(
AgentComposerImpactResponse,
AgentComposerService.calculate_impact(
session=session,
tenant_id=tenant_id,
current_snapshot_id=current_snapshot_id,
session=db.session(),
),
)
@ -423,17 +441,18 @@ class SnippetAgentComposerSaveToRosterApi(Resource):
)
@with_current_user_id
@with_current_tenant_id
def post(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
@with_session
def post(self, session: Session, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
return dump_response(
WorkflowAgentComposerResponse,
AgentComposerService.save_workflow_composer(
session=session,
tenant_id=tenant_id,
app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
app_id=_require_snippet_app_id(session=session, tenant_id=tenant_id, snippet_id=snippet_id),
node_id=node_id,
account_id=account_id,
payload=payload,
session=db.session(),
),
)
@ -445,10 +464,11 @@ class AgentComposerApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, agent_id: UUID):
return dump_response(
AgentAppComposerResponse,
AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=str(agent_id), session=db.session()),
AgentComposerService.load_agent_composer(session=session, tenant_id=tenant_id, agent_id=str(agent_id)),
)
@console_ns.expect(console_ns.models[ComposerSavePayload.__name__])
@ -460,16 +480,17 @@ class AgentComposerApi(Resource):
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_current_user_id
@with_current_tenant_id
def put(self, tenant_id: str, account_id: str, agent_id: UUID):
@with_session
def put(self, session: Session, tenant_id: str, account_id: str, agent_id: UUID):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
return dump_response(
AgentAppComposerResponse,
AgentComposerService.save_agent_composer(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=account_id,
payload=payload,
session=db.session(),
),
)
@ -484,16 +505,19 @@ class AgentComposerValidateApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def post(self, tenant_id: str, agent_id: UUID):
AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=str(agent_id), session=db.session())
@with_session
def post(self, session: Session, tenant_id: str, agent_id: UUID):
AgentComposerService.load_agent_composer(session=session, tenant_id=tenant_id, agent_id=str(agent_id))
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
ComposerConfigValidator.validate_publish_payload(payload)
AgentComposerService.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul)
AgentComposerService.validate_knowledge_datasets(
session=session, tenant_id=tenant_id, agent_soul=payload.agent_soul
)
findings = AgentComposerService.collect_validation_findings(
session=session,
tenant_id=tenant_id,
payload=payload,
agent_id=str(agent_id),
session=db.session(),
)
return dump_response(AgentComposerValidateResponse, {"result": "success", "errors": [], **findings})
@ -508,13 +532,14 @@ class AgentComposerCandidatesApi(Resource):
@account_initialization_required
@with_current_user_id
@with_current_tenant_id
def get(self, tenant_id: str, current_user_id: str, agent_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user_id: str, agent_id: UUID):
return dump_response(
AgentComposerCandidatesResponse,
AgentComposerService.get_agent_app_candidates(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
user_id=current_user_id,
session=db.session(),
),
)

228
api/controllers/console/agent/roster.py
View File

@ -4,6 +4,7 @@ from flask import abort, request
from flask_restx import Resource
from pydantic import AliasChoices, BaseModel, Field, field_validator
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from controllers.common.schema import (
query_params_from_model,
@ -11,8 +12,8 @@ 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.agent.app_helpers import resolve_agent_app_model, resolve_agent_runtime_app_model
from controllers.console.apikey import ApiKeyItem, ApiKeyList, BaseApiKeyListResource, BaseApiKeyResource
from controllers.console.app.app import (
APP_LIST_QUERY_ARRAY_FIELDS,
@ -42,7 +43,6 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from extensions.ext_database import db
from fields.agent_fields import (
AgentConfigDraftSummaryResponse,
AgentConfigSnapshotDetailResponse,
@ -339,11 +339,13 @@ register_response_schema_models(
)
def _agent_roster_service() -> AgentRosterService:
return AgentRosterService(db.session)
def _agent_roster_service(session: Session) -> AgentRosterService:
return AgentRosterService(session)
def _serialize_agent_app_detail(app_model, *, current_user: Account, agent_id: str | None = None) -> dict:
def _serialize_agent_app_detail(
session: Session, app_model, *, current_user: Account, agent_id: str | None = None
) -> dict:
"""Serialize an Agent App detail using roster-only DTOs.
`/agent` responses are roster-shaped rather than raw app-shaped: `id`
@ -353,15 +355,19 @@ def _serialize_agent_app_detail(app_model, *, current_user: Account, agent_id: s
roster persona fields without widening the shared /apps detail schema.
"""
app_model = AppService().get_app(app_model)
app_model = AppService().get_app(app_model, session=session)
if FeatureService.get_system_features().webapp_auth.enabled:
app_setting = EnterpriseService.WebAppAuth.get_app_access_mode_by_id(app_id=str(app_model.id))
app_model.access_mode = app_setting.access_mode # type: ignore[attr-defined]
roster_service = _agent_roster_service()
payload = AgentAppDetailWithSite.model_validate(app_model, from_attributes=True).model_dump(mode="json")
roster_service = _agent_roster_service(session)
payload = AgentAppDetailWithSite.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
agent = (
db.session.scalar(
session.scalar(
select(Agent).where(
Agent.tenant_id == app_model.tenant_id,
Agent.id == agent_id,
@ -382,6 +388,7 @@ def _serialize_agent_app_detail(app_model, *, current_user: Account, agent_id: s
tenant_id=app_model.tenant_id,
agent_id=agent.id,
account_id=current_user.id,
commit=False,
)
message_count = roster_service.count_agent_app_debug_conversation_messages(
conversation_id=debug_conversation_id,
@ -397,7 +404,7 @@ def _serialize_agent_app_detail(app_model, *, current_user: Account, agent_id: s
return payload
def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str, current_user: Account) -> dict:
def _serialize_agent_app_pagination(session: Session, app_pagination, *, tenant_id: str, current_user: Account) -> dict:
"""Serialize Agent App lists with roster-shaped items.
Each item starts from the shared App list shape, then drops
@ -407,7 +414,7 @@ def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str, current_u
"""
app_ids = [str(app.id) for app in app_pagination.items]
roster_service = _agent_roster_service()
roster_service = _agent_roster_service(session)
agents_by_app_id = roster_service.load_app_backing_agents_by_app_id(
tenant_id=tenant_id,
app_ids=app_ids,
@ -425,7 +432,11 @@ def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str, current_u
agents=list(agents_by_app_id.values()),
account_id=current_user.id,
)
payload = AgentAppPagination.model_validate(app_pagination, from_attributes=True).model_dump(mode="json")
payload = AgentAppPagination.model_validate(
app_pagination,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
for item in payload["data"]:
app_id = item["id"]
item.pop("bound_agent_id", None)
@ -456,13 +467,17 @@ def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str, current_u
)
def _resolve_agent_app_model(*, tenant_id: str, agent_id: UUID):
return resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
def _resolve_agent_app_model(session: Session, *, tenant_id: str, agent_id: UUID) -> App:
return _agent_roster_service(session).get_agent_app_model(tenant_id=tenant_id, agent_id=str(agent_id))
def _agent_api_key_count(app_id: str) -> int:
def _resolve_agent_runtime_app_model(session: Session, *, tenant_id: str, agent_id: UUID) -> App:
return _agent_roster_service(session).get_agent_runtime_app_model(tenant_id=tenant_id, agent_id=str(agent_id))
def _agent_api_key_count(session: Session, app_id: str) -> int:
return (
db.session.scalar(
session.scalar(
select(func.count(ApiToken.id)).where(
ApiToken.type == ApiTokenType.APP,
ApiToken.app_id == app_id,
@ -472,7 +487,7 @@ def _agent_api_key_count(app_id: str) -> int:
)
def _serialize_agent_api_access(app_model: App) -> dict:
def _serialize_agent_api_access(session: Session, app_model: App) -> dict:
base_url = app_model.api_base_url
response = AgentApiAccessResponse(
enabled=bool(app_model.enable_api),
@ -487,13 +502,13 @@ def _serialize_agent_api_access(app_model: App) -> dict:
meta_endpoint=f"{base_url}/meta",
api_rpm=app_model.api_rpm or 0,
api_rph=app_model.api_rph or 0,
api_key_count=_agent_api_key_count(str(app_model.id)),
api_key_count=_agent_api_key_count(session, str(app_model.id)),
)
return response.model_dump(mode="json")
def _agent_observability_service() -> AgentObservabilityService:
return AgentObservabilityService(db.session)
def _agent_observability_service(session: Session) -> AgentObservabilityService:
return AgentObservabilityService(session)
def _parse_observability_time_range(start: str | None, end: str | None, account: Account):
@ -520,7 +535,8 @@ class AgentAppListApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account):
@with_session
def get(self, session: Session, current_tenant_id: str, current_user: Account):
args = query_params_from_request(AppListQuery, list_fields=APP_LIST_QUERY_ARRAY_FIELDS)
params = AppListParams(
page=args.page,
@ -534,12 +550,13 @@ class AgentAppListApi(Resource):
status="normal",
)
app_pagination = AppService().get_paginate_apps(current_user.id, current_tenant_id, params, db.session())
app_pagination = AppService().get_paginate_apps(current_user.id, current_tenant_id, params, session)
if app_pagination is None:
empty = AgentAppPagination(page=args.page, limit=args.limit, total=0, has_more=False, data=[])
return empty.model_dump(mode="json")
return _serialize_agent_app_pagination(
session,
app_pagination,
tenant_id=current_tenant_id,
current_user=current_user,
@ -555,7 +572,8 @@ class AgentAppListApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, current_tenant_id: str, current_user: Account):
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account):
args = AgentAppCreatePayload.model_validate(console_ns.payload)
params = CreateAppParams(
name=args.name,
@ -567,8 +585,8 @@ class AgentAppListApi(Resource):
icon_background=args.icon_background,
)
app = AppService().create_app(current_tenant_id, params, current_user, session=db.session())
return _serialize_agent_app_detail(app, current_user=current_user), 201
app = AppService().create_app(current_tenant_id, params, current_user, session=session)
return _serialize_agent_app_detail(session, app, current_user=current_user), 201
@console_ns.route("/agent/<uuid:agent_id>")
@ -580,9 +598,10 @@ class AgentAppApi(Resource):
@enterprise_license_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _serialize_agent_app_detail(app_model, current_user=current_user, agent_id=str(agent_id))
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = _resolve_agent_runtime_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
return _serialize_agent_app_detail(session, app_model, current_user=current_user, agent_id=str(agent_id))
@console_ns.expect(console_ns.models[AgentAppUpdatePayload.__name__])
@console_ns.response(200, "Agent app updated successfully", console_ns.models[AgentAppDetailWithSite.__name__])
@ -594,8 +613,9 @@ class AgentAppApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def put(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session
def put(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
args = AgentAppUpdatePayload.model_validate(console_ns.payload)
args_dict: AppService.ArgsDict = {
"name": args.name,
@ -607,8 +627,8 @@ class AgentAppApi(Resource):
"max_active_requests": args.max_active_requests or 0,
"role": args.role,
}
updated = AppService().update_app(app_model, args_dict, session=db.session())
return _serialize_agent_app_detail(updated, current_user=current_user)
updated = AppService().update_app(app_model, args_dict, session=session)
return _serialize_agent_app_detail(session, updated, current_user=current_user)
@console_ns.response(204, "Agent app deleted successfully")
@console_ns.response(403, "Insufficient permissions")
@ -617,9 +637,10 @@ class AgentAppApi(Resource):
@account_initialization_required
@edit_permission_required
@with_current_tenant_id
def delete(self, tenant_id: str, agent_id: UUID):
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
AppService().delete_app(app_model, session=db.session())
@with_session
def delete(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
AppService().delete_app(app_model, session=session)
return "", 204
@ -637,8 +658,9 @@ class AgentDebugConversationRefreshApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
debug_conversation_id = _agent_roster_service().refresh_agent_app_debug_conversation_id(
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
debug_conversation_id = _agent_roster_service(session).refresh_agent_app_debug_conversation_id(
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
@ -661,14 +683,15 @@ class AgentPublishApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
args = AgentPublishPayload.model_validate(console_ns.payload or {})
return AgentComposerService.publish_agent_app_draft(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
version_note=args.version_note,
session=db.session(),
)
@ -682,14 +705,15 @@ class AgentBuildDraftCheckoutApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
args = AgentBuildDraftCheckoutPayload.model_validate(console_ns.payload or {})
return AgentComposerService.checkout_agent_app_build_draft(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
force=args.force,
session=db.session(),
)
@ -702,12 +726,13 @@ class AgentBuildDraftApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return AgentComposerService.load_agent_app_build_draft(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
session=db.session(),
)
@console_ns.expect(console_ns.models[ComposerSavePayload.__name__])
@ -718,14 +743,15 @@ class AgentBuildDraftApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def put(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def put(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
return AgentComposerService.save_agent_app_build_draft(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
payload=payload,
session=db.session(),
)
@console_ns.response(200, "Agent build draft discarded", console_ns.models[AgentSimpleResultResponse.__name__])
@ -735,12 +761,13 @@ class AgentBuildDraftApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def delete(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return AgentComposerService.discard_agent_app_build_draft(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
session=db.session(),
)
@ -753,12 +780,13 @@ class AgentBuildDraftApplyApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return AgentComposerService.apply_agent_app_build_draft(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
session=db.session(),
)
@ -774,9 +802,10 @@ class AgentAppCopyApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
args = AgentAppCopyPayload.model_validate(console_ns.payload or {})
copied_app = _agent_roster_service().duplicate_agent_app(
copied_app = _agent_roster_service(session).duplicate_agent_app(
tenant_id=tenant_id,
agent_id=str(agent_id),
account=current_user,
@ -787,7 +816,7 @@ class AgentAppCopyApi(Resource):
icon=args.icon,
icon_background=args.icon_background,
)
return _serialize_agent_app_detail(copied_app, current_user=current_user), 201
return _serialize_agent_app_detail(session, copied_app, current_user=current_user), 201
@console_ns.route("/agent/<uuid:agent_id>/api-access")
@ -797,9 +826,10 @@ class AgentApiAccessApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _serialize_agent_api_access(app_model)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
return _serialize_agent_api_access(session, app_model)
@console_ns.route("/agent/<uuid:agent_id>/api-enable")
@ -813,11 +843,12 @@ class AgentApiStatusApi(Resource):
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
@with_current_tenant_id
def post(self, tenant_id: str, agent_id: UUID):
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session
def post(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
args = AgentApiStatusPayload.model_validate(console_ns.payload)
app_model = AppService().update_app_api_status(app_model, args.enable_api, session=db.session())
return _serialize_agent_api_access(app_model)
app_model = AppService().update_app_api_status(app_model, args.enable_api, session=session)
return _serialize_agent_api_access(session, app_model)
@console_ns.route("/agent/<uuid:agent_id>/api-keys")
@ -829,18 +860,23 @@ class AgentApiKeyListApi(BaseApiKeyListResource):
@console_ns.response(200, "Agent service API keys", console_ns.models[ApiKeyList.__name__])
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID) -> dict[str, object]:
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
return dump_response(ApiKeyList, self._get_api_key_list(str(app_model.id), tenant_id))
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID) -> dict[str, object]:
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
return dump_response(ApiKeyList, self._get_api_key_list(str(app_model.id), tenant_id, session=session))
@console_ns.response(201, "Agent service API key created", console_ns.models[ApiKeyItem.__name__])
@console_ns.response(400, "Maximum keys exceeded")
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
def post(self, tenant_id: str, agent_id: UUID) -> tuple[dict[str, object], int]:
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
return dump_response(ApiKeyItem, self._create_api_key(str(app_model.id), tenant_id)), 201
@with_session
def post(self, session: Session, tenant_id: str, agent_id: UUID) -> tuple[dict[str, object], int]:
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
return dump_response(
ApiKeyItem,
self._create_api_key(str(app_model.id), tenant_id, session=session),
), 201
@console_ns.route("/agent/<uuid:agent_id>/api-keys/<uuid:api_key_id>")
@ -853,9 +889,17 @@ class AgentApiKeyApi(BaseApiKeyResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID, api_key_id: UUID) -> tuple[str, int]:
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
self._delete_api_key(str(app_model.id), str(api_key_id), tenant_id, current_user)
@with_session
def delete(
self,
session: Session,
tenant_id: str,
current_user: Account,
agent_id: UUID,
api_key_id: UUID,
) -> tuple[str, int]:
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
self._delete_api_key(str(app_model.id), str(api_key_id), tenant_id, current_user, session=session)
return "", 204
@ -867,11 +911,12 @@ class AgentInviteOptionsApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str):
@with_session(write=False)
def get(self, session: Session, tenant_id: str):
query = AgentInviteOptionsQuery.model_validate(request.args.to_dict(flat=True))
return dump_response(
AgentInviteOptionsResponse,
_agent_roster_service().list_invite_options(
_agent_roster_service(session).list_invite_options(
tenant_id=tenant_id,
page=query.page,
limit=query.limit,
@ -890,15 +935,16 @@ class AgentLogsApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = _resolve_agent_runtime_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
query_data: dict[str, object] = dict(request.args.to_dict(flat=True))
query_data["sources"] = _query_values("sources", "source")
query_data["statuses"] = _query_values("statuses", "status")
query = AgentLogsQuery.model_validate(query_data)
start, end = _parse_observability_time_range(query.start, query.end, current_user)
try:
payload = _agent_observability_service().list_logs(
payload = _agent_observability_service(session).list_logs(
app=app_model,
agent_id=str(agent_id),
params=AgentLogQueryParams(
@ -927,15 +973,16 @@ class AgentLogMessagesApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, conversation_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, conversation_id: UUID):
app_model = _resolve_agent_runtime_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
query_data: dict[str, object] = dict(request.args.to_dict(flat=True))
query_data["sources"] = _query_values("sources", "source")
query_data["statuses"] = _query_values("statuses", "status")
query = AgentLogsQuery.model_validate(query_data)
start, end = _parse_observability_time_range(query.start, query.end, current_user)
try:
payload = _agent_observability_service().list_log_messages(
payload = _agent_observability_service(session).list_log_messages(
app=app_model,
agent_id=str(agent_id),
conversation_id=str(conversation_id),
@ -964,9 +1011,10 @@ class AgentLogSourcesApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
payload = _agent_observability_service().list_log_sources(app=app_model, agent_id=str(agent_id))
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = _resolve_agent_runtime_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
payload = _agent_observability_service(session).list_log_sources(app=app_model, agent_id=str(agent_id))
return dump_response(AgentLogSourceListResponse, payload)
@ -983,13 +1031,14 @@ class AgentStatisticsSummaryApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = _resolve_agent_runtime_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
query = AgentStatisticsQuery.model_validate(request.args.to_dict(flat=True))
timezone = current_user.timezone or "UTC"
start, end = _parse_observability_time_range(query.start, query.end, current_user)
try:
payload = _agent_observability_service().get_statistics_summary(
payload = _agent_observability_service(session).get_statistics_summary(
app=app_model,
agent_id=str(agent_id),
params=AgentStatisticsQueryParams(source=query.source, start=start, end=end, timezone=timezone),
@ -1006,10 +1055,11 @@ class AgentRosterVersionsApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
return dump_response(
AgentConfigSnapshotListResponse,
{"data": _agent_roster_service().list_agent_versions(tenant_id=tenant_id, agent_id=str(agent_id))},
{"data": _agent_roster_service(session).list_agent_versions(tenant_id=tenant_id, agent_id=str(agent_id))},
)
@ -1020,10 +1070,11 @@ class AgentRosterVersionDetailApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID, version_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID, version_id: UUID):
return dump_response(
AgentConfigSnapshotDetailResponse,
_agent_roster_service().get_agent_version_detail(
_agent_roster_service(session).get_agent_version_detail(
tenant_id=tenant_id,
agent_id=str(agent_id),
version_id=str(version_id),
@ -1040,10 +1091,11 @@ class AgentRosterVersionRestoreApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID, version_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, version_id: UUID):
return dump_response(
AgentConfigSnapshotRestoreResponse,
_agent_roster_service().restore_agent_version(
_agent_roster_service(session).restore_agent_version(
tenant_id=tenant_id,
agent_id=str(agent_id),
version_id=str(version_id),

131
api/controllers/console/apikey.py
View File

@ -6,12 +6,12 @@ from flask_restx import Resource
from flask_restx._http import HTTPStatus
from pydantic import field_validator
from sqlalchemy import delete, func, select
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
from configs import dify_config
from controllers.common.schema import register_response_schema_models
from extensions.ext_database import db
from controllers.common.session import with_session
from fields.base import ResponseModel
from libs.helper import dump_response, to_timestamp
from libs.login import login_required
@ -54,11 +54,10 @@ class ApiKeyList(ResponseModel):
register_response_schema_models(console_ns, ApiKeyItem, ApiKeyList)
def _get_resource(resource_id, tenant_id, resource_model):
with sessionmaker(db.engine).begin() as session:
resource = session.execute(
select(resource_model).filter_by(id=resource_id, tenant_id=tenant_id)
).scalar_one_or_none()
def _get_resource(resource_id, tenant_id, resource_model, *, session: Session):
resource = session.execute(
select(resource_model).filter_by(id=resource_id, tenant_id=tenant_id)
).scalar_one_or_none()
if resource is None:
flask_restx.abort(HTTPStatus.NOT_FOUND, message=f"{resource_model.__name__} not found.")
@ -75,14 +74,18 @@ class BaseApiKeyListResource(Resource):
token_prefix: str | None = None
max_keys = 10
def get(self, resource_id: str, current_tenant_id: str) -> dict[str, object]:
return dump_response(ApiKeyList, self._get_api_key_list(resource_id, current_tenant_id))
@with_session(write=False)
def get(self, session: Session, resource_id: str, current_tenant_id: str) -> dict[str, object]:
return dump_response(
ApiKeyList,
self._get_api_key_list(resource_id, current_tenant_id, session=session),
)
def _get_api_key_list(self, resource_id: str, current_tenant_id: str) -> ApiKeyList:
def _get_api_key_list(self, resource_id: str, current_tenant_id: str, *, session: Session) -> ApiKeyList:
assert self.resource_id_field is not None, "resource_id_field must be set"
_get_resource(resource_id, current_tenant_id, self.resource_model)
keys = db.session.scalars(
_get_resource(resource_id, current_tenant_id, self.resource_model, session=session)
keys = session.scalars(
select(ApiToken).where(
ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id
)
@ -90,14 +93,18 @@ class BaseApiKeyListResource(Resource):
return ApiKeyList.model_validate({"data": keys}, from_attributes=True)
@edit_permission_required
def post(self, resource_id: str, current_tenant_id: str) -> tuple[dict[str, object], int]:
return dump_response(ApiKeyItem, self._create_api_key(resource_id, current_tenant_id)), 201
@with_session
def post(self, session: Session, resource_id: str, current_tenant_id: str) -> tuple[dict[str, object], int]:
return dump_response(
ApiKeyItem,
self._create_api_key(resource_id, current_tenant_id, session=session),
), 201
def _create_api_key(self, resource_id: str, current_tenant_id: str) -> ApiToken:
def _create_api_key(self, resource_id: str, current_tenant_id: str, *, session: Session) -> ApiToken:
assert self.resource_id_field is not None, "resource_id_field must be set"
_get_resource(resource_id, current_tenant_id, self.resource_model)
_get_resource(resource_id, current_tenant_id, self.resource_model, session=session)
current_key_count: int = (
db.session.scalar(
session.scalar(
select(func.count(ApiToken.id)).where(
ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id
)
@ -112,15 +119,15 @@ class BaseApiKeyListResource(Resource):
custom="max_keys_exceeded",
)
key = ApiToken.generate_api_key(self.token_prefix or "", 24)
key = ApiToken.generate_api_key(self.token_prefix or "", 24, session=session)
assert self.resource_type is not None, "resource_type must be set"
api_token = ApiToken()
setattr(api_token, self.resource_id_field, resource_id)
api_token.tenant_id = current_tenant_id
api_token.token = key
api_token.type = self.resource_type
db.session.add(api_token)
db.session.commit()
session.add(api_token)
session.commit()
return api_token
@ -131,10 +138,16 @@ class BaseApiKeyResource(Resource):
resource_model: type | None = None
resource_id_field: str | None = None
@with_session
def delete(
self, resource_id: str, api_key_id: str, current_tenant_id: str, current_user: Account
self,
session: Session,
resource_id: str,
api_key_id: str,
current_tenant_id: str,
current_user: Account,
) -> tuple[str, int]:
self._delete_api_key(resource_id, api_key_id, current_tenant_id, current_user)
self._delete_api_key(resource_id, api_key_id, current_tenant_id, current_user, session=session)
return "", 204
def _delete_api_key(
@ -143,14 +156,16 @@ class BaseApiKeyResource(Resource):
api_key_id: str,
current_tenant_id: str,
current_user: Account,
*,
session: Session,
) -> None:
assert self.resource_id_field is not None, "resource_id_field must be set"
_get_resource(resource_id, current_tenant_id, self.resource_model)
_get_resource(resource_id, current_tenant_id, self.resource_model, session=session)
if not dify_config.RBAC_ENABLED and not current_user.is_admin_or_owner:
raise Forbidden()
key = db.session.scalar(
key = session.scalar(
select(ApiToken)
.where(
getattr(ApiToken, self.resource_id_field) == resource_id,
@ -168,8 +183,8 @@ class BaseApiKeyResource(Resource):
assert key is not None # nosec - for type checker only
ApiTokenCache.delete(key.token, key.type)
db.session.execute(delete(ApiToken).where(ApiToken.id == api_key_id))
db.session.commit()
session.execute(delete(ApiToken).where(ApiToken.id == api_key_id))
session.commit()
@console_ns.route("/apps/<uuid:resource_id>/api-keys")
@ -179,9 +194,13 @@ class AppApiKeyListResource(BaseApiKeyListResource):
@console_ns.doc(params={"resource_id": "App ID"})
@console_ns.response(200, "API keys retrieved successfully", console_ns.models[ApiKeyList.__name__])
@with_current_tenant_id
def get(self, current_tenant_id: str, resource_id: UUID) -> dict[str, object]:
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, resource_id: UUID) -> dict[str, object]:
"""Get all API keys for an app"""
return dump_response(ApiKeyList, self._get_api_key_list(str(resource_id), current_tenant_id))
return dump_response(
ApiKeyList,
self._get_api_key_list(str(resource_id), current_tenant_id, session=session),
)
@console_ns.doc("create_app_api_key")
@console_ns.doc(description="Create a new API key for an app")
@ -191,9 +210,13 @@ class AppApiKeyListResource(BaseApiKeyListResource):
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
def post(self, current_tenant_id: str, resource_id: UUID) -> tuple[dict[str, object], int]:
@with_session
def post(self, session: Session, current_tenant_id: str, resource_id: UUID) -> tuple[dict[str, object], int]:
"""Create a new API key for an app"""
return dump_response(ApiKeyItem, self._create_api_key(str(resource_id), current_tenant_id)), 201
return dump_response(
ApiKeyItem,
self._create_api_key(str(resource_id), current_tenant_id, session=session),
), 201
resource_type = ApiTokenType.APP
resource_model = App
@ -210,11 +233,23 @@ class AppApiKeyResource(BaseApiKeyResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
@with_session
def delete(
self, current_tenant_id: str, current_user: Account, resource_id: UUID, api_key_id: UUID
self,
session: Session,
current_tenant_id: str,
current_user: Account,
resource_id: UUID,
api_key_id: UUID,
) -> tuple[str, int]:
"""Delete an API key for an app"""
self._delete_api_key(str(resource_id), str(api_key_id), current_tenant_id, current_user)
self._delete_api_key(
str(resource_id),
str(api_key_id),
current_tenant_id,
current_user,
session=session,
)
return "", 204
resource_type = ApiTokenType.APP
@ -229,9 +264,13 @@ class DatasetApiKeyListResource(BaseApiKeyListResource):
@console_ns.doc(params={"resource_id": "Dataset ID"})
@console_ns.response(200, "API keys retrieved successfully", console_ns.models[ApiKeyList.__name__])
@with_current_tenant_id
def get(self, current_tenant_id: str, resource_id: UUID) -> dict[str, object]:
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, resource_id: UUID) -> dict[str, object]:
"""Get all API keys for a dataset"""
return dump_response(ApiKeyList, self._get_api_key_list(str(resource_id), current_tenant_id))
return dump_response(
ApiKeyList,
self._get_api_key_list(str(resource_id), current_tenant_id, session=session),
)
@console_ns.doc("create_dataset_api_key")
@console_ns.doc(description="Create a new API key for a dataset")
@ -241,9 +280,13 @@ class DatasetApiKeyListResource(BaseApiKeyListResource):
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_API_KEY_MANAGE)
def post(self, current_tenant_id: str, resource_id: UUID) -> tuple[dict[str, object], int]:
@with_session
def post(self, session: Session, current_tenant_id: str, resource_id: UUID) -> tuple[dict[str, object], int]:
"""Create a new API key for a dataset"""
return dump_response(ApiKeyItem, self._create_api_key(str(resource_id), current_tenant_id)), 201
return dump_response(
ApiKeyItem,
self._create_api_key(str(resource_id), current_tenant_id, session=session),
), 201
resource_type = ApiTokenType.DATASET
resource_model = Dataset
@ -260,11 +303,23 @@ class DatasetApiKeyResource(BaseApiKeyResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_API_KEY_MANAGE)
@with_session
def delete(
self, current_tenant_id: str, current_user: Account, resource_id: UUID, api_key_id: UUID
self,
session: Session,
current_tenant_id: str,
current_user: Account,
resource_id: UUID,
api_key_id: UUID,
) -> tuple[str, int]:
"""Delete an API key for a dataset"""
self._delete_api_key(str(resource_id), str(api_key_id), current_tenant_id, current_user)
self._delete_api_key(
str(resource_id),
str(api_key_id),
current_tenant_id,
current_user,
session=session,
)
return "", 204
resource_type = ApiTokenType.DATASET

122
api/controllers/console/app/agent.py
View File

@ -5,6 +5,7 @@ from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import select
from sqlalchemy.orm import Session
from controllers.common.schema import (
query_params_from_model,
@ -12,6 +13,7 @@ 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.agent.app_helpers import resolve_agent_runtime_app_model
from controllers.console.app.wraps import get_app_model
@ -24,7 +26,6 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.helper import uuid_value
from libs.login import login_required
@ -169,23 +170,23 @@ register_response_schema_models(
)
def _resolve_agent_id(app_model: App, node_id: str | None) -> str | None:
def _resolve_agent_id(session: Session, app_model: App, node_id: str | None) -> str | None:
if node_id and app_model.mode != AppMode.AGENT:
return AgentComposerService.resolve_workflow_node_agent_id(
tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id, session=db.session()
session=session, tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id
)
return app_model.bound_agent_id
return app_model.bound_agent_id_with_session(session=session)
def _agent_not_bound() -> tuple[dict[str, str], int]:
return {"code": "agent_not_bound", "message": "no agent is bound for this app/node"}, 400
def _upload_skill_for_app(*, current_user: Account, app_model: App):
def _upload_skill_for_app(*, session: Session, current_user: Account, app_model: App):
"""Upload one skill package and commit its normalized files into the agent drive."""
query = query_params_from_request(AgentDriveMutationQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
if "file" not in request.files:
@ -202,22 +203,22 @@ def _upload_skill_for_app(*, current_user: Account, app_model: App):
tenant_id=app_model.tenant_id,
user_id=current_user.id,
agent_id=agent_id,
session=db.session(),
session=session,
)
except (SkillPackageError, AgentDriveError) as exc:
return {"code": exc.code, "message": exc.message}, exc.status_code
return result, 201
def _commit_drive_file_for_app(*, current_user: Account, app_model: App, allow_node_id: bool = True):
def _commit_drive_file_for_app(*, session: Session, current_user: Account, app_model: App, allow_node_id: bool = True):
query = query_params_from_request(AgentDriveMutationQuery)
node_id = query.node_id if allow_node_id else None
agent_id = _resolve_agent_id(app_model, node_id)
agent_id = _resolve_agent_id(session, app_model, node_id)
if not agent_id:
return _agent_not_bound()
payload = AgentDriveFilePayload.model_validate(console_ns.payload or {})
upload_file = db.session.scalar(
upload_file = session.scalar(
select(UploadFile).where(
UploadFile.id == payload.upload_file_id,
UploadFile.tenant_id == app_model.tenant_id,
@ -241,7 +242,7 @@ def _commit_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
value_owned_by_drive=True,
)
],
session=db.session(),
session=session,
)
except AgentDriveError as exc:
return {"code": exc.code, "message": exc.message}, exc.status_code
@ -258,10 +259,10 @@ def _commit_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
}, 201
def _delete_drive_file_for_app(*, current_user: Account, app_model: App, allow_node_id: bool = True):
def _delete_drive_file_for_app(*, session: Session, current_user: Account, app_model: App, allow_node_id: bool = True):
query = query_params_from_request(AgentDriveDeleteFileQuery)
node_id = query.node_id if allow_node_id else None
agent_id = _resolve_agent_id(app_model, node_id)
agent_id = _resolve_agent_id(session, app_model, node_id)
if not agent_id:
return _agent_not_bound()
try:
@ -275,7 +276,7 @@ def _delete_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
user_id=current_user.id,
agent_id=agent_id,
items=[DriveCommitItem(key=key, file_ref=None)],
session=db.session(),
session=session,
)
except AgentDriveError as exc:
return {"code": exc.code, "message": exc.message}, exc.status_code
@ -283,10 +284,12 @@ def _delete_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
return {"result": "success", "removed_keys": removed_keys}
def _delete_skill_for_app(*, current_user: Account, app_model: App, slug: str, allow_node_id: bool = True):
def _delete_skill_for_app(
*, session: Session, current_user: Account, app_model: App, slug: str, allow_node_id: bool = True
):
query = query_params_from_request(AgentDriveMutationQuery)
node_id = query.node_id if allow_node_id else None
agent_id = _resolve_agent_id(app_model, node_id)
agent_id = _resolve_agent_id(session, app_model, node_id)
if not agent_id:
return _agent_not_bound()
if "/" in slug or not slug.strip():
@ -301,7 +304,7 @@ def _delete_skill_for_app(*, current_user: Account, app_model: App, slug: str, a
DriveCommitItem(key=f"{slug}/SKILL.md", file_ref=None),
DriveCommitItem(key=f"{slug}/.DIFY-SKILL-FULL.zip", file_ref=None),
],
session=db.session(),
session=session,
)
except AgentDriveError as exc:
return {"code": exc.code, "message": exc.message}, exc.status_code
@ -309,16 +312,16 @@ def _delete_skill_for_app(*, current_user: Account, app_model: App, slug: str, a
return {"result": "success", "removed_keys": removed_keys}
def _infer_skill_tools_for_app(*, app_model: App, slug: str):
def _infer_skill_tools_for_app(*, session: Session, app_model: App, slug: str):
query = query_params_from_request(AgentDriveMutationQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
if "/" in slug or not slug.strip():
return {"code": "drive_key_invalid", "message": "skill slug must be a single path segment"}, 400
try:
return SkillToolInferenceService().infer(
tenant_id=app_model.tenant_id, agent_id=agent_id, slug=slug, session=db.session()
tenant_id=app_model.tenant_id, agent_id=agent_id, slug=slug, session=session
)
except SkillToolInferenceError as exc:
return {"code": exc.code, "message": exc.message}, exc.status_code
@ -336,12 +339,13 @@ class AgentLogApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model(mode=[AppMode.AGENT_CHAT])
def get(self, app_model: App):
def get(self, session: Session, app_model: App):
"""Get agent logs"""
args = AgentLogQuery.model_validate(request.args.to_dict(flat=True))
return AgentService.get_agent_logs(app_model, args.conversation_id, args.message_id, db.session())
return AgentService.get_agent_logs(app_model, args.conversation_id, args.message_id, session)
@console_ns.route("/agent/<uuid:agent_id>/skills/upload")
@ -356,9 +360,10 @@ class AgentSkillUploadByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _upload_skill_for_app(current_user=current_user, app_model=app_model)
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
return _upload_skill_for_app(session=session, current_user=current_user, app_model=app_model)
@console_ns.route("/apps/<uuid:app_id>/agent/skills/upload")
@ -378,11 +383,12 @@ class AgentSkillUploadApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
@with_current_user
def post(self, current_user: Account, app_model: App):
@with_session
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
def post(self, session: Session, current_user: Account, app_model: App):
"""Upload a Skill, validate it, and commit drive-backed skill files."""
return _upload_skill_for_app(current_user=current_user, app_model=app_model)
return _upload_skill_for_app(session=session, current_user=current_user, app_model=app_model)
@console_ns.route("/agent/<uuid:agent_id>/files")
@ -399,9 +405,12 @@ class AgentDriveFilesByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _commit_drive_file_for_app(current_user=current_user, app_model=app_model, allow_node_id=False)
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
return _commit_drive_file_for_app(
session=session, current_user=current_user, app_model=app_model, allow_node_id=False
)
@console_ns.doc("delete_agent_drive_file_by_agent")
@console_ns.doc(description="Delete one Agent App drive file by key")
@ -412,9 +421,12 @@ class AgentDriveFilesByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _delete_drive_file_for_app(current_user=current_user, app_model=app_model, allow_node_id=False)
@with_session
def delete(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
return _delete_drive_file_for_app(
session=session, current_user=current_user, app_model=app_model, allow_node_id=False
)
@console_ns.route("/apps/<uuid:app_id>/agent/files")
@ -429,11 +441,12 @@ class AgentDriveFilesApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
@with_current_user
def post(self, current_user: Account, app_model: App):
@with_session
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
def post(self, session: Session, current_user: Account, app_model: App):
"""ADD FILE: commit one uploaded file into the bound agent's drive."""
return _commit_drive_file_for_app(current_user=current_user, app_model=app_model)
return _commit_drive_file_for_app(session=session, current_user=current_user, app_model=app_model)
@console_ns.doc("delete_agent_drive_file")
@console_ns.doc(description="Delete one drive file by key via drive commit-null semantics")
@ -442,10 +455,11 @@ class AgentDriveFilesApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
@with_current_user
def delete(self, current_user: Account, app_model: App):
return _delete_drive_file_for_app(current_user=current_user, app_model=app_model)
@with_session
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
def delete(self, session: Session, current_user: Account, app_model: App):
return _delete_drive_file_for_app(session=session, current_user=current_user, app_model=app_model)
@console_ns.route("/agent/<uuid:agent_id>/skills/<string:slug>")
@ -459,9 +473,12 @@ class AgentSkillByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID, slug: str):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _delete_skill_for_app(current_user=current_user, app_model=app_model, slug=slug, allow_node_id=False)
@with_session
def delete(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, slug: str):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
return _delete_skill_for_app(
session=session, current_user=current_user, app_model=app_model, slug=slug, allow_node_id=False
)
@console_ns.route("/apps/<uuid:app_id>/agent/skills/<string:slug>")
@ -479,10 +496,11 @@ class AgentSkillApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
@with_current_user
def delete(self, current_user: Account, app_model: App, slug: str):
return _delete_skill_for_app(current_user=current_user, app_model=app_model, slug=slug)
@with_session
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
def delete(self, session: Session, current_user: Account, app_model: App, slug: str):
return _delete_skill_for_app(session=session, current_user=current_user, app_model=app_model, slug=slug)
@console_ns.route("/agent/<uuid:agent_id>/skills/<string:slug>/infer-tools")
@ -499,9 +517,10 @@ class AgentSkillInferToolsByAgentApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def post(self, tenant_id: str, agent_id: UUID, slug: str):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
return _infer_skill_tools_for_app(app_model=app_model, slug=slug)
@with_session(write=False)
def post(self, session: Session, tenant_id: str, agent_id: UUID, slug: str):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
return _infer_skill_tools_for_app(session=session, app_model=app_model, slug=slug)
@console_ns.route("/apps/<uuid:app_id>/agent/skills/<string:slug>/infer-tools")
@ -525,7 +544,8 @@ class AgentSkillInferToolsApi(Resource):
@setup_required
@login_required
@account_initialization_required
@with_session(write=False)
@get_app_model(mode=_WORKFLOW_AGENT_DRIVE_APP_MODES)
def post(self, app_model: App, slug: str):
def post(self, session: Session, app_model: App, slug: str):
"""Suggest CLI tools/env for a skill. Saving still goes through composer validation."""
return _infer_skill_tools_for_app(app_model=app_model, slug=slug)
return _infer_skill_tools_for_app(session=session, app_model=app_model, slug=slug)

10
api/controllers/console/app/agent_app_access.py
View File

@ -9,12 +9,13 @@ from uuid import UUID
from flask_restx import Resource
from pydantic import Field
from sqlalchemy.orm import Session
from controllers.common.schema import register_response_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.agent.app_helpers import resolve_agent_app_model
from controllers.console.wraps import account_initialization_required, setup_required, with_current_tenant_id
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.login import login_required
from services.agent.roster_service import AgentRosterService
@ -55,9 +56,10 @@ class AgentAppReferencingWorkflowsResource(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
workflows = AgentRosterService(db.session).list_workflows_referencing_app_agent(
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
workflows = AgentRosterService(session).list_workflows_referencing_app_agent(
tenant_id=tenant_id, app_id=app_model.id
)
return AgentReferencingWorkflowsResponse(

17
api/controllers/console/app/agent_app_feature.py
View File

@ -13,9 +13,11 @@ from uuid import UUID
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from controllers.common.fields import SimpleResultResponse
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.agent.app_helpers import resolve_agent_runtime_app_model
from controllers.console.wraps import (
@ -29,7 +31,6 @@ from controllers.console.wraps import (
with_current_user,
)
from events.app_event import app_model_config_was_updated
from extensions.ext_database import db
from libs.login import login_required
from models import Account
from models.agent_config_entities import (
@ -85,17 +86,23 @@ class AgentAppFeatureConfigResource(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
args = AgentAppFeaturesPayload.model_validate(console_ns.payload or {})
new_app_model_config = AgentAppFeatureConfigService.update_features(
app_model=app_model,
account=current_user,
config=args.model_dump(exclude_none=True),
session=db.session(),
session=session,
)
app_model_config_was_updated.send(app_model, app_model_config=new_app_model_config)
app_model_config_was_updated.send(
app_model,
app_model_config=new_app_model_config,
session=session,
)
session.commit()
return SimpleResultResponse(result="success").model_dump(mode="json")

22
api/controllers/console/app/agent_app_sandbox.py
View File

@ -14,6 +14,7 @@ from dify_agent.client import DifyAgentClientError, DifyAgentHTTPError, DifyAgen
from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from controllers.common.schema import (
query_params_from_model,
@ -21,6 +22,7 @@ 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.agent.app_helpers import resolve_agent_runtime_app_model
from controllers.console.app.wraps import get_app_model
@ -153,8 +155,9 @@ class AgentAppSandboxInfoResource(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
query = query_params_from_request(AgentSandboxInfoQuery)
try:
result = AgentAppSandboxService().get_info(
@ -177,8 +180,9 @@ class AgentAppSandboxListResource(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
query = query_params_from_request(AgentSandboxListQuery)
try:
result = AgentAppSandboxService().list_files(
@ -202,8 +206,9 @@ class AgentAppSandboxReadResource(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
query = query_params_from_request(AgentSandboxFileQuery)
try:
result = AgentAppSandboxService().read_file(
@ -227,8 +232,9 @@ class AgentAppSandboxUploadResource(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def post(self, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def post(self, session: Session, tenant_id: str, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
payload = AgentSandboxUploadPayload.model_validate(request.get_json(silent=True) or {})
try:
result = AgentAppSandboxService().upload_file(

194
api/controllers/console/app/agent_config_inspector.py
View File

@ -13,6 +13,7 @@ from uuid import UUID
from flask import Response, request, send_file, url_for
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from controllers.common.schema import (
query_params_from_model,
@ -20,6 +21,7 @@ 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.agent.app_helpers import resolve_agent_runtime_app_model
from controllers.console.app.wraps import get_app_model
@ -33,7 +35,6 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.login import login_required
from models.account import Account
@ -247,15 +248,15 @@ def _service() -> AgentConfigService:
return AgentConfigService()
def _resolve_agent_id(app_model: App, node_id: str | None) -> str | None:
def _resolve_agent_id(session: Session, app_model: App, node_id: str | None) -> str | None:
if node_id:
return AgentComposerService.resolve_workflow_node_agent_id(
session=session,
tenant_id=app_model.tenant_id,
app_id=app_model.id,
node_id=node_id,
session=db.session(),
)
return app_model.bound_agent_id
return app_model.bound_agent_id_with_session(session=session)
def _agent_not_bound() -> tuple[dict[str, object], int]:
@ -275,6 +276,7 @@ def _json_response(data: Mapping[str, Any]) -> Response:
def _resolve_console_version(
*,
session: Session,
tenant_id: str,
agent_id: str,
account_id: str,
@ -286,26 +288,24 @@ def _resolve_console_version(
try:
if draft_type == "debug_build":
state = AgentComposerService.load_agent_app_build_draft(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
account_id=account_id,
session=db.session(),
)
draft = state.get("draft") or {}
draft_id = draft.get("id")
if isinstance(draft_id, str) and draft_id:
return draft_id, AgentConfigVersionKind.BUILD_DRAFT
else:
state = AgentComposerService.load_agent_composer(
tenant_id=tenant_id, agent_id=agent_id, session=db.session()
)
state = AgentComposerService.load_agent_composer(session=session, tenant_id=tenant_id, agent_id=agent_id)
draft = state.get("draft") or {}
draft_id = draft.get("id")
if isinstance(draft_id, str) and draft_id:
# load_agent_composer creates the normal draft on first access.
# Config asset services use their own SQLAlchemy session, so the
# draft must be visible before we hand its id across that boundary.
db.session.commit()
session.commit()
return draft_id, AgentConfigVersionKind.DRAFT
except AgentVersionNotFoundError as exc:
raise AgentConfigServiceError(
@ -322,6 +322,7 @@ def _resolve_console_version(
def _resolve_target(
*,
session: Session,
tenant_id: str,
agent_id: str,
account_id: str,
@ -329,6 +330,7 @@ def _resolve_target(
draft_type: str | None,
) -> _ResolvedConsoleTarget:
resolved_version_id, version_kind = _resolve_console_version(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
account_id=account_id,
@ -346,13 +348,15 @@ def _resolve_target(
def _resolve_agent_route_target(
*,
session: Session,
tenant_id: str,
agent_id: UUID,
current_user: Account,
query: AgentConfigByAgentQuery,
) -> _ResolvedConsoleTarget:
resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
return _resolve_target(
session=session,
tenant_id=tenant_id,
agent_id=str(agent_id),
account_id=current_user.id,
@ -363,14 +367,16 @@ def _resolve_agent_route_target(
def _resolve_app_route_target(
*,
session: Session,
app_model: App,
current_user: Account,
query: AgentConfigQuery,
) -> _ResolvedConsoleTarget | tuple[dict[str, object], int]:
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
return _resolve_target(
session=session,
tenant_id=app_model.tenant_id,
agent_id=agent_id,
account_id=current_user.id,
@ -381,6 +387,7 @@ def _resolve_app_route_target(
def _with_agent_route_target(
*,
session: Session,
tenant_id: str,
agent_id: UUID,
current_user: Account,
@ -389,6 +396,7 @@ def _with_agent_route_target(
query = query_params_from_request(AgentConfigByAgentQuery)
try:
target = _resolve_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -401,13 +409,14 @@ def _with_agent_route_target(
def _with_app_route_target(
*,
session: Session,
app_model: App,
current_user: Account,
action: Callable[[_ResolvedConsoleTarget], Any],
) -> Any:
query = query_params_from_request(AgentConfigQuery)
try:
target = _resolve_app_route_target(app_model=app_model, current_user=current_user, query=query)
target = _resolve_app_route_target(session=session, app_model=app_model, current_user=current_user, query=query)
if isinstance(target, tuple):
return target
return action(target)
@ -649,8 +658,10 @@ class AgentConfigManifestByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -666,10 +677,13 @@ class AgentConfigManifestApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App):
return _with_app_route_target(app_model=app_model, current_user=current_user, action=_manifest_response)
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App):
return _with_app_route_target(
session=session, app_model=app_model, current_user=current_user, action=_manifest_response
)
@console_ns.route("/agent/<uuid:agent_id>/config/skills/upload")
@ -689,8 +703,10 @@ class AgentConfigSkillUploadByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -715,10 +731,13 @@ class AgentConfigSkillUploadApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def post(self, current_user: Account, app_model: App):
return _with_app_route_target(app_model=app_model, current_user=current_user, action=_skill_upload_response)
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def post(self, session: Session, current_user: Account, app_model: App):
return _with_app_route_target(
session=session, app_model=app_model, current_user=current_user, action=_skill_upload_response
)
@console_ns.route("/agent/<uuid:agent_id>/config/skills")
@ -731,8 +750,10 @@ class AgentConfigSkillsByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -748,10 +769,13 @@ class AgentConfigSkillsApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App):
return _with_app_route_target(app_model=app_model, current_user=current_user, action=_skill_list_response)
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App):
return _with_app_route_target(
session=session, app_model=app_model, current_user=current_user, action=_skill_list_response
)
@console_ns.route("/agent/<uuid:agent_id>/config/files")
@ -764,8 +788,10 @@ class AgentConfigFilesByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -781,9 +807,11 @@ class AgentConfigFilesByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID):
payload = AgentConfigFileUploadPayload.model_validate(console_ns.payload or {})
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -799,10 +827,13 @@ class AgentConfigFilesApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App):
return _with_app_route_target(app_model=app_model, current_user=current_user, action=_file_list_response)
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App):
return _with_app_route_target(
session=session, app_model=app_model, current_user=current_user, action=_file_list_response
)
@console_ns.doc("upload_agent_config_file")
@console_ns.doc(params={"app_id": "Application ID", **query_params_from_model(AgentConfigQuery)})
@ -813,11 +844,13 @@ class AgentConfigFilesApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def post(self, current_user: Account, app_model: App):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def post(self, session: Session, current_user: Account, app_model: App):
payload = AgentConfigFileUploadPayload.model_validate(console_ns.payload or {})
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _file_upload_response(target, payload),
@ -836,8 +869,10 @@ class AgentConfigSkillInspectByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -855,10 +890,12 @@ class AgentConfigSkillInspectApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _skill_inspect_response(target, name),
@ -883,10 +920,12 @@ class AgentConfigSkillFilePreviewByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
query = query_params_from_request(AgentConfigSkillFileByAgentQuery)
try:
target = _resolve_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -913,12 +952,15 @@ class AgentConfigSkillFilePreviewApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
query = query_params_from_request(AgentConfigSkillFileQuery)
try:
target = _resolve_app_route_target(app_model=app_model, current_user=current_user, query=query)
target = _resolve_app_route_target(
session=session, app_model=app_model, current_user=current_user, query=query
)
if isinstance(target, tuple):
return target
return _skill_file_preview_response(target, name, query.path)
@ -938,8 +980,10 @@ class AgentConfigSkillDownloadByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -957,10 +1001,12 @@ class AgentConfigSkillDownloadApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _skill_download_response(target, name),
@ -983,10 +1029,12 @@ class AgentConfigSkillFileDownloadByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
query = query_params_from_request(AgentConfigSkillFileByAgentQuery)
try:
target = _resolve_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -1017,12 +1065,15 @@ class AgentConfigSkillFileDownloadApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
query = query_params_from_request(AgentConfigSkillFileQuery)
try:
target = _resolve_app_route_target(app_model=app_model, current_user=current_user, query=query)
target = _resolve_app_route_target(
session=session, app_model=app_model, current_user=current_user, query=query
)
if isinstance(target, tuple):
return target
return _skill_file_download_response(
@ -1047,10 +1098,12 @@ class AgentConfigSkillFileDownloadContentByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
query = query_params_from_request(AgentConfigSkillFileByAgentQuery)
try:
target = _resolve_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -1069,12 +1122,15 @@ class AgentConfigSkillFileDownloadContentApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
query = query_params_from_request(AgentConfigSkillFileQuery)
try:
target = _resolve_app_route_target(app_model=app_model, current_user=current_user, query=query)
target = _resolve_app_route_target(
session=session, app_model=app_model, current_user=current_user, query=query
)
if isinstance(target, tuple):
return target
return _skill_file_raw_download_response(target, name, query.path)
@ -1094,8 +1150,10 @@ class AgentConfigSkillByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def delete(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -1115,10 +1173,12 @@ class AgentConfigSkillApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def delete(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def delete(self, session: Session, current_user: Account, app_model: App, name: str):
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _skill_delete_response(target, name),
@ -1137,8 +1197,10 @@ class AgentConfigFilePreviewByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -1156,10 +1218,12 @@ class AgentConfigFilePreviewApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _file_preview_response(target, name),
@ -1178,8 +1242,10 @@ class AgentConfigFileDownloadByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def get(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -1197,10 +1263,12 @@ class AgentConfigFileDownloadApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def get(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, session: Session, current_user: Account, app_model: App, name: str):
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _file_download_response(target, name),
@ -1219,8 +1287,10 @@ class AgentConfigFileByAgentApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
@with_session
def delete(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID, name: str):
return _with_agent_route_target(
session=session,
tenant_id=tenant_id,
agent_id=agent_id,
current_user=current_user,
@ -1240,10 +1310,12 @@ class AgentConfigFileApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model(mode=_WORKFLOW_APP_MODES)
@with_current_user
def delete(self, current_user: Account, app_model: App, name: str):
@with_session
@get_app_model(mode=_WORKFLOW_APP_MODES)
def delete(self, session: Session, current_user: Account, app_model: App, name: str):
return _with_app_route_target(
session=session,
app_model=app_model,
current_user=current_user,
action=lambda target: _file_delete_response(target, name),

81
api/controllers/console/app/agent_drive_inspector.py
View File

@ -18,17 +18,18 @@ from uuid import UUID
from flask import Response
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from controllers.common.schema import (
query_params_from_model,
query_params_from_request,
register_response_schema_models,
)
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.agent.app_helpers import resolve_agent_runtime_app_model
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import account_initialization_required, setup_required, with_current_tenant_id
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.login import login_required
from models.model import App, AppMode
@ -144,13 +145,13 @@ register_response_schema_models(
)
def _resolve_agent_id(app_model: App, node_id: str | None) -> str | None:
def _resolve_agent_id(session: Session, app_model: App, node_id: str | None) -> str | None:
"""Agent identity for the drive: app-bound agent, or the workflow node binding."""
if node_id:
return AgentComposerService.resolve_workflow_node_agent_id(
tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id, session=db.session()
session=session, tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id
)
return app_model.bound_agent_id
return app_model.bound_agent_id_with_session(session=session)
def _agent_not_bound() -> tuple[dict[str, object], int]:
@ -181,12 +182,13 @@ class AgentDriveListByAgentApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
query = query_params_from_request(AgentDriveListByAgentQuery)
resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
try:
items = AgentDriveService().manifest(
tenant_id=tenant_id, agent_id=str(agent_id), prefix=query.prefix, session=db.session()
tenant_id=tenant_id, agent_id=str(agent_id), prefix=query.prefix, session=session
)
except AgentDriveError as exc:
return _handle(exc)
@ -203,10 +205,11 @@ class AgentDriveSkillListByAgentApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
try:
items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=str(agent_id), session=db.session())
items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=str(agent_id), session=session)
except AgentDriveError as exc:
return _handle(exc)
return {"items": items}
@ -222,15 +225,16 @@ class AgentDriveSkillInspectByAgentApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID, skill_path: str):
resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID, skill_path: str):
resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
try:
return _json_response(
AgentDriveService().inspect_skill(
tenant_id=tenant_id,
agent_id=str(agent_id),
skill_path=skill_path,
session=db.session(),
session=session,
)
)
except AgentDriveError as exc:
@ -247,12 +251,13 @@ class AgentDrivePreviewByAgentApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
query = query_params_from_request(AgentDriveFileByAgentQuery)
resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
try:
return AgentDriveService().preview(
tenant_id=tenant_id, agent_id=str(agent_id), key=query.key, session=db.session()
tenant_id=tenant_id, agent_id=str(agent_id), key=query.key, session=session
)
except AgentDriveError as exc:
return _handle(exc)
@ -268,12 +273,13 @@ class AgentDriveDownloadByAgentApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, tenant_id: str, agent_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID):
query = query_params_from_request(AgentDriveFileByAgentQuery)
resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id)
resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id)
try:
url = AgentDriveService().download_url(
tenant_id=tenant_id, agent_id=str(agent_id), key=query.key, session=db.session()
tenant_id=tenant_id, agent_id=str(agent_id), key=query.key, session=session
)
except AgentDriveError as exc:
return _handle(exc)
@ -289,15 +295,16 @@ class AgentDriveListApi(Resource):
@setup_required
@login_required
@account_initialization_required
@with_session(write=False)
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, app_model: App):
def get(self, session: Session, app_model: App):
query = query_params_from_request(AgentDriveListQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
try:
items = AgentDriveService().manifest(
tenant_id=app_model.tenant_id, agent_id=agent_id, prefix=query.prefix, session=db.session()
tenant_id=app_model.tenant_id, agent_id=agent_id, prefix=query.prefix, session=session
)
except AgentDriveError as exc:
return _handle(exc)
@ -315,16 +322,15 @@ class AgentDriveSkillListApi(Resource):
@setup_required
@login_required
@account_initialization_required
@with_session(write=False)
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, app_model: App):
def get(self, session: Session, app_model: App):
query = query_params_from_request(AgentDriveListQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
try:
items = AgentDriveService().list_skills(
tenant_id=app_model.tenant_id, agent_id=agent_id, session=db.session()
)
items = AgentDriveService().list_skills(tenant_id=app_model.tenant_id, agent_id=agent_id, session=session)
except AgentDriveError as exc:
return _handle(exc)
return {"items": items}
@ -345,10 +351,11 @@ class AgentDriveSkillInspectApi(Resource):
@setup_required
@login_required
@account_initialization_required
@with_session(write=False)
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, app_model: App, skill_path: str):
def get(self, session: Session, app_model: App, skill_path: str):
query = query_params_from_request(AgentDriveSkillInspectQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
try:
@ -357,7 +364,7 @@ class AgentDriveSkillInspectApi(Resource):
tenant_id=app_model.tenant_id,
agent_id=agent_id,
skill_path=skill_path,
session=db.session(),
session=session,
)
)
except AgentDriveError as exc:
@ -373,15 +380,16 @@ class AgentDrivePreviewApi(Resource):
@setup_required
@login_required
@account_initialization_required
@with_session(write=False)
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, app_model: App):
def get(self, session: Session, app_model: App):
query = query_params_from_request(AgentDriveFileQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
try:
return AgentDriveService().preview(
tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key, session=db.session()
tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key, session=session
)
except AgentDriveError as exc:
return _handle(exc)
@ -396,15 +404,16 @@ class AgentDriveDownloadApi(Resource):
@setup_required
@login_required
@account_initialization_required
@with_session(write=False)
@get_app_model(mode=_WORKFLOW_APP_MODES)
def get(self, app_model: App):
def get(self, session: Session, app_model: App):
query = query_params_from_request(AgentDriveFileQuery)
agent_id = _resolve_agent_id(app_model, query.node_id)
agent_id = _resolve_agent_id(session, app_model, query.node_id)
if not agent_id:
return _agent_not_bound()
try:
url = AgentDriveService().download_url(
tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key, session=db.session()
tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key, session=session
)
except AgentDriveError as exc:
return _handle(exc)

69
api/controllers/console/app/annotation.py
View File

@ -5,10 +5,12 @@ from flask import abort, request
from flask_restx import Resource
from pydantic import BaseModel, Field, TypeAdapter, field_validator
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.errors import NoFileUploadedError, TooManyFilesError
from controllers.common.schema import query_params_from_model, 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,
@ -21,7 +23,6 @@ from controllers.console.wraps import (
rbac_permission_required,
setup_required,
)
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from fields.annotation_fields import (
Annotation,
@ -46,9 +47,9 @@ from services.annotation_service import (
from services.app_ref_service import AppRef, AppRefService
def _get_app_ref(app_id: str) -> AppRef:
def _get_app_ref(session: Session, app_id: str) -> AppRef:
_, current_tenant_id = current_account_with_tenant()
app = db.session.scalar(
app = session.scalar(
select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1)
)
if app is None:
@ -210,8 +211,9 @@ class AppAnnotationSettingDetailApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
def get(self, app_id: UUID):
result = AppAnnotationService.get_app_annotation_setting_by_app_id(str(app_id), session=db.session())
@with_session(write=False)
def get(self, session: Session, app_id: UUID):
result = AppAnnotationService.get_app_annotation_setting_by_app_id(str(app_id), session)
return dump_response(AnnotationSettingResponse, result), 200
@ -228,14 +230,15 @@ class AppAnnotationSettingUpdateApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
def post(self, app_id: UUID, annotation_setting_id: UUID):
@with_session
def post(self, session: Session, app_id: UUID, annotation_setting_id: UUID):
annotation_setting_id_str = str(annotation_setting_id)
args = AnnotationSettingUpdatePayload.model_validate(console_ns.payload)
setting_args: UpdateAnnotationSettingArgs = {"score_threshold": args.score_threshold}
result = AppAnnotationService.update_app_annotation_setting(
str(app_id), annotation_setting_id_str, setting_args, session=db.session()
str(app_id), annotation_setting_id_str, setting_args, session
)
return dump_response(AnnotationSettingResponse, result), 200
@ -286,14 +289,15 @@ class AnnotationApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
def get(self, app_id: UUID):
@with_session(write=False)
def get(self, session: Session, app_id: UUID):
args = AnnotationListQuery.model_validate(request.args.to_dict(flat=True))
page = args.page
limit = args.limit
keyword = args.keyword
annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id(
str(app_id), page, limit, keyword, session=db.session()
str(app_id), page, limit, keyword, session
)
annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True)
return AnnotationList(
@ -312,7 +316,8 @@ class AnnotationApi(Resource):
@cloud_edition_billing_resource_check("annotation")
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
def post(self, app_id: UUID):
@with_session
def post(self, session: Session, app_id: UUID):
args = CreateAnnotationPayload.model_validate(console_ns.payload)
upsert_args: UpsertAnnotationArgs = {}
if args.answer is not None:
@ -323,9 +328,7 @@ class AnnotationApi(Resource):
upsert_args["message_id"] = args.message_id
if args.question is not None:
upsert_args["question"] = args.question
annotation = AppAnnotationService.up_insert_app_annotation_from_message(
upsert_args, str(app_id), session=db.session()
)
annotation = AppAnnotationService.up_insert_app_annotation_from_message(upsert_args, str(app_id), session)
return dump_response(Annotation, annotation), 201
@setup_required
@ -334,7 +337,8 @@ class AnnotationApi(Resource):
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_CREATE_AND_MANAGEMENT)
@console_ns.response(204, "Annotations deleted successfully")
def delete(self, app_id: UUID):
@with_session
def delete(self, session: Session, app_id: UUID):
# Use request.args.getlist to get annotation_ids array directly
annotation_ids = request.args.getlist("annotation_id")
@ -348,12 +352,12 @@ class AnnotationApi(Resource):
"message": "annotation_ids are required if the parameter is provided.",
}, 400
app_ref = _get_app_ref(str(app_id))
AppAnnotationService.delete_app_annotations_in_batch(app_ref, annotation_ids, session=db.session())
app_ref = _get_app_ref(session, str(app_id))
AppAnnotationService.delete_app_annotations_in_batch(app_ref, annotation_ids, session)
return "", 204
# If no annotation_ids are provided, handle clearing all annotations
else:
AppAnnotationService.clear_all_annotations(str(app_id), session=db.session())
AppAnnotationService.clear_all_annotations(str(app_id), session)
return "", 204
@ -373,8 +377,9 @@ class AnnotationExportApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
def get(self, app_id: UUID):
annotation_list = AppAnnotationService.export_annotation_list_by_app_id(str(app_id), session=db.session())
@with_session(write=False)
def get(self, session: Session, app_id: UUID):
annotation_list = AppAnnotationService.export_annotation_list_by_app_id(str(app_id), session)
annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True)
return (
AnnotationExportList(data=annotation_models).model_dump(mode="json"),
@ -401,16 +406,17 @@ class AnnotationUpdateDeleteApi(Resource):
@cloud_edition_billing_resource_check("annotation")
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
def post(self, app_id: UUID, annotation_id: UUID):
@with_session
def post(self, session: Session, app_id: UUID, annotation_id: UUID):
args = UpdateAnnotationPayload.model_validate(console_ns.payload)
update_args: UpdateAnnotationArgs = {}
if args.answer is not None:
update_args["answer"] = args.answer
if args.question is not None:
update_args["question"] = args.question
app_ref = _get_app_ref(str(app_id))
app_ref = _get_app_ref(session, str(app_id))
annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id))
annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, db.session())
annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, session)
return Annotation.model_validate(annotation, from_attributes=True).model_dump(mode="json")
@setup_required
@ -419,10 +425,11 @@ class AnnotationUpdateDeleteApi(Resource):
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@console_ns.response(204, "Annotation deleted successfully")
def delete(self, app_id: UUID, annotation_id: UUID):
app_ref = _get_app_ref(str(app_id))
@with_session
def delete(self, session: Session, app_id: UUID, annotation_id: UUID):
app_ref = _get_app_ref(session, str(app_id))
annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id))
AppAnnotationService.delete_app_annotation(annotation_ref, db.session())
AppAnnotationService.delete_app_annotation(annotation_ref, session)
return "", 204
@ -446,7 +453,8 @@ class AnnotationBatchImportApi(Resource):
@annotation_import_concurrency_limit
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
def post(self, app_id: UUID):
@with_session
def post(self, session: Session, app_id: UUID):
from configs import dify_config
# check file
@ -481,7 +489,7 @@ class AnnotationBatchImportApi(Resource):
return dump_response(
AnnotationBatchImportResponse,
AppAnnotationService.batch_import_app_annotations(str(app_id), file, session=db.session()),
AppAnnotationService.batch_import_app_annotations(str(app_id), file, session),
)
@ -533,16 +541,17 @@ class AnnotationHitHistoryListApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
def get(self, app_id: UUID, annotation_id: UUID):
@with_session(write=False)
def get(self, session: Session, app_id: UUID, annotation_id: UUID):
page = request.args.get("page", default=1, type=int)
limit = request.args.get("limit", default=20, type=int)
app_ref = _get_app_ref(str(app_id))
app_ref = _get_app_ref(session, str(app_id))
annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id))
annotation_hit_history_list, total = AppAnnotationService.get_annotation_hit_histories(
annotation_ref,
page,
limit,
session=db.session(),
session,
)
history_models = TypeAdapter(list[AnnotationHitHistory]).validate_python(
annotation_hit_history_list, from_attributes=True

161
api/controllers/console/app/app.py
View File

@ -6,7 +6,7 @@ from typing import Any, Literal
from flask import request
from flask_restx import Resource
from pydantic import AliasChoices, BaseModel, Field, computed_field, field_validator
from pydantic import AliasChoices, BaseModel, Field, ValidationInfo, computed_field, field_validator, model_validator
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import BadRequest, NotFound
@ -52,7 +52,14 @@ from libs.login import login_required
from models import Account, App, DatasetPermissionEnum, Workflow
from models.model import IconType
from services.app_dsl_service import AppDslService
from services.app_service import AppListParams, AppListSortBy, AppService, CreateAppParams, StarredAppListParams
from services.app_service import (
AppListParams,
AppListSortBy,
AppResponseView,
AppService,
CreateAppParams,
StarredAppListParams,
)
from services.enterprise import rbac_service as enterprise_rbac_service
from services.enterprise.enterprise_service import EnterpriseService
from services.entities.dsl_entities import DslImportWarning, ImportMode, ImportStatus
@ -348,7 +355,18 @@ class DeletedTool(ResponseModel):
provider_id: str
class AppPartial(ResponseModel):
class AppResponseModel(ResponseModel):
@model_validator(mode="before")
@classmethod
def _use_request_session(cls, value: Any, info: ValidationInfo) -> Any:
if not isinstance(value, App):
return value
if info.context is None or "session" not in info.context:
raise ValueError("session context is required to serialize an App")
return AppResponseView(value, session=info.context["session"])
class AppPartial(AppResponseModel):
id: str
name: str
max_active_requests: int | None = None
@ -392,7 +410,7 @@ class AppPartial(ResponseModel):
return to_timestamp(value)
class AppDetail(ResponseModel):
class AppDetail(AppResponseModel):
id: str
name: str
description: str | None = None
@ -587,12 +605,13 @@ class AppListApi(Resource):
permissions = enterprise_rbac_service.RBACService.MyPermissions.get(
str(current_tenant_id),
current_user_id,
session=db.session(),
session=session,
)
if dify_config.RBAC_ENABLED:
access_filter = resolve_app_access_filter(
str(current_tenant_id),
current_user_id,
session=session,
permissions=permissions,
)
access_filter.apply_to_params(params)
@ -608,7 +627,11 @@ class AppListApi(Resource):
permission_keys_map = permissions.app.permission_keys_by_resource_ids(app_ids)
_enrich_app_list_items(session, apps=app_pagination.items, tenant_id=current_tenant_id)
pagination_model = AppPagination.model_validate(app_pagination, from_attributes=True)
pagination_model = AppPagination.model_validate(
app_pagination,
from_attributes=True,
context={"session": session},
)
if app_pagination.items:
pagination_model = pagination_model.model_copy(
update={
@ -634,7 +657,8 @@ class AppListApi(Resource):
@edit_permission_required
@with_current_user
@with_current_tenant_id
def post(self, current_tenant_id: str, current_user: Account):
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account):
"""Create app"""
args = CreateAppPayload.model_validate(console_ns.payload)
params = CreateAppParams(
@ -647,7 +671,7 @@ class AppListApi(Resource):
)
app_service = AppService()
app = app_service.create_app(current_tenant_id, params, current_user, session=db.session())
app = app_service.create_app(current_tenant_id, params, current_user, session=session)
if dify_config.RBAC_ENABLED:
enterprise_rbac_service.RBACService.AppAccess.replace_whitelist(
tenant_id=str(current_tenant_id),
@ -660,11 +684,13 @@ class AppListApi(Resource):
str(current_tenant_id),
current_user.id,
[str(app.id)],
session=db.session(),
)
app_detail = AppDetailWithSite.model_validate(app, from_attributes=True).model_copy(
update={"permission_keys": permission_keys_map.get(str(app.id), [])}
session=session,
)
app_detail = AppDetailWithSite.model_validate(
app,
from_attributes=True,
context={"session": session},
).model_copy(update={"permission_keys": permission_keys_map.get(str(app.id), [])})
return app_detail.model_dump(mode="json"), 201
@ -700,7 +726,14 @@ class StarredAppListApi(Resource):
return empty.model_dump(mode="json"), 200
_enrich_app_list_items(session, apps=app_pagination.items, tenant_id=current_tenant_id)
return AppPagination.model_validate(app_pagination, from_attributes=True).model_dump(mode="json"), 200
return (
AppPagination.model_validate(
app_pagination,
from_attributes=True,
context={"session": session},
).model_dump(mode="json"),
200,
)
@console_ns.route("/apps/<uuid:app_id>/star")
@ -751,12 +784,13 @@ class AppApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model(mode=None)
def get(self, current_tenant_id: str, current_user: Account, app_model: App):
def get(self, session: Session, current_tenant_id: str, current_user: Account, app_model: App):
"""Get app detail"""
app_service = AppService()
app_model = app_service.get_app(app_model)
app_model = app_service.get_app(app_model, session=session)
if FeatureService.get_system_features().webapp_auth.enabled:
app_setting = EnterpriseService.WebAppAuth.get_app_access_mode_by_id(app_id=str(app_model.id))
@ -766,13 +800,15 @@ class AppApi(Resource):
str(current_tenant_id),
current_user.id,
app_id=str(app_model.id),
session=db.session(),
session=session,
)
permission_keys_map = permissions.app.permission_keys_by_resource_ids([str(app_model.id)])
response_model = AppDetailWithSite.model_validate(app_model, from_attributes=True).model_copy(
update={"permission_keys": permission_keys_map.get(str(app_model.id), [])}
)
response_model = AppDetailWithSite.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_copy(update={"permission_keys": permission_keys_map.get(str(app_model.id), [])})
return response_model.model_dump(mode="json")
@console_ns.doc("update_app")
@ -787,8 +823,9 @@ class AppApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_session
@get_app_model(mode=None)
def put(self, app_model: App):
def put(self, session: Session, app_model: App):
"""Update app"""
args = UpdateAppPayload.model_validate(console_ns.payload)
@ -803,8 +840,12 @@ class AppApi(Resource):
"use_icon_as_answer_icon": args.use_icon_as_answer_icon or False,
"max_active_requests": args.max_active_requests or 0,
}
app_model = app_service.update_app(app_model, args_dict, session=db.session())
return dump_response(AppDetailWithSite, app_model)
app_model = app_service.update_app(app_model, args_dict, session=session)
return AppDetailWithSite.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
@console_ns.doc("delete_app")
@console_ns.doc(description="Delete application")
@ -816,11 +857,12 @@ class AppApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_DELETE)
@with_session
@get_app_model
def delete(self, app_model: App):
def delete(self, session: Session, app_model: App):
"""Delete app"""
app_service = AppService()
app_service.delete_app(app_model, session=db.session())
app_service.delete_app(app_model, session=session)
return "", 204
@ -883,20 +925,21 @@ class AppCopyApi(Resource):
stmt = select(App).where(App.id == result.app_id)
app = session.scalar(stmt)
if not app:
raise NotFound("App not found")
if not app:
raise NotFound("App not found")
permission_keys_map = enterprise_rbac_service.RBACService.AppPermissions.batch_get(
str(current_tenant_id),
current_user.id,
[str(app.id)],
session=db.session(),
)
response_model = AppDetailWithSite.model_validate(app, from_attributes=True).model_copy(
update={"permission_keys": permission_keys_map.get(str(app.id), [])}
)
return response_model.model_dump(mode="json"), 201
permission_keys_map = enterprise_rbac_service.RBACService.AppPermissions.batch_get(
str(current_tenant_id),
current_user.id,
[str(app.id)],
session=session,
)
response_model = AppDetailWithSite.model_validate(
app,
from_attributes=True,
context={"session": session},
).model_copy(update={"permission_keys": permission_keys_map.get(str(app.id), [])})
return response_model.model_dump(mode="json"), 201
@console_ns.route("/apps/<uuid:app_id>/export")
@ -966,13 +1009,18 @@ class AppNameApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_session
@get_app_model(mode=None)
def post(self, app_model: App):
def post(self, session: Session, app_model: App):
args = AppNamePayload.model_validate(console_ns.payload)
app_service = AppService()
app_model = app_service.update_app_name(app_model, args.name, session=db.session())
return dump_response(AppDetail, app_model)
app_model = app_service.update_app_name(app_model, args.name, session=session)
return AppDetail.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/icon")
@ -988,8 +1036,9 @@ class AppIconApi(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_session
@get_app_model(mode=None)
def post(self, app_model: App):
def post(self, session: Session, app_model: App):
args = AppIconPayload.model_validate(console_ns.payload or {})
app_service = AppService()
@ -998,9 +1047,13 @@ class AppIconApi(Resource):
args.icon or "",
args.icon_background or "",
args.icon_type,
session=db.session(),
session=session,
)
return dump_response(AppDetail, app_model)
return AppDetail.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/site-enable")
@ -1016,13 +1069,18 @@ class AppSiteStatus(Resource):
@account_initialization_required
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
@with_session
@get_app_model(mode=None)
def post(self, app_model: App):
def post(self, session: Session, app_model: App):
args = AppSiteStatusPayload.model_validate(console_ns.payload)
app_service = AppService()
app_model = app_service.update_app_site_status(app_model, args.enable_site, session=db.session())
return dump_response(AppDetail, app_model)
app_model = app_service.update_app_site_status(app_model, args.enable_site, session=session)
return AppDetail.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/api-enable")
@ -1038,13 +1096,18 @@ class AppApiStatus(Resource):
@is_admin_or_owner_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
@with_session
@get_app_model(mode=None)
def post(self, app_model: App):
def post(self, session: Session, app_model: App):
args = AppApiStatusPayload.model_validate(console_ns.payload)
app_service = AppService()
app_model = app_service.update_app_api_status(app_model, args.enable_api, session=db.session())
return dump_response(AppDetail, app_model)
app_model = app_service.update_app_api_status(app_model, args.enable_api, session=session)
return AppDetail.model_validate(
app_model,
from_attributes=True,
context={"session": session},
).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/trace")

16
api/controllers/console/app/audio.py
View File

@ -127,6 +127,7 @@ def _transcribe_audio_to_text(
*,
app_model: App,
file: FileStorage | None,
session: Session,
agent_soul: AgentSoulConfig | None = None,
) -> dict[str, str]:
try:
@ -134,6 +135,7 @@ def _transcribe_audio_to_text(
response = AudioService.transcript_asr(
app_model=app_model,
file=file,
session=session,
end_user=None,
)
else:
@ -141,6 +143,7 @@ def _transcribe_audio_to_text(
app_model=app_model,
agent_soul=agent_soul,
file=file,
session=session,
end_user=None,
)
return dump_response(AudioTranscriptResponse, response)
@ -194,7 +197,11 @@ class ChatMessageAudioApi(Resource):
@account_initialization_required
@get_app_model(mode=_CONSOLE_AUDIO_TRANSCRIPT_APP_MODES)
def post(self, app_model: App):
return _transcribe_audio_to_text(app_model=app_model, file=request.files.get("file"))
return _transcribe_audio_to_text(
app_model=app_model,
file=request.files.get("file"),
session=db.session(),
)
@console_ns.route("/agent/<uuid:agent_id>/audio-to-text")
@ -228,7 +235,11 @@ class AgentChatMessageAudioApi(Resource):
agent_id: UUID,
):
payload = AgentAudioTranscriptFormPayload.model_validate(request.form.to_dict(flat=True))
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
app_model = resolve_agent_runtime_app_model(
session=session,
tenant_id=current_tenant_id,
agent_id=agent_id,
)
# Agent routes expose Agent ids, while APP RBAC is keyed by the resolved runtime App id.
enforce_rbac_access(
tenant_id=current_tenant_id,
@ -248,6 +259,7 @@ class AgentChatMessageAudioApi(Resource):
app_model=app_model,
agent_soul=agent_soul,
file=request.files.get("file"),
session=session,
)

28
api/controllers/console/app/completion.py
View File

@ -44,7 +44,6 @@ from core.errors.error import (
QuotaExceededError,
)
from core.helper.trace_id_helper import get_external_trace_id
from extensions.ext_database import db
from graphon.model_runtime.errors.invoke import InvokeError
from libs import helper
from libs.helper import uuid_value
@ -158,8 +157,8 @@ class CompletionMessageApi(Resource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
@get_app_model(mode=AppMode.COMPLETION)
@with_session
@get_app_model(mode=AppMode.COMPLETION)
def post(self, session: Session, current_user: Account, app_model: App):
args_model = CompletionMessagePayload.model_validate(console_ns.payload)
args = args_model.model_dump(exclude_none=True, by_alias=True)
@ -240,8 +239,8 @@ class ChatMessageApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.AGENT])
@with_session
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.AGENT])
def post(self, session: Session, current_tenant_id: str, current_user: Account, app_model: App):
return _create_chat_message(
session=session, current_tenant_id=current_tenant_id, current_user=current_user, app_model=app_model
@ -266,7 +265,9 @@ class AgentChatMessageApi(Resource):
@with_current_tenant_id
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
app_model = AgentRosterService(session).get_agent_runtime_app_model(
tenant_id=current_tenant_id, agent_id=str(agent_id)
)
return _create_chat_message(
session=session,
current_tenant_id=current_tenant_id,
@ -293,7 +294,9 @@ class AgentBuildChatFinalizeApi(Resource):
@with_current_tenant_id
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
app_model = AgentRosterService(session).get_agent_runtime_app_model(
tenant_id=current_tenant_id, agent_id=str(agent_id)
)
return _create_build_chat_finalization_message(
session=session,
current_tenant_id=current_tenant_id,
@ -329,15 +332,20 @@ class AgentChatMessageStopApi(Resource):
@account_initialization_required
@with_current_user_id
@with_current_tenant_id
def post(self, current_tenant_id: str, current_user_id: str, agent_id: UUID, task_id: str):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
@with_session(write=False)
def post(self, session: Session, current_tenant_id: str, current_user_id: str, agent_id: UUID, task_id: str):
app_model = resolve_agent_runtime_app_model(
session=session,
tenant_id=current_tenant_id,
agent_id=agent_id,
)
return _stop_chat_message(current_user_id=current_user_id, app_model=app_model, task_id=task_id)
def _resolve_current_user_agent_debug_conversation_id(
*, current_tenant_id: str, current_user: Account, app_model: App, agent_id: str | None
*, session: Session, current_tenant_id: str, current_user: Account, app_model: App, agent_id: str | None
) -> str:
roster_service = AgentRosterService(db.session)
roster_service = AgentRosterService(session)
if agent_id:
return roster_service.get_or_create_agent_app_debug_conversation_id(
tenant_id=current_tenant_id,
@ -369,6 +377,7 @@ def _create_chat_message(
if AppMode.value_of(app_model.mode) == AppMode.AGENT:
debug_conversation_id = _resolve_current_user_agent_debug_conversation_id(
session=session,
current_tenant_id=current_tenant_id or app_model.tenant_id,
current_user=current_user,
app_model=app_model,
@ -404,6 +413,7 @@ def _create_build_chat_finalization_message(
*, session: Session, current_user: Account, app_model: App, current_tenant_id: str, agent_id: str
):
debug_conversation_id = _resolve_current_user_agent_debug_conversation_id(
session=session,
current_tenant_id=current_tenant_id,
current_user=current_user,
app_model=app_model,

73
api/controllers/console/app/conversation.py
View File

@ -6,10 +6,11 @@ from flask import abort, request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import func, or_
from sqlalchemy.orm import selectinload
from sqlalchemy.orm import Session, selectinload
from werkzeug.exceptions import NotFound
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import (
@ -22,7 +23,6 @@ from controllers.console.wraps import (
with_current_user,
)
from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from fields.conversation_fields import (
Conversation as ConversationResponse,
)
@ -35,6 +35,7 @@ from fields.conversation_fields import (
from fields.conversation_fields import (
ConversationPagination as ConversationPaginationResponse,
)
from fields.conversation_fields import ConversationResponseSource
from fields.conversation_fields import (
ConversationWithSummaryPagination as ConversationWithSummaryPaginationResponse,
)
@ -105,8 +106,9 @@ class CompletionConversationApi(Resource):
@edit_permission_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model(mode=AppMode.COMPLETION)
def get(self, current_user: Account, app_model: App):
def get(self, session: Session, current_user: Account, app_model: App):
args = CompletionConversationQuery.model_validate(request.args.to_dict(flat=True))
query = sa.select(Conversation).where(
@ -157,9 +159,18 @@ class CompletionConversationApi(Resource):
query = query.order_by(Conversation.created_at.desc())
conversations = paginate_query(query, page=args.page, per_page=args.limit)
conversations = paginate_query(query, session=session, page=args.page, per_page=args.limit)
return dump_response(ConversationPaginationResponse, conversations)
return dump_response(
ConversationPaginationResponse,
{
"page": conversations.page,
"per_page": conversations.per_page,
"total": conversations.total,
"has_next": conversations.has_next,
"items": [ConversationResponseSource(item, session=session) for item in conversations.items],
},
)
@console_ns.route("/apps/<uuid:app_id>/completion-conversations/<uuid:conversation_id>")
@ -176,11 +187,15 @@ class CompletionConversationDetailApi(Resource):
@edit_permission_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session
@get_app_model(mode=AppMode.COMPLETION)
def get(self, current_user: Account, app_model: App, conversation_id: UUID):
def get(self, session: Session, current_user: Account, app_model: App, conversation_id: UUID):
conversation_id_str = str(conversation_id)
return dump_response(
ConversationMessageDetailResponse, _get_conversation(current_user, app_model, conversation_id_str)
ConversationMessageDetailResponse,
ConversationResponseSource(
_get_conversation(session, current_user, app_model, conversation_id_str), session=session
),
)
@console_ns.doc("delete_completion_conversation")
@ -195,12 +210,13 @@ class CompletionConversationDetailApi(Resource):
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_current_user
@with_session
@get_app_model(mode=AppMode.COMPLETION)
def delete(self, current_user: Account, app_model: App, conversation_id: UUID):
def delete(self, session: Session, current_user: Account, app_model: App, conversation_id: UUID):
conversation_id_str = str(conversation_id)
try:
ConversationService.delete(app_model, conversation_id_str, current_user, session=db.session())
ConversationService.delete(app_model, conversation_id_str, current_user, session=session)
except ConversationNotExistsError:
raise NotFound("Conversation Not Exists.")
@ -220,8 +236,9 @@ class ChatConversationApi(Resource):
@edit_permission_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT])
def get(self, current_user: Account, app_model: App):
def get(self, session: Session, current_user: Account, app_model: App):
args = ChatConversationQuery.model_validate(request.args.to_dict(flat=True))
subquery = (
@ -311,9 +328,18 @@ class ChatConversationApi(Resource):
case _:
query = query.order_by(Conversation.created_at.desc())
conversations = paginate_query(query, page=args.page, per_page=args.limit)
conversations = paginate_query(query, session=session, page=args.page, per_page=args.limit)
return dump_response(ConversationWithSummaryPaginationResponse, conversations)
return dump_response(
ConversationWithSummaryPaginationResponse,
{
"page": conversations.page,
"per_page": conversations.per_page,
"total": conversations.total,
"has_next": conversations.has_next,
"items": [ConversationResponseSource(item, session=session) for item in conversations.items],
},
)
@console_ns.route("/apps/<uuid:app_id>/chat-conversations/<uuid:conversation_id>")
@ -330,11 +356,15 @@ class ChatConversationDetailApi(Resource):
@edit_permission_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT])
def get(self, current_user: Account, app_model: App, conversation_id: UUID):
def get(self, session: Session, current_user: Account, app_model: App, conversation_id: UUID):
conversation_id_str = str(conversation_id)
return dump_response(
ConversationDetailResponse, _get_conversation(current_user, app_model, conversation_id_str)
ConversationDetailResponse,
ConversationResponseSource(
_get_conversation(session, current_user, app_model, conversation_id_str), session=session
),
)
@console_ns.doc("delete_chat_conversation")
@ -349,27 +379,28 @@ class ChatConversationDetailApi(Resource):
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_current_user
@with_session
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT])
def delete(self, current_user: Account, app_model: App, conversation_id: UUID):
def delete(self, session: Session, current_user: Account, app_model: App, conversation_id: UUID):
conversation_id_str = str(conversation_id)
try:
ConversationService.delete(app_model, conversation_id_str, current_user, session=db.session())
ConversationService.delete(app_model, conversation_id_str, current_user, session=session)
except ConversationNotExistsError:
raise NotFound("Conversation Not Exists.")
return "", 204
def _get_conversation(current_user: Account, app_model, conversation_id):
conversation = db.session.scalar(
def _get_conversation(session: Session, current_user: Account, app_model, conversation_id):
conversation = session.scalar(
sa.select(Conversation).where(Conversation.id == conversation_id, Conversation.app_id == app_model.id).limit(1)
)
if not conversation:
raise NotFound("Conversation Not Exists.")
db.session.execute(
session.execute(
sa.update(Conversation)
.where(Conversation.id == conversation_id, Conversation.read_at.is_(None))
# Keep updated_at unchanged when only marking a conversation as read.
@ -379,7 +410,7 @@ def _get_conversation(current_user: Account, app_model, conversation_id):
updated_at=Conversation.updated_at,
)
)
db.session.commit()
db.session.refresh(conversation)
session.flush()
session.refresh(conversation)
return conversation

115
api/controllers/console/app/message.py
View File

@ -6,11 +6,13 @@ from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import exists, func, select
from sqlalchemy.orm import Session
from werkzeug.exceptions import InternalServerError, NotFound
from controllers.common.controller_schemas import MessageFeedbackPayload as _MessageFeedbackPayloadBase
from controllers.common.fields import SimpleResultResponse, TextFileResponse
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.agent.app_helpers import resolve_agent_runtime_app_model
from controllers.console.app.error import (
@ -39,6 +41,7 @@ from fields.base import ResponseModel
from fields.conversation_fields import (
MessageDetail as BaseMessageDetailResponse,
)
from fields.conversation_fields import MessageResponseSource
from graphon.model_runtime.errors.invoke import InvokeError
from libs.helper import dump_response, uuid_value
from libs.infinite_scroll_pagination import InfiniteScrollPagination
@ -152,9 +155,10 @@ class ChatMessageListApi(Resource):
@edit_permission_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT])
def get(self, current_user: Account, app_model: App):
return _list_chat_messages(app_model=app_model, current_user=current_user)
def get(self, session: Session, current_user: Account, app_model: App):
return _list_chat_messages(session=session, app_model=app_model, current_user=current_user)
@console_ns.route("/agent/<uuid:agent_id>/chat-messages")
@ -172,9 +176,14 @@ class AgentChatMessageListApi(Resource):
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
return _list_chat_messages(app_model=app_model, current_user=current_user)
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(
session=session,
tenant_id=current_tenant_id,
agent_id=agent_id,
)
return _list_chat_messages(session=session, app_model=app_model, current_user=current_user)
@console_ns.route("/apps/<uuid:app_id>/feedbacks")
@ -190,9 +199,10 @@ class MessageFeedbackApi(Resource):
@login_required
@account_initialization_required
@with_current_user
@with_session
@get_app_model
def post(self, current_user: Account, app_model: App):
return _update_message_feedback(current_user=current_user, app_model=app_model)
def post(self, session: Session, current_user: Account, app_model: App):
return _update_message_feedback(session=session, current_user=current_user, app_model=app_model)
@console_ns.route("/agent/<uuid:agent_id>/feedbacks")
@ -208,9 +218,14 @@ class AgentMessageFeedbackApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def post(self, current_tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
return _update_message_feedback(current_user=current_user, app_model=app_model)
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account, agent_id: UUID):
app_model = resolve_agent_runtime_app_model(
session=session,
tenant_id=current_tenant_id,
agent_id=agent_id,
)
return _update_message_feedback(session=session, current_user=current_user, app_model=app_model)
@console_ns.route("/apps/<uuid:app_id>/annotations/count")
@ -252,9 +267,12 @@ class MessageSuggestedQuestionApi(Resource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT])
def get(self, current_user: Account, app_model: App, message_id: UUID):
return _get_message_suggested_questions(current_user=current_user, app_model=app_model, message_id=message_id)
def get(self, session: Session, current_user: Account, app_model: App, message_id: UUID):
return _get_message_suggested_questions(
session=session, current_user=current_user, app_model=app_model, message_id=message_id
)
@console_ns.route("/agent/<uuid:agent_id>/chat-messages/<uuid:message_id>/suggested-questions")
@ -273,9 +291,16 @@ class AgentMessageSuggestedQuestionApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account, agent_id: UUID, message_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
return _get_message_suggested_questions(current_user=current_user, app_model=app_model, message_id=message_id)
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account, agent_id: UUID, message_id: UUID):
app_model = resolve_agent_runtime_app_model(
session=session,
tenant_id=current_tenant_id,
agent_id=agent_id,
)
return _get_message_suggested_questions(
session=session, current_user=current_user, app_model=app_model, message_id=message_id
)
@console_ns.route("/apps/<uuid:app_id>/feedbacks/export")
@ -333,9 +358,10 @@ class MessageApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
@with_session(write=False)
@get_app_model
def get(self, app_model: App, message_id: UUID):
return _get_message_detail(app_model=app_model, message_id=message_id)
def get(self, session: Session, app_model: App, message_id: UUID):
return _get_message_detail(session=session, app_model=app_model, message_id=message_id)
@console_ns.route("/agent/<uuid:agent_id>/messages/<uuid:message_id>")
@ -349,12 +375,17 @@ class AgentMessageApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, current_tenant_id: str, agent_id: UUID, message_id: UUID):
app_model = resolve_agent_runtime_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
return _get_message_detail(app_model=app_model, message_id=message_id)
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, agent_id: UUID, message_id: UUID):
app_model = resolve_agent_runtime_app_model(
session=session,
tenant_id=current_tenant_id,
agent_id=agent_id,
)
return _get_message_detail(session=session, app_model=app_model, message_id=message_id)
def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
def _list_chat_messages(*, session: Session, app_model: App, current_user: Account | None = None):
args = ChatMessagesQuery.model_validate(request.args.to_dict())
if AppMode.value_of(app_model.mode) == AppMode.AGENT and current_user is not None:
@ -363,12 +394,12 @@ def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
app_model=app_model,
conversation_id=args.conversation_id,
user=current_user,
session=db.session(),
session=session,
)
except ConversationNotExistsError:
raise NotFound("Conversation Not Exists.")
else:
conversation = db.session.scalar(
conversation = session.scalar(
select(Conversation)
.where(Conversation.id == args.conversation_id, Conversation.app_id == app_model.id)
.limit(1)
@ -378,14 +409,14 @@ def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
raise NotFound("Conversation Not Exists.")
if args.first_id:
first_message = db.session.scalar(
first_message = session.scalar(
select(Message).where(Message.conversation_id == conversation.id, Message.id == args.first_id).limit(1)
)
if not first_message:
raise NotFound("First message not found")
history_messages = db.session.scalars(
history_messages = session.scalars(
select(Message)
.where(
Message.conversation_id == conversation.id,
@ -396,7 +427,7 @@ def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
.limit(args.limit)
).all()
else:
history_messages = db.session.scalars(
history_messages = session.scalars(
select(Message)
.where(Message.conversation_id == conversation.id)
.order_by(Message.created_at.desc())
@ -407,7 +438,7 @@ def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
if len(history_messages) == args.limit:
current_page_first_message = history_messages[-1]
# Check if there are more messages before the current page
has_more = db.session.scalar(
has_more = session.scalar(
select(
exists().where(
Message.conversation_id == conversation.id,
@ -425,26 +456,28 @@ def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
return dump_response(
MessageInfiniteScrollPaginationResponse,
InfiniteScrollPagination(data=history_messages, limit=args.limit, has_more=has_more),
InfiniteScrollPagination(
data=[MessageResponseSource(message, session=session) for message in history_messages],
limit=args.limit,
has_more=has_more,
),
)
def _update_message_feedback(*, current_user: Account, app_model: App):
def _update_message_feedback(*, session: Session, current_user: Account, app_model: App):
args = MessageFeedbackPayload.model_validate(console_ns.payload)
message_id = args.message_id
message = db.session.scalar(
select(Message).where(Message.id == message_id, Message.app_id == app_model.id).limit(1)
)
message = session.scalar(select(Message).where(Message.id == message_id, Message.app_id == app_model.id).limit(1))
if not message:
raise NotFound("Message Not Exists.")
feedback = message.admin_feedback
feedback = message.admin_feedback_with_session(session=session)
if not args.rating and feedback:
db.session.delete(feedback)
session.delete(feedback)
elif args.rating and feedback:
feedback.rating = FeedbackRating(args.rating)
feedback.content = args.content
@ -463,14 +496,14 @@ def _update_message_feedback(*, current_user: Account, app_model: App):
from_source=FeedbackFromSource.ADMIN,
from_account_id=current_user.id,
)
db.session.add(feedback)
session.add(feedback)
db.session.commit()
session.commit()
return SimpleResultResponse(result="success").model_dump(mode="json")
def _get_message_suggested_questions(*, current_user: Account, app_model: App, message_id: UUID):
def _get_message_suggested_questions(*, session: Session, current_user: Account, app_model: App, message_id: UUID):
message_id_str = str(message_id)
try:
@ -479,7 +512,7 @@ def _get_message_suggested_questions(*, current_user: Account, app_model: App, m
message_id=message_id_str,
user=current_user,
invoke_from=InvokeFrom.DEBUGGER,
session=db.session(),
session=session,
)
except MessageNotExistsError:
raise NotFound("Message not found")
@ -502,10 +535,10 @@ def _get_message_suggested_questions(*, current_user: Account, app_model: App, m
return dump_response(SuggestedQuestionsResponse, {"data": questions})
def _get_message_detail(*, app_model: App, message_id: UUID):
def _get_message_detail(*, session: Session, app_model: App, message_id: UUID):
message_id_str = str(message_id)
message = db.session.scalar(
message = session.scalar(
select(Message).where(Message.id == message_id_str, Message.app_id == app_model.id).limit(1)
)
@ -513,4 +546,4 @@ def _get_message_detail(*, app_model: App, message_id: UUID):
raise NotFound("Message Not Exists.")
attach_message_extra_contents([message])
return dump_response(MessageDetailResponse, message)
return dump_response(MessageDetailResponse, MessageResponseSource(message, session=session))

25
api/controllers/console/app/model_config.py
View File

@ -4,9 +4,11 @@ from typing import Any, cast
from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from controllers.common.fields import SimpleResultResponse
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.app.wraps import get_app_model
from controllers.console.wraps import (
@ -23,7 +25,6 @@ from core.agent.entities import AgentToolEntity
from core.tools.tool_manager import ToolManager
from core.tools.utils.configuration import ToolParameterConfigurationManager
from events.app_event import app_model_config_was_updated
from extensions.ext_database import db
from libs.datetime_utils import naive_utc_now
from libs.login import login_required
from models.model import App, AppMode, AppModelConfig
@ -93,14 +94,16 @@ class ModelConfigResource(Resource):
@account_initialization_required
@with_current_user_id
@with_current_tenant_id
@with_session
@get_app_model(mode=[AppMode.AGENT_CHAT, AppMode.CHAT, AppMode.COMPLETION])
def post(self, current_tenant_id: str, current_user_id: str, app_model: App):
"""Modify app model config"""
def post(self, session: Session, current_tenant_id: str, current_user_id: str, app_model: App):
"""Modify the app model config and dataset joins in one request transaction."""
# validate config
model_configuration = AppModelConfigService.validate_configuration(
tenant_id=current_tenant_id,
config=cast(dict, request.json),
app_mode=AppMode.value_of(app_model.mode),
session=session,
)
new_app_model_config = AppModelConfig(
@ -110,9 +113,8 @@ class ModelConfigResource(Resource):
)
new_app_model_config = new_app_model_config.from_model_config_dict(model_configuration)
if app_model.mode == AppMode.AGENT_CHAT or app_model.is_agent:
# get original app model config
original_app_model_config = db.session.get(AppModelConfig, app_model.app_model_config_id)
if app_model.mode == AppMode.AGENT_CHAT or app_model.is_agent_with_session(session=session):
original_app_model_config = app_model.app_model_config_with_session(session=session)
if original_app_model_config is None:
raise ValueError("Original app model config not found")
agent_mode = original_app_model_config.agent_mode_dict
@ -204,14 +206,17 @@ class ModelConfigResource(Resource):
# update app model config
new_app_model_config.agent_mode = json.dumps(agent_mode)
db.session.add(new_app_model_config)
db.session.flush()
session.add(new_app_model_config)
session.flush()
app_model.app_model_config_id = new_app_model_config.id
app_model.updated_by = current_user_id
app_model.updated_at = naive_utc_now()
db.session.commit()
app_model_config_was_updated.send(app_model, app_model_config=new_app_model_config)
app_model_config_was_updated.send(
app_model,
app_model_config=new_app_model_config,
session=session,
)
return {"result": "success"}

19
api/controllers/console/app/site.py
View File

@ -3,10 +3,12 @@ from typing import Literal
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from constants.languages import supported_language
from controllers.common.schema import register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import (
@ -19,7 +21,6 @@ from controllers.console.wraps import (
setup_required,
with_current_user,
)
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
@ -94,10 +95,11 @@ class AppSite(Resource):
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
@account_initialization_required
@with_current_user
@with_session
@get_app_model
def post(self, current_user: Account, app_model: App):
def post(self, session: Session, current_user: Account, app_model: App):
args = AppSiteUpdatePayload.model_validate(console_ns.payload or {})
site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
site = session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
if not site:
raise NotFound
@ -126,7 +128,7 @@ class AppSite(Resource):
site.updated_by = current_user.id
site.updated_at = naive_utc_now()
db.session.commit()
session.flush()
return dump_response(AppSiteResponse, site)
@ -145,16 +147,17 @@ class AppSiteAccessTokenReset(Resource):
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
@account_initialization_required
@with_current_user
@with_session
@get_app_model
def post(self, current_user: Account, app_model: App):
site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
def post(self, session: Session, current_user: Account, app_model: App):
site = session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
if not site:
raise NotFound
site.code = Site.generate_code(16)
site.code = Site.generate_code(16, session=session)
site.updated_by = current_user.id
site.updated_at = naive_utc_now()
db.session.commit()
session.flush()
return dump_response(AppSiteResponse, site)

39
api/controllers/console/app/workflow.py
View File

@ -324,6 +324,27 @@ class WorkflowResponse(ResponseModel):
return [_serialize_environment_variable(item) for item in value]
class _WorkflowResponseSource:
def __init__(self, workflow: Workflow, *, session: Session) -> None:
self._workflow = workflow
self._session = session
def __getattr__(self, name: str) -> object:
return getattr(self._workflow, name) # noqa: no-new-getattr response adapter delegates model fields
@property
def created_by_account(self) -> Account | None:
return self._workflow.get_created_by_account(session=self._session)
@property
def updated_by_account(self) -> Account | None:
return self._workflow.get_updated_by_account(session=self._session)
@property
def tool_published(self) -> bool:
return self._workflow.get_tool_published(session=self._session)
class WorkflowPaginationResponse(ResponseModel):
items: list[WorkflowResponse]
page: int
@ -629,10 +650,10 @@ class AdvancedChatDraftWorkflowRunApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
@get_app_model(mode=[AppMode.ADVANCED_CHAT])
@with_current_user
@edit_permission_required
@with_session
@get_app_model(mode=[AppMode.ADVANCED_CHAT])
def post(self, session: Session, current_user: Account, app_model: App):
"""
Run draft workflow
@ -1074,10 +1095,10 @@ class DraftWorkflowRunApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
@get_app_model(mode=[AppMode.WORKFLOW])
@with_current_user
@edit_permission_required
@with_session
@get_app_model(mode=[AppMode.WORKFLOW])
def post(self, session: Session, current_user: Account, app_model: App):
"""
Run draft workflow
@ -1438,7 +1459,7 @@ class PublishedAllWorkflowApi(Resource):
)
return WorkflowPaginationResponse.model_validate(
{
"items": workflows,
"items": [_WorkflowResponseSource(workflow, session=session) for workflow in workflows],
"page": page,
"limit": limit,
"has_more": has_more,
@ -1532,7 +1553,9 @@ class WorkflowByIdApi(Resource):
if not workflow:
raise NotFound("Workflow not found")
return dump_response(WorkflowResponse, workflow)
response = dump_response(WorkflowResponse, _WorkflowResponseSource(workflow, session=session))
return response
@setup_required
@login_required
@ -1626,10 +1649,10 @@ class DraftWorkflowTriggerRunApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
@get_app_model(mode=[AppMode.WORKFLOW])
@with_current_user
@edit_permission_required
@with_session
@get_app_model(mode=[AppMode.WORKFLOW])
def post(self, session: Session, current_user: Account, app_model: App):
"""
Poll for trigger events and execute full workflow when event arrives
@ -1637,7 +1660,7 @@ class DraftWorkflowTriggerRunApi(Resource):
args = DraftWorkflowTriggerRunPayload.model_validate(console_ns.payload or {})
node_id = args.node_id
workflow_service = WorkflowService()
draft_workflow = workflow_service.get_draft_workflow(app_model, session=db.session())
draft_workflow = workflow_service.get_draft_workflow(app_model, session=session)
if not draft_workflow:
raise ValueError("Workflow not found")
@ -1777,11 +1800,11 @@ class DraftWorkflowTriggerRunAllApi(Resource):
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW])
@with_current_user
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
@with_session
@get_app_model(mode=[AppMode.WORKFLOW])
def post(self, session: Session, current_user: Account, app_model: App):
"""
Full workflow debug when the start node is a trigger
@ -1790,7 +1813,7 @@ class DraftWorkflowTriggerRunAllApi(Resource):
args = DraftWorkflowTriggerRunAllPayload.model_validate(console_ns.payload or {})
node_ids = args.node_ids
workflow_service = WorkflowService()
draft_workflow = workflow_service.get_draft_workflow(app_model, session=db.session())
draft_workflow = workflow_service.get_draft_workflow(app_model, session=session)
if not draft_workflow:
raise ValueError("Workflow not found")

20
api/controllers/console/app/wraps.py
View File

@ -1,9 +1,8 @@
"""Controller decorators for console app resources.
App-loading decorators prefer a session injected by
`controllers.common.session.with_session` when present, while still supporting
existing handlers that have not been migrated yet and still rely on
Flask-SQLAlchemy's scoped `db.session`.
`get_app_model` still supports legacy handlers backed by Flask-SQLAlchemy's
scoped session. Trial app handlers compose `get_app_model_with_trial` under
`controllers.common.session.with_session` and always reuse that request session.
"""
from collections.abc import Callable
@ -41,9 +40,9 @@ def _load_app_model_from_scoped_session(app_id: str) -> App | None:
return app_model
def _load_app_model_with_trial(app_id: str) -> App | None:
def _load_app_model_with_trial(session: Session, app_id: str) -> App | None:
"""Load a normal app through its trial registration without applying current-tenant scope."""
app_model = db.session.scalar(
app_model = session.scalar(
select(App).join(TrialApp, TrialApp.app_id == App.id).where(App.id == app_id, App.status == "normal").limit(1)
)
return app_model
@ -157,7 +156,7 @@ def get_app_model_with_trial[**P, R](
*,
mode: AppMode | list[AppMode] | None = None,
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
"""Inject an app registered for trial or available from the recommended catalog."""
"""Inject a trial-registered or recommended App using the Session supplied by `with_session`."""
def decorator(view_func: Callable[P, R]) -> Callable[P, R]:
@wraps(view_func)
@ -170,9 +169,12 @@ def get_app_model_with_trial[**P, R](
del kwargs["app_id"]
app_model = _load_app_model_with_trial(app_id)
session = _get_injected_session(args)
if session is None:
raise RuntimeError("get_app_model_with_trial requires @with_session")
app_model = _load_app_model_with_trial(session, app_id)
if app_model is None:
app_model = RecommendedAppService.get_app(app_id, session=db.session())
app_model = RecommendedAppService.get_app(app_id, session=session)
if not app_model:
raise AppNotFoundError()

2
api/controllers/console/auth/forgot_password.py
View File

@ -203,5 +203,5 @@ class ForgotPasswordResetApi(Resource):
):
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
TenantService.create_tenant_member(tenant, account, db.session(), role="owner")
account.current_tenant = tenant
account.set_current_tenant_with_session(tenant, session=db.session())
tenant_was_created.send(tenant)

9
api/controllers/console/auth/login.py
View File

@ -4,6 +4,7 @@ import flask_login
from flask import make_response, request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy.orm import Session
from werkzeug.exceptions import Unauthorized
import services
@ -16,6 +17,7 @@ from controllers.common.fields import (
SimpleResultResponse,
)
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.auth.error import (
AuthenticationFailedError,
@ -317,7 +319,7 @@ class EmailCodeLoginApi(Resource):
else:
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner")
account.current_tenant = new_tenant
account.set_current_tenant_with_session(new_tenant, session=db.session())
tenant_was_created.send(new_tenant)
if account is None:
@ -356,7 +358,8 @@ class EmailCodeLoginApi(Resource):
class RefreshTokenApi(Resource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@console_ns.response(401, "Unauthorized", console_ns.models[SimpleResultMessageResponse.__name__])
def post(self):
@with_session(write=False)
def post(self, session: Session):
# Get refresh token from cookie instead of request body
refresh_token = extract_refresh_token(request)
@ -366,7 +369,7 @@ class RefreshTokenApi(Resource):
), 401
try:
new_token_pair = AccountService.refresh_token(refresh_token, session=db.session())
new_token_pair = AccountService.refresh_token(refresh_token, session=session)
except Unauthorized as exc:
return SimpleResultMessageResponse(result="fail", message=exc.description or "Unauthorized.").model_dump(
mode="json"

2
api/controllers/console/auth/oauth.py
View File

@ -284,7 +284,7 @@ def _generate_account(
else:
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner")
account.current_tenant = new_tenant
account.set_current_tenant_with_session(new_tenant, session=db.session())
tenant_was_created.send(new_tenant)
if not account:

9
api/controllers/console/auth/oauth_server.py
View File

@ -10,7 +10,7 @@ from werkzeug.exceptions import BadRequest, NotFound
from controllers.common.schema import register_response_schema_models, register_schema_models
from controllers.console.wraps import account_initialization_required, setup_required, with_current_user
from extensions.ext_database import db
from core.db.session_factory import session_factory
from graphon.model_runtime.utils.encoders import jsonable_encoder
from libs.login import login_required
from models import Account
@ -132,9 +132,10 @@ def oauth_server_access_token_required[T, **P, R](
response.headers["WWW-Authenticate"] = "Bearer"
return response
account = OAuthServerService.validate_oauth_access_token(
oauth_provider_app.client_id, access_token, db.session()
)
with session_factory.create_session() as session:
account = OAuthServerService.validate_oauth_access_token(
oauth_provider_app.client_id, access_token, session
)
if not account:
response = jsonify({"error": "access_token or client_id is invalid"})
response.status_code = 401

56
api/controllers/console/datasets/data_source.py
View File

@ -8,11 +8,12 @@ from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_serializer
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.fields import SimpleResultResponse, TextContentResponse
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
@ -188,16 +189,16 @@ class DataSourceApi(Resource):
@account_initialization_required
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@with_current_tenant_id
@with_session
def patch(
self, current_tenant_id: str, binding_id: UUID, action: Literal["enable", "disable"]
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)
with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
data_source_binding = session.execute(
select(DataSourceOauthBinding).where(
DataSourceOauthBinding.id == binding_id_str, DataSourceOauthBinding.tenant_id == current_tenant_id
)
).scalar_one_or_none()
data_source_binding = session.scalar(
select(DataSourceOauthBinding).where(
DataSourceOauthBinding.id == binding_id_str, DataSourceOauthBinding.tenant_id == current_tenant_id
)
)
if data_source_binding is None:
raise NotFound("Data source binding not found.")
# enable binding
@ -206,8 +207,6 @@ class DataSourceApi(Resource):
if data_source_binding.disabled:
data_source_binding.disabled = False
data_source_binding.updated_at = naive_utc_now()
db.session.add(data_source_binding)
db.session.commit()
else:
raise ValueError("Data source is not disabled.")
# disable binding
@ -215,8 +214,6 @@ class DataSourceApi(Resource):
if not data_source_binding.disabled:
data_source_binding.disabled = True
data_source_binding.updated_at = naive_utc_now()
db.session.add(data_source_binding)
db.session.commit()
else:
raise ValueError("Data source is disabled.")
return {"result": "success"}, 200
@ -231,7 +228,8 @@ class DataSourceNotionListApi(Resource):
@console_ns.response(200, "Success", console_ns.models[NotionIntegrateInfoListResponse.__name__])
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account) -> tuple[dict[str, Any], int]:
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account) -> tuple[dict[str, Any], int]:
query = DataSourceNotionListQuery.model_validate(request.args.to_dict(flat=True))
datasource_provider_service = DatasourceProviderService()
credential = datasource_provider_service.get_datasource_credentials(
@ -245,13 +243,13 @@ class DataSourceNotionListApi(Resource):
exist_page_ids = []
# import notion in the exist dataset
if query.dataset_id:
dataset = DatasetService.get_dataset(query.dataset_id, db.session())
dataset = DatasetService.get_dataset(query.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 = db.session.scalars(
documents = session.scalars(
select(Document).where(
Document.dataset_id == query.dataset_id,
Document.tenant_id == current_tenant_id,
@ -355,7 +353,8 @@ class DataSourceNotionIndexingEstimateApi(Resource):
@console_ns.expect(console_ns.models[NotionEstimatePayload.__name__])
@console_ns.response(200, "Success", console_ns.models[IndexingEstimate.__name__])
@with_current_tenant_id
def post(self, current_tenant_id: str) -> tuple[dict[str, Any], int]:
@with_session
def post(self, session: Session, current_tenant_id: str) -> tuple[dict[str, Any], int]:
payload = NotionEstimatePayload.model_validate(console_ns.payload or {})
args = payload.model_dump()
# validate args
@ -382,11 +381,12 @@ class DataSourceNotionIndexingEstimateApi(Resource):
extract_settings.append(extract_setting)
indexing_runner = IndexingRunner()
response = indexing_runner.indexing_estimate(
current_tenant_id,
extract_settings,
args["process_rule"],
args["doc_form"],
args["doc_language"],
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,
)
return dump_response(IndexingEstimate, response), 200
@ -398,13 +398,14 @@ class DataSourceNotionDatasetSyncApi(Resource):
@account_initialization_required
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, dataset_id: UUID) -> tuple[dict[str, str], int]:
@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, db.session())
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, db.session())
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
@ -417,14 +418,15 @@ class DataSourceNotionDocumentSyncApi(Resource):
@account_initialization_required
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, dataset_id: UUID, document_id: UUID) -> tuple[dict[str, str], int]:
@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, db.session())
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=db.session())
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)

204
api/controllers/console/datasets/datasets.py
View File

@ -1,3 +1,4 @@
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from uuid import UUID
@ -13,10 +14,10 @@ import services
from configs import dify_config
from controllers.common.fields import ApiBaseUrlResponse, SimpleResultResponse, UsageCheckResponse
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.apikey import ApiKeyItem, ApiKeyList
from controllers.console.app.error import ProviderNotInitializeError
from controllers.console.app.wraps import with_session
from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError
from controllers.console.wraps import (
RBACPermission,
@ -39,18 +40,18 @@ 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 IndexTechniqueType
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from extensions.ext_database import db
from fields.base import ResponseModel
from fields.dataset_fields import DatasetDetailResponse
from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from graphon.model_runtime.entities.model_entities import ModelType
from libs.helper import build_icon_url, dump_response, to_timestamp
from libs.login import login_required
from libs.url_utils import normalize_api_base_url
from models import Account, ApiToken, Dataset, Document, DocumentSegment, UploadFile
from models.dataset import DatasetPermission, DatasetPermissionEnum
from models import Account, ApiToken, App, Dataset, Document, DocumentSegment, UploadFile
from models.dataset import DatasetPermission, DatasetPermissionEnum, DatasetQuery
from models.enums import ApiTokenType, SegmentStatus
from models.provider_ids import ModelProviderID
from services.api_token_service import ApiTokenCache
from services.app_service import AppService
from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
from services.enterprise import rbac_service as enterprise_rbac_service
from services.enterprise.rbac_service import RBACResourceWhitelistScope, ReplaceMemberBindings
@ -205,6 +206,21 @@ class DatasetQueryDetailResponse(ResponseModel):
return to_timestamp(value)
@dataclass(frozen=True)
class _DatasetQueryResponseSource:
"""Expose query content through the request's database session."""
query: DatasetQuery
session: Session
@property
def queries(self) -> list[dict[str, Any]]:
return self.query.get_queries(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.query, name) # noqa: no-new-getattr response adapter delegates model fields
class DatasetQueryListResponse(ResponseModel):
data: list[DatasetQueryDetailResponse]
has_more: bool
@ -229,6 +245,21 @@ class RelatedAppResponse(ResponseModel):
return self
@dataclass(frozen=True)
class _RelatedAppResponseSource:
"""Expose the compatible app mode through the request's database session."""
app: App
session: Session
@property
def mode_compatible_with_agent(self) -> str:
return self.app.mode_compatible_with_agent_with_session(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.app, name) # noqa: no-new-getattr response adapter delegates model fields
class RelatedAppListResponse(ResponseModel):
data: list[RelatedAppResponse]
total: int
@ -397,7 +428,8 @@ class DatasetListApi(Resource):
@enterprise_license_required
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account):
# Convert query parameters to dict, handling list parameters correctly
query_params: dict[str, str | list[str]] = dict(request.args.to_dict())
# Handle ids and tag_ids as lists (Flask request.args.getlist returns list even for single value)
@ -410,7 +442,7 @@ class DatasetListApi(Resource):
permissions = enterprise_rbac_service.RBACService.MyPermissions.get(
str(current_tenant_id),
current_user.id,
session=db.session(),
session=session,
)
accessible_dataset_ids: list[str] | None = None
@ -449,12 +481,13 @@ class DatasetListApi(Resource):
user=current_user,
accessible_dataset_ids=accessible_dataset_ids,
include_own_datasets=include_own_datasets,
session=session,
)
else:
datasets, total = DatasetService.get_datasets(
query.page,
query.limit,
db.session(),
session,
current_tenant_id,
current_user,
query.keyword,
@ -479,11 +512,14 @@ class DatasetListApi(Resource):
for embedding_model in embedding_models:
model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
data = [dump_response(DatasetDetailResponse, dataset) for dataset in datasets]
data = [
dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session))
for dataset in datasets
]
dataset_ids = [item["id"] for item in data if item.get("permission") == "partial_members"]
partial_members_map: dict[str, list[str]] = {}
if dataset_ids:
partial_member_rows = db.session.execute(
partial_member_rows = session.execute(
select(DatasetPermission.dataset_id, DatasetPermission.account_id).where(
DatasetPermission.dataset_id.in_(dataset_ids)
)
@ -579,9 +615,9 @@ class DatasetListApi(Resource):
session=session,
)
item = DatasetDetailWithPartialMembersResponse.model_validate(dataset, from_attributes=True).model_dump(
mode="json"
)
item = DatasetDetailWithPartialMembersResponse.model_validate(
dataset_detail_response_source(dataset, session=session), from_attributes=True
).model_dump(mode="json")
item["permission_keys"] = permission_keys_map.get(dataset.id, [])
return item, 201
@ -604,30 +640,31 @@ class DatasetApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID):
@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(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
permissions = enterprise_rbac_service.RBACService.MyPermissions.get(
current_tenant_id,
current_user.id,
dataset_id=dataset_id_str,
session=db.session(),
session=session,
)
permission_keys_map = permissions.dataset.permission_keys_by_resource_ids([dataset_id_str])
data = dump_response(DatasetDetailResponse, dataset)
data = dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session))
data["permission_keys"] = permission_keys_map.get(dataset_id_str, [])
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
if dataset.embedding_model_provider:
provider_id = ModelProviderID(dataset.embedding_model_provider)
data["embedding_model_provider"] = str(provider_id)
if data.get("permission") == "partial_members":
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session())
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session)
data.update({"partial_member_list": part_users_list})
# check embedding setting
@ -671,7 +708,7 @@ class DatasetApi(Resource):
@with_session
def patch(self, 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, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
@ -704,19 +741,19 @@ class DatasetApi(Resource):
[dataset_id_str],
session=session,
)
result_data = dump_response(DatasetDetailResponse, dataset)
result_data = dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session))
result_data["permission_keys"] = permission_keys_map.get(dataset_id_str, [])
tenant_id = current_tenant_id
if payload.partial_member_list is not None and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM:
DatasetPermissionService.update_partial_member_list(
tenant_id, dataset_id_str, payload.partial_member_list, db.session()
tenant_id, dataset_id_str, payload.partial_member_list, session
)
# clear partial member list when permission is only_me or all_team_members
elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}:
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session())
DatasetPermissionService.clear_partial_member_list(dataset_id_str, session)
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session())
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session)
result_data.update({"partial_member_list": partial_member_list})
return dump_response(DatasetDetailWithPartialMembersResponse, result_data), 200
@ -728,15 +765,16 @@ class DatasetApi(Resource):
@console_ns.response(204, "Dataset deleted successfully")
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def delete(self, current_user: Account, dataset_id: UUID):
@with_session
def delete(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
if not (current_user.has_edit_permission or current_user.is_dataset_operator):
raise Forbidden()
try:
if DatasetService.delete_dataset(dataset_id_str, current_user, db.session()):
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session())
if DatasetService.delete_dataset(dataset_id_str, current_user, session):
DatasetPermissionService.clear_partial_member_list(dataset_id_str, session)
return "", 204
else:
raise NotFound("Dataset not found.")
@ -758,10 +796,11 @@ class DatasetUseCheckApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, db.session())
dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, session)
return UsageCheckResponse(is_using=dataset_is_using).model_dump(mode="json"), 200
@ -780,24 +819,27 @@ class DatasetQueryApi(Resource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_user: Account, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
page = request.args.get("page", default=1, type=int)
limit = request.args.get("limit", default=20, type=int)
dataset_queries, total = DatasetService.get_dataset_queries(dataset_id=dataset.id, page=page, per_page=limit)
dataset_queries, total = DatasetService.get_dataset_queries(
dataset_id=dataset.id, page=page, per_page=limit, session=session
)
response = {
"data": dataset_queries,
"data": [_DatasetQueryResponseSource(query=query, session=session) for query in dataset_queries],
"has_more": len(dataset_queries) == limit,
"limit": limit,
"total": total,
@ -820,7 +862,8 @@ class DatasetIndexingEstimateApi(Resource):
@account_initialization_required
@console_ns.expect(console_ns.models[IndexingEstimatePayload.__name__])
@with_current_tenant_id
def post(self, current_tenant_id: str):
@with_session
def post(self, session: Session, current_tenant_id: str):
payload = IndexingEstimatePayload.model_validate(console_ns.payload or {})
args = payload.model_dump()
# validate args
@ -829,7 +872,7 @@ class DatasetIndexingEstimateApi(Resource):
match args["info_list"]["data_source_type"]:
case "upload_file":
file_ids = args["info_list"]["file_info_list"]["file_ids"]
file_details = db.session.scalars(
file_details = session.scalars(
select(UploadFile).where(UploadFile.tenant_id == current_tenant_id, UploadFile.id.in_(file_ids))
).all()
if file_details is None:
@ -886,13 +929,14 @@ class DatasetIndexingEstimateApi(Resource):
indexing_runner = IndexingRunner()
try:
response = indexing_runner.indexing_estimate(
current_tenant_id,
extract_settings,
args["process_rule"],
args["doc_form"],
args["doc_language"],
args["dataset_id"],
args["indexing_technique"],
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"],
dataset_id=args["dataset_id"],
indexing_technique=args["indexing_technique"],
session=session,
)
except LLMBadRequestError:
raise ProviderNotInitializeError(
@ -931,24 +975,25 @@ class DatasetRelatedAppListApi(Resource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_user: Account, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
app_dataset_joins = DatasetService.get_related_apps(dataset.id, db.session())
app_dataset_joins = DatasetService.get_related_apps(dataset.id, session)
related_apps = []
for app_dataset_join in app_dataset_joins:
app_model = app_dataset_join.app
app_model = AppService.get_app_by_id(app_dataset_join.app_id, session)
if app_model:
related_apps.append(app_model)
related_apps.append(_RelatedAppResponseSource(app=app_model, session=session))
return dump_response(RelatedAppListResponse, {"data": related_apps, "total": len(related_apps)}), 200
@ -968,15 +1013,16 @@ class DatasetIndexingStatusApi(Resource):
@account_initialization_required
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_tenant_id: str, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, dataset_id: UUID):
dataset_id_str = str(dataset_id)
documents = db.session.scalars(
documents = session.scalars(
select(Document).where(Document.dataset_id == dataset_id_str, Document.tenant_id == current_tenant_id)
).all()
documents_status = []
for document in documents:
completed_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.completed_at.isnot(None),
DocumentSegment.document_id == str(document.id),
@ -986,7 +1032,7 @@ class DatasetIndexingStatusApi(Resource):
or 0
)
total_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.document_id == str(document.id),
DocumentSegment.status != SegmentStatus.RE_SEGMENT,
@ -1026,8 +1072,9 @@ class DatasetApiKeyApi(Resource):
@login_required
@account_initialization_required
@with_current_tenant_id
def get(self, current_tenant_id: str):
keys = db.session.scalars(
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str):
keys = session.scalars(
select(ApiToken).where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id)
).all()
return dump_response(ApiKeyList, {"data": keys})
@ -1040,9 +1087,10 @@ class DatasetApiKeyApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_API_KEY_MANAGE, resource_required=False)
@account_initialization_required
@with_current_tenant_id
def post(self, current_tenant_id: str):
@with_session
def post(self, session: Session, current_tenant_id: str):
current_key_count = (
db.session.scalar(
session.scalar(
select(func.count(ApiToken.id)).where(
ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id
)
@ -1057,13 +1105,13 @@ class DatasetApiKeyApi(Resource):
custom="max_keys_exceeded",
)
key = ApiToken.generate_api_key(self.token_prefix, 24)
key = ApiToken.generate_api_key(self.token_prefix, 24, session=session)
api_token = ApiToken()
api_token.tenant_id = current_tenant_id
api_token.token = key
api_token.type = self.resource_type
db.session.add(api_token)
db.session.commit()
session.add(api_token)
session.flush()
return dump_response(ApiKeyItem, api_token), 200
@ -1081,9 +1129,10 @@ class DatasetApiDeleteApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_API_KEY_MANAGE, resource_required=False)
@account_initialization_required
@with_current_tenant_id
def delete(self, current_tenant_id: str, api_key_id: UUID):
@with_session
def delete(self, session: Session, current_tenant_id: str, api_key_id: UUID):
api_key_id_str = str(api_key_id)
key = db.session.scalar(
key = session.scalar(
select(ApiToken)
.where(
ApiToken.tenant_id == current_tenant_id,
@ -1101,8 +1150,7 @@ class DatasetApiDeleteApi(Resource):
assert key is not None # nosec - for type checker only
ApiTokenCache.delete(key.token, key.type)
db.session.delete(key)
db.session.commit()
session.delete(key)
return "", 204
@ -1114,10 +1162,11 @@ class DatasetEnableApiApi(Resource):
@account_initialization_required
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, dataset_id: UUID, status: str):
@with_session
def post(self, session: Session, dataset_id: UUID, status: str):
dataset_id_str = str(dataset_id)
DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", db.session())
DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", session)
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@ -1184,12 +1233,13 @@ class DatasetErrorDocs(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str, db.session())
results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str, session)
return dump_response(ErrorDocsResponse, {"data": results, "total": len(results)}), 200
@ -1211,17 +1261,18 @@ class DatasetPermissionUserListApi(Resource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_user: Account, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session())
partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session)
return dump_response(PartialMemberListResponse, {"data": partial_members_list}), 200
@ -1241,10 +1292,11 @@ class DatasetAutoDisableLogApi(Resource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(dataset_id_str, db.session())
auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(dataset_id_str, session)
return dump_response(AutoDisableLogsResponse, auto_disable_logs), 200

333
api/controllers/console/datasets/datasets_document.py
View File

@ -1,7 +1,7 @@
import json
import logging
from argparse import ArgumentTypeError
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from contextlib import ExitStack
from datetime import datetime
from typing import Any, Literal, cast
@ -12,12 +12,14 @@ from flask import request, send_file
from flask_restx import Resource
from pydantic import BaseModel, Field, JsonValue, field_validator
from sqlalchemy import asc, desc, func, select
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, NotFound
import services
from controllers.common.controller_schemas import DocumentBatchDownloadZipPayload
from controllers.common.fields import SimpleResultMessageResponse, SimpleResultResponse, UrlResponse
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, RBACResourceScope, rbac_permission_required
from core.entities.knowledge_entities import IndexingEstimate
@ -34,13 +36,15 @@ from core.rag.entities import Rule
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 IndexTechniqueType
from extensions.ext_database import db
from fields.base import ResponseModel
from fields.document_fields import (
DocumentMetadataResponse,
DocumentResponse,
DocumentStatusListResponse,
DocumentStatusResponse,
DocumentWithSession,
document_response,
document_responses,
normalize_enum,
)
from graphon.model_runtime.entities.model_entities import ModelType
@ -49,7 +53,7 @@ from libs.datetime_utils import naive_utc_now
from libs.helper import dump_response, to_timestamp
from libs.login import login_required
from libs.pagination import paginate_query
from models import Account, DatasetProcessRule, Document, DocumentSegment, UploadFile
from models import Account, Document, DocumentSegment, UploadFile
from models.dataset import DocumentPipelineExecutionLog
from models.enums import IndexingStatus, ProcessRuleMode, SegmentStatus
from services.dataset_ref_service import DatasetRefService
@ -110,6 +114,24 @@ class DocumentWithSegmentsResponse(DocumentResponse):
total_segments: int | None = Field(default=None, exclude_if=lambda value: value is None)
class DocumentWithSegmentsSession(DocumentWithSession):
@property
def process_rule_dict(self) -> Any:
process_rule = self.document.get_dataset_process_rule(session=self.session)
return process_rule.to_dict() if process_rule else None
def document_with_segments_responses(
documents: Sequence[Document], *, session: Session
) -> list[DocumentWithSegmentsResponse]:
return [
DocumentWithSegmentsResponse.model_validate(
DocumentWithSegmentsSession(document=document, session=session), from_attributes=True
)
for document in documents
]
class DatasetAndDocumentResponse(ResponseModel):
dataset: DatasetResponse
documents: list[DocumentResponse]
@ -269,18 +291,18 @@ register_response_schema_models(
class DocumentResource(Resource):
def get_document(
self, dataset_id: str, document_id: str, current_user: Account, current_tenant_id: str
self, session: Session, dataset_id: str, document_id: str, current_user: Account, current_tenant_id: str
) -> Document:
dataset = DatasetService.get_dataset(dataset_id, db.session())
dataset = DatasetService.get_dataset(dataset_id, session)
if not dataset:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
document = DocumentService.get_document(dataset_id, document_id, session=db.session())
document = DocumentService.get_document(dataset_id, document_id, session=session)
if not document:
raise NotFound("Document not found.")
@ -290,17 +312,19 @@ class DocumentResource(Resource):
return document
def get_batch_documents(self, dataset_id: str, batch: str, current_user: Account) -> Sequence[Document]:
dataset = DatasetService.get_dataset(dataset_id, db.session())
def get_batch_documents(
self, session: Session, dataset_id: str, batch: str, current_user: Account
) -> Sequence[Document]:
dataset = DatasetService.get_dataset(dataset_id, session)
if not dataset:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
documents = DocumentService.get_batch_documents(dataset_id, batch, db.session())
documents = DocumentService.get_batch_documents(dataset_id, batch, session)
if not documents:
raise NotFound("Documents not found.")
@ -318,7 +342,8 @@ class GetProcessRuleApi(Resource):
@login_required
@account_initialization_required
@with_current_user
def get(self, current_user: Account):
@with_session(write=False)
def get(self, session: Session, current_user: Account):
req_data = request.args
document_id = req_data.get("document_id")
@ -328,26 +353,21 @@ class GetProcessRuleApi(Resource):
rules = DocumentService.DEFAULT_RULES["rules"]
limits = DocumentService.DEFAULT_RULES["limits"]
if document_id:
# get the latest process rule
document = db.get_or_404(Document, document_id)
document = DocumentService.get_document_by_id(document_id, session)
if document is None:
raise NotFound("Document not found.")
dataset = DatasetService.get_dataset(document.dataset_id, db.session())
dataset = DatasetService.get_dataset(document.dataset_id, session)
if not dataset:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
# get the latest process rule
dataset_process_rule = db.session.scalar(
select(DatasetProcessRule)
.where(DatasetProcessRule.dataset_id == document.dataset_id)
.order_by(DatasetProcessRule.created_at.desc())
.limit(1)
)
dataset_process_rule = dataset.get_latest_process_rule(session=session)
if dataset_process_rule:
mode = dataset_process_rule.mode
rules = dataset_process_rule.rules_dict
@ -381,7 +401,8 @@ class DatasetDocumentListApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID):
@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)
raw_args = request.args.to_dict()
param = DocumentDatasetListParam.model_validate(raw_args)
@ -407,12 +428,12 @@ class DatasetDocumentListApi(Resource):
)
except (ArgumentTypeError, ValueError, Exception):
fetch = False
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
@ -457,20 +478,20 @@ class DatasetDocumentListApi(Resource):
desc(Document.position),
)
paginated_documents = paginate_query(query, page=page, per_page=limit, max_per_page=100)
paginated_documents = paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100)
documents = paginated_documents.items
DocumentService.enrich_documents_with_summary_index_status(
documents=documents,
dataset=dataset,
tenant_id=current_tenant_id,
session=db.session(),
session=session,
)
if fetch:
for document in documents:
completed_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.completed_at.isnot(None),
DocumentSegment.document_id == str(document.id),
@ -480,7 +501,7 @@ class DatasetDocumentListApi(Resource):
or 0
)
total_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.document_id == str(document.id),
DocumentSegment.status != SegmentStatus.RE_SEGMENT,
@ -491,7 +512,7 @@ class DatasetDocumentListApi(Resource):
document.completed_segments = completed_segments
document.total_segments = total_segments
response = {
"data": documents,
"data": document_with_segments_responses(documents, session=session),
"has_more": len(documents) == limit,
"limit": limit,
"total": paginated_documents.total,
@ -509,10 +530,11 @@ class DatasetDocumentListApi(Resource):
@console_ns.response(200, "Documents created successfully", console_ns.models[DatasetAndDocumentResponse.__name__])
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID):
@with_session
def post(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
@ -522,7 +544,7 @@ class DatasetDocumentListApi(Resource):
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
@ -536,9 +558,9 @@ class DatasetDocumentListApi(Resource):
try:
documents, batch = DocumentService.save_document_with_dataset_id(
dataset, knowledge_config, current_user, session=db.session()
dataset, knowledge_config, current_user, session=session
)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@ -547,7 +569,10 @@ class DatasetDocumentListApi(Resource):
except ModelCurrentlyNotSupportError:
raise ProviderModelCurrentlyNotSupportError()
return dump_response(DatasetAndDocumentResponse, {"dataset": dataset, "documents": documents, "batch": batch})
return dump_response(
DatasetAndDocumentResponse,
{"dataset": dataset, "documents": document_responses(documents, session=session), "batch": batch},
)
@setup_required
@login_required
@ -555,9 +580,10 @@ class DatasetDocumentListApi(Resource):
@cloud_edition_billing_rate_limit_check("knowledge")
@console_ns.response(204, "Documents deleted successfully")
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def delete(self, dataset_id: UUID):
@with_session
def delete(self, session: Session, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
# check user's model setting
@ -566,7 +592,7 @@ class DatasetDocumentListApi(Resource):
try:
document_ids = request.args.getlist("document_id")
dataset_ref = DatasetRefService.create_dataset_ref(dataset)
DocumentService.delete_documents(dataset_ref, document_ids, dataset.doc_form, db.session())
DocumentService.delete_documents(dataset_ref, document_ids, dataset.get_doc_form(session=session), session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
@ -589,7 +615,8 @@ class DatasetInitApi(Resource):
@cloud_edition_billing_rate_limit_check("knowledge")
@with_current_user
@with_current_tenant_id
def post(self, current_tenant_id: str, current_user: Account):
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account):
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
raise Forbidden()
@ -625,7 +652,7 @@ class DatasetInitApi(Resource):
tenant_id=current_tenant_id,
knowledge_config=knowledge_config,
account=current_user,
session=db.session(),
session=session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@ -634,7 +661,10 @@ class DatasetInitApi(Resource):
except ModelCurrentlyNotSupportError:
raise ProviderModelCurrentlyNotSupportError()
return dump_response(DatasetAndDocumentResponse, {"dataset": dataset, "documents": documents, "batch": batch})
return dump_response(
DatasetAndDocumentResponse,
{"dataset": dataset, "documents": document_responses(documents, session=session), "batch": batch},
)
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-estimate")
@ -655,23 +685,24 @@ class DocumentIndexingEstimateApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
if document.indexing_status in {IndexingStatus.COMPLETED, IndexingStatus.ERROR}:
raise DocumentAlreadyFinishedError()
data_process_rule = document.dataset_process_rule
data_process_rule_dict = data_process_rule.to_dict() if data_process_rule else {}
data_process_rule = document.get_dataset_process_rule(session=session)
data_process_rule_dict: Mapping[str, Any] = data_process_rule.to_dict() if data_process_rule else {}
if document.data_source_type == "upload_file":
data_source_info = document.data_source_info_dict
if data_source_info and "upload_file_id" in data_source_info:
file_id = data_source_info["upload_file_id"]
file = db.session.scalar(
file = session.scalar(
select(UploadFile)
.where(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id)
.limit(1)
@ -689,12 +720,13 @@ class DocumentIndexingEstimateApi(DocumentResource):
try:
estimate_response = indexing_runner.indexing_estimate(
current_tenant_id,
[extract_setting],
data_process_rule_dict,
document.doc_form,
"English",
dataset_id_str,
tenant_id=current_tenant_id,
extract_settings=[extract_setting],
tmp_processing_rule=data_process_rule_dict,
doc_form=document.doc_form,
doc_language="English",
dataset_id=dataset_id_str,
session=session,
)
return (
# TODO: why using zero here? the same for the below endpoint
@ -745,9 +777,10 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, batch: str):
@with_session
def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, batch: str):
dataset_id_str = str(dataset_id)
documents = self.get_batch_documents(dataset_id_str, batch, current_user)
documents = self.get_batch_documents(session, dataset_id_str, batch, current_user)
if not documents:
return (
IndexingEstimateResponse(
@ -759,8 +792,8 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
).model_dump(mode="json", exclude_none=True),
200,
)
data_process_rule = documents[0].dataset_process_rule
data_process_rule_dict = data_process_rule.to_dict() if data_process_rule else {}
data_process_rule = documents[0].get_dataset_process_rule(session=session)
data_process_rule_dict: Mapping[str, Any] = data_process_rule.to_dict() if data_process_rule else {}
extract_settings = []
for document in documents:
if document.indexing_status in {IndexingStatus.COMPLETED, IndexingStatus.ERROR}:
@ -771,7 +804,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
if not data_source_info:
continue
file_id = data_source_info["upload_file_id"]
file_detail = db.session.scalar(
file_detail = session.scalar(
select(UploadFile)
.where(UploadFile.tenant_id == current_tenant_id, UploadFile.id == file_id)
.limit(1)
@ -825,12 +858,13 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
indexing_runner = IndexingRunner()
try:
response = indexing_runner.indexing_estimate(
current_tenant_id,
extract_settings,
data_process_rule_dict,
document.doc_form,
"English",
dataset_id_str,
tenant_id=current_tenant_id,
extract_settings=extract_settings,
tmp_processing_rule=data_process_rule_dict,
doc_form=document.doc_form,
doc_language="English",
dataset_id=dataset_id_str,
session=session,
)
return (
IndexingEstimateResponse(
@ -865,13 +899,14 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_user: Account, dataset_id: UUID, batch: str):
@with_session(write=False)
def get(self, session: Session, current_user: Account, dataset_id: UUID, batch: str):
dataset_id_str = str(dataset_id)
documents = self.get_batch_documents(dataset_id_str, batch, current_user)
documents = self.get_batch_documents(session, dataset_id_str, batch, current_user)
documents_status = []
for document in documents:
completed_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.completed_at.isnot(None),
DocumentSegment.document_id == str(document.id),
@ -881,7 +916,7 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
or 0
)
total_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.document_id == str(document.id),
DocumentSegment.status != SegmentStatus.RE_SEGMENT,
@ -923,13 +958,14 @@ class DocumentIndexingStatusApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
completed_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.completed_at.isnot(None),
DocumentSegment.document_id == document_id_str,
@ -939,7 +975,7 @@ class DocumentIndexingStatusApi(DocumentResource):
or 0
)
total_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.document_id == document_id_str,
DocumentSegment.status != SegmentStatus.RE_SEGMENT,
@ -987,10 +1023,11 @@ class DocumentApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
metadata = request.args.get("metadata", "all")
if metadata not in self.METADATA_CHOICES:
@ -1002,20 +1039,22 @@ class DocumentApi(DocumentResource):
{
"id": document.id,
"doc_type": document.doc_type,
"doc_metadata": document.doc_metadata_details,
"doc_metadata": document.get_doc_metadata_details(session=session),
}
)
return response.model_dump(mode="json", include={"id", *metadata_fields}, exclude_unset=True), 200
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session())
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session)
document_process_rule = document.get_dataset_process_rule(session=session)
document_process_rules: Mapping[str, Any] = document_process_rule.to_dict() if document_process_rule else {}
segment_count = document.get_segment_count(session=session)
response = DocumentDetailResponse.model_validate(
{
"id": document.id,
"position": document.position,
"data_source_type": document.data_source_type,
"data_source_info": document.data_source_info_dict,
"data_source_detail_dict": document.data_source_detail_dict,
"data_source_detail_dict": document.get_data_source_detail_dict(session=session),
"dataset_process_rule_id": document.dataset_process_rule_id,
"dataset_process_rule": dataset_process_rules,
"document_process_rule": document_process_rules,
@ -1034,10 +1073,10 @@ class DocumentApi(DocumentResource):
"disabled_by": document.disabled_by,
"archived": document.archived,
"doc_type": document.doc_type,
"doc_metadata": document.doc_metadata_details,
"segment_count": document.segment_count,
"average_segment_length": document.average_segment_length,
"hit_count": document.hit_count,
"doc_metadata": document.get_doc_metadata_details(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),
"display_status": document.display_status,
"doc_form": document.doc_form,
"doc_language": document.doc_language,
@ -1055,19 +1094,22 @@ class DocumentApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def delete(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def delete(
self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID
):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
try:
DocumentService.delete_document(document, db.session())
DocumentService.delete_document(document, session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
@ -1088,12 +1130,13 @@ class DocumentDownloadApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_DOCUMENT_DOWNLOAD)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID) -> dict[str, Any]:
@with_session(write=False)
def get(
self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID
) -> dict[str, Any]:
# Reuse the shared permission/tenant checks implemented in DocumentResource.
document = self.get_document(str(dataset_id), str(document_id), current_user, current_tenant_id)
return UrlResponse(url=DocumentService.get_document_download_url(document, db.session())).model_dump(
mode="json"
)
document = self.get_document(session, str(dataset_id), str(document_id), current_user, current_tenant_id)
return UrlResponse(url=DocumentService.get_document_download_url(document, session)).model_dump(mode="json")
@console_ns.route("/datasets/<uuid:dataset_id>/documents/download-zip")
@ -1111,7 +1154,8 @@ class DocumentBatchDownloadZipApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID):
@with_session(write=False)
def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
"""Stream a ZIP archive containing the requested uploaded documents."""
# Parse and validate request payload.
payload = DocumentBatchDownloadZipPayload.model_validate(console_ns.payload or {})
@ -1123,7 +1167,7 @@ class DocumentBatchDownloadZipApi(DocumentResource):
document_ids=document_ids,
tenant_id=current_tenant_id,
current_user=current_user,
session=db.session(),
session=session,
)
# Delegate ZIP packing to FileService, but keep Flask response+cleanup in the route.
@ -1162,8 +1206,10 @@ class DocumentProcessingApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def patch(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
@ -1172,7 +1218,7 @@ class DocumentProcessingApi(DocumentResource):
):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
@ -1186,7 +1232,6 @@ class DocumentProcessingApi(DocumentResource):
document.paused_by = current_user.id
document.paused_at = naive_utc_now()
document.is_paused = True
db.session.commit()
case "resume":
if document.indexing_status not in {IndexingStatus.PAUSED, IndexingStatus.ERROR}:
@ -1195,7 +1240,6 @@ class DocumentProcessingApi(DocumentResource):
document.paused_by = None
document.paused_at = None
document.is_paused = False
db.session.commit()
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@ -1219,10 +1263,11 @@ class DocumentMetadataApi(DocumentResource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def put(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def put(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id)
req_data = DocumentMetadataUpdatePayload.model_validate(request.get_json() or {})
@ -1254,7 +1299,6 @@ class DocumentMetadataApi(DocumentResource):
document.doc_type = doc_type
document.updated_at = naive_utc_now()
db.session.commit()
return SimpleResultMessageResponse(result="success", message="Document metadata updated.").model_dump(
mode="json"
@ -1271,11 +1315,16 @@ class DocumentStatusApi(DocumentResource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def patch(
self, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable", "archive", "un_archive"]
self,
session: Session,
current_user: Account,
dataset_id: UUID,
action: Literal["enable", "disable", "archive", "un_archive"],
):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
@ -1287,12 +1336,12 @@ class DocumentStatusApi(DocumentResource):
DatasetService.check_dataset_model_setting(dataset)
# check user's permission
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
document_ids = request.args.getlist("document_id")
try:
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session())
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, session)
except services.errors.document.DocumentIndexingError as e:
raise InvalidActionError(str(e))
except ValueError as e:
@ -1311,16 +1360,17 @@ class DocumentPauseApi(DocumentResource):
@cloud_edition_billing_rate_limit_check("knowledge")
@console_ns.response(204, "Document paused successfully")
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def patch(self, dataset_id: UUID, document_id: UUID):
@with_session
def patch(self, session: Session, dataset_id: UUID, document_id: UUID):
"""pause document."""
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
# 404 if document not found
if document is None:
@ -1332,7 +1382,7 @@ class DocumentPauseApi(DocumentResource):
try:
# pause document
DocumentService.pause_document(document, db.session())
DocumentService.pause_document(document, session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot pause completed document.")
@ -1347,14 +1397,15 @@ class DocumentRecoverApi(DocumentResource):
@cloud_edition_billing_rate_limit_check("knowledge")
@console_ns.response(204, "Document resumed successfully")
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def patch(self, dataset_id: UUID, document_id: UUID):
@with_session
def patch(self, session: Session, dataset_id: UUID, document_id: UUID):
"""recover document."""
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
# 404 if document not found
if document is None:
@ -1365,7 +1416,7 @@ class DocumentRecoverApi(DocumentResource):
raise ArchivedDocumentImmutableError()
try:
# pause document
DocumentService.recover_document(document, db.session())
DocumentService.recover_document(document, session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Document is not in paused status.")
@ -1381,17 +1432,18 @@ class DocumentRetryApi(DocumentResource):
@console_ns.expect(console_ns.models[DocumentRetryPayload.__name__])
@console_ns.response(204, "Documents retry started successfully")
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, dataset_id: UUID):
@with_session
def post(self, session: Session, dataset_id: UUID):
"""retry document."""
payload = DocumentRetryPayload.model_validate(console_ns.payload or {})
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
retry_documents = []
if not dataset:
raise NotFound("Dataset not found.")
for document_id in payload.document_ids:
try:
document = DocumentService.get_document(dataset.id, document_id, session=db.session())
document = DocumentService.get_document(dataset.id, document_id, session=session)
# 404 if document not found
if document is None:
@ -1409,7 +1461,7 @@ class DocumentRetryApi(DocumentResource):
logger.exception("Failed to retry document, document id: %s", document_id)
continue
# retry document
DocumentService.retry_document(dataset_id_str, retry_documents, db.session())
DocumentService.retry_document(dataset_id_str, retry_documents, session)
return "", 204
@ -1423,22 +1475,23 @@ class DocumentRenameApi(DocumentResource):
@console_ns.expect(console_ns.models[DocumentRenamePayload.__name__])
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def post(self, session: Session, current_user: Account, dataset_id: UUID, document_id: UUID):
# The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
if not current_user.is_dataset_editor:
raise Forbidden()
dataset = DatasetService.get_dataset(str(dataset_id), db.session())
dataset = DatasetService.get_dataset(dataset_id, session)
if not dataset:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_operator_permission(current_user, dataset, session=db.session())
DatasetService.check_dataset_operator_permission(current_user, dataset, session=session)
payload = DocumentRenamePayload.model_validate(console_ns.payload or {})
try:
document = DocumentService.rename_document(str(dataset_id), str(document_id), payload.name, db.session())
document = DocumentService.rename_document(str(dataset_id), str(document_id), payload.name, session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
return dump_response(DocumentResponse, document)
return dump_response(DocumentResponse, document_response(document, session=session))
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/website-sync")
@ -1449,14 +1502,15 @@ class WebsiteDocumentSyncApi(DocumentResource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def get(self, session: Session, current_tenant_id: str, dataset_id: UUID, document_id: UUID):
"""sync website document."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
if document.tenant_id != current_tenant_id:
@ -1467,7 +1521,7 @@ class WebsiteDocumentSyncApi(DocumentResource):
if DocumentService.check_archived(document):
raise ArchivedDocumentImmutableError()
# sync document
DocumentService.sync_website_document(dataset_id_str, document, db.session())
DocumentService.sync_website_document(dataset_id_str, document, session)
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@ -1483,17 +1537,18 @@ class DocumentPipelineExecutionLogApi(DocumentResource):
@login_required
@account_initialization_required
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, dataset_id: UUID, document_id: UUID):
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
log = db.session.scalar(
log = session.scalar(
select(DocumentPipelineExecutionLog)
.where(DocumentPipelineExecutionLog.document_id == document_id_str)
.order_by(DocumentPipelineExecutionLog.created_at.desc())
@ -1532,7 +1587,8 @@ class DocumentGenerateSummaryApi(Resource):
@cloud_edition_billing_rate_limit_check("knowledge")
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID):
@with_session
def post(self, session: Session, current_user: Account, dataset_id: UUID):
"""
Generate summary index for specified documents.
@ -1543,7 +1599,7 @@ class DocumentGenerateSummaryApi(Resource):
dataset_id_str = str(dataset_id)
# Get dataset
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
@ -1552,7 +1608,7 @@ class DocumentGenerateSummaryApi(Resource):
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
@ -1577,7 +1633,7 @@ class DocumentGenerateSummaryApi(Resource):
raise ValueError("Summary index is not enabled for this dataset. Please enable it in the dataset settings.")
# Verify all documents exist and belong to the dataset
documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list, db.session())
documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list, session)
if len(documents) != len(document_list):
found_ids = {doc.id for doc in documents}
@ -1593,7 +1649,7 @@ class DocumentGenerateSummaryApi(Resource):
DocumentService.update_documents_need_summary(
dataset_id=dataset_id_str,
document_ids=document_ids_to_update,
session=db.session(),
session=session,
need_summary=True,
)
@ -1631,7 +1687,8 @@ class DocumentSummaryStatusApi(DocumentResource):
@account_initialization_required
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_user: Account, dataset_id: UUID, document_id: UUID):
"""
Get summary index generation status for a document.
@ -1649,13 +1706,13 @@ class DocumentSummaryStatusApi(DocumentResource):
document_id_str = str(document_id)
# Get dataset
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# Check permissions
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
@ -1665,7 +1722,7 @@ class DocumentSummaryStatusApi(DocumentResource):
result = SummaryIndexService.get_document_summary_status_detail(
document_id=document_id_str,
dataset_id=dataset_id_str,
session=db.session(),
session=session,
)
return dump_response(DocumentSummaryStatusResponse, result), 200

213
api/controllers/console/datasets/datasets_segments.py
View File

@ -8,6 +8,7 @@ from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy import String, case, cast, func, literal, or_, select
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, NotFound
import services
@ -20,6 +21,7 @@ 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.app.error import ProviderNotInitializeError
from controllers.console.datasets.error import (
@ -42,7 +44,6 @@ from controllers.console.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_database import db
from extensions.ext_redis import redis_client
from fields.base import ResponseModel
from fields.segment_fields import (
@ -165,7 +166,7 @@ register_response_schema_models(
def _get_segment_for_document(
dataset: Dataset, document: Document, segment_id: str
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)
@ -173,7 +174,7 @@ def _get_segment_for_document(
raise NotFound("Document not found.")
segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id)
segment = SegmentService.get_segment_by_ref(segment_ref, db.session())
segment = SegmentService.get_segment_by_ref(segment_ref, session=session)
if not segment:
raise NotFound("Segment not found.")
return segment_ref, segment
@ -190,19 +191,20 @@ class DatasetDocumentSegmentListApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
@ -271,19 +273,19 @@ class DatasetDocumentSegmentListApi(Resource):
elif args.enabled.lower() == "false":
query = query.where(DocumentSegment.enabled == False)
segments = paginate_query(query, page=page, per_page=limit, max_per_page=100)
segments = paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100)
segment_list = list(segments.items)
segment_ids = [segment.id for segment in segment_list]
summaries: dict[str, str | None] = {}
if segment_ids:
summary_records = SummaryIndexService.get_segments_summaries(
segment_ids=segment_ids, dataset_id=dataset_id_str, session=db.session()
segment_ids=segment_ids, dataset_id=dataset_id_str, session=session
)
summaries = {chunk_id: summary.summary_content for chunk_id, summary in summary_records.items()}
response = {
"data": segment_responses_with_summaries(segment_list, summaries),
"data": segment_responses_with_summaries(segment_list, summaries, session=session),
"limit": limit,
"total": segments.total,
"total_pages": segments.pages,
@ -300,17 +302,18 @@ class DatasetDocumentSegmentListApi(Resource):
@console_ns.response(204, "Segments deleted successfully")
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def delete(self, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def delete(self, session: Session, current_user: Account, dataset_id: UUID, document_id: UUID):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_ids = request.args.getlist("segment_id")
@ -319,10 +322,10 @@ class DatasetDocumentSegmentListApi(Resource):
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
SegmentService.delete_segments(segment_ids, document, dataset, db.session())
SegmentService.delete_segments(segment_ids, document, dataset, session)
return "", 204
@ -339,8 +342,10 @@ class DatasetDocumentSegmentApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def patch(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
@ -348,11 +353,11 @@ class DatasetDocumentSegmentApi(Resource):
action: Literal["enable", "disable"],
):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# check user's model setting
@ -362,7 +367,7 @@ class DatasetDocumentSegmentApi(Resource):
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
@ -388,7 +393,7 @@ class DatasetDocumentSegmentApi(Resource):
if cache_result is not None:
raise InvalidActionError("Document is being indexed, please try again later")
try:
SegmentService.update_segments_status(segment_ids, action, dataset, document, db.session())
SegmentService.update_segments_status(segment_ids, action, dataset, document, session)
except Exception as e:
raise InvalidActionError(str(e))
return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@ -408,15 +413,23 @@ class DatasetDocumentSegmentAddApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def post(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
document_id: UUID,
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
if not current_user.is_dataset_editor:
@ -438,22 +451,21 @@ class DatasetDocumentSegmentAddApi(Resource):
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
# validate args
payload = SegmentCreatePayload.model_validate(console_ns.payload or {})
payload_dict = payload.model_dump(exclude_none=True)
SegmentService.segment_create_args_validate(payload_dict, document)
segment = type_cast(
DocumentSegment,
SegmentService.create_segment(payload_dict, document, dataset, db.session()),
)
segment = type_cast(DocumentSegment, SegmentService.create_segment(payload_dict, document, dataset, session))
summary = SummaryIndexService.get_segment_summary(
segment_id=segment.id, dataset_id=dataset_id_str, session=db.session()
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),
"data": segment_response_with_summary(
segment, summary.summary_content if summary else None, session=session
),
"doc_form": document.doc_form,
}
return dump_response(SegmentDetailResponse, response), 200
@ -472,26 +484,33 @@ class DatasetDocumentSegmentUpdateApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def patch(
self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID, segment_id: UUID
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
document_id: UUID,
segment_id: UUID,
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
@ -511,7 +530,7 @@ class DatasetDocumentSegmentUpdateApi(Resource):
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
# validate args
payload = SegmentUpdatePayload.model_validate(console_ns.payload or {})
payload_dict = payload.model_dump(exclude_none=True)
@ -523,13 +542,15 @@ class DatasetDocumentSegmentUpdateApi(Resource):
segment,
document,
dataset,
db.session(),
session,
)
summary = SummaryIndexService.get_segment_summary(
segment_id=segment.id, dataset_id=dataset_id_str, session=db.session()
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),
"data": segment_response_with_summary(
segment, summary.summary_content if summary else None, session=session
),
"doc_form": document.doc_form,
}
return dump_response(SegmentDetailResponse, response), 200
@ -543,31 +564,38 @@ class DatasetDocumentSegmentUpdateApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def delete(
self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID, segment_id: UUID
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
document_id: UUID,
segment_id: UUID,
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
SegmentService.delete_segment(segment, document, dataset, db.session())
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
SegmentService.delete_segment(segment, document, dataset, session)
return "", 204
@ -587,22 +615,30 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
@with_session
def post(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
document_id: UUID,
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
payload = BatchImportPayload.model_validate(console_ns.payload or {})
upload_file_id = payload.upload_file_id
upload_file = db.session.scalar(select(UploadFile).where(UploadFile.id == upload_file_id).limit(1))
upload_file = session.scalar(select(UploadFile).where(UploadFile.id == upload_file_id).limit(1))
if not upload_file:
raise NotFound("UploadFile not found.")
@ -660,23 +696,30 @@ class ChildChunkAddApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def post(
self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID, segment_id: UUID
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
document_id: UUID,
segment_id: UUID,
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
# check embedding model setting
@ -696,11 +739,11 @@ class ChildChunkAddApi(Resource):
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
# validate args
try:
payload = ChildChunkCreatePayload.model_validate(console_ns.payload or {})
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session())
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, session)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200
@ -713,21 +756,22 @@ class ChildChunkAddApi(Resource):
@account_initialization_required
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
_get_segment_for_document(dataset, document, segment_id_str)
_get_segment_for_document(session, dataset, document, segment_id_str)
args = query_params_from_request(ChildChunkListQuery, use_defaults_for_malformed_ints=True)
page = args.page
@ -735,7 +779,13 @@ class ChildChunkAddApi(Resource):
keyword = args.keyword
child_chunks = SegmentService.get_child_chunks(
segment_id_str, document_id_str, dataset_id_str, page, limit, keyword
segment_id_str,
document_id_str,
dataset_id_str,
page,
limit,
keyword,
session=session,
)
response = {
"data": child_chunks.items,
@ -761,34 +811,41 @@ class ChildChunkAddApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def patch(
self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID, segment_id: UUID
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
document_id: UUID,
segment_id: UUID,
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
# validate args
payload = ChildChunkBatchUpdatePayload.model_validate(console_ns.payload or {})
try:
child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset, db.session())
child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset, session)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
return dump_response(ChildChunkBatchUpdateResponse, {"data": child_chunks}), 200
@ -807,8 +864,10 @@ class ChildChunkUpdateApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def delete(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
@ -818,31 +877,31 @@ class ChildChunkUpdateApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
segment_id_str = str(segment_id)
segment_ref, _ = _get_segment_for_document(dataset, document, segment_id_str)
segment_ref, _ = _get_segment_for_document(session, dataset, document, segment_id_str)
child_chunk_id_str = str(child_chunk_id)
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session())
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, session=session)
if not child_chunk:
raise NotFound("Child chunk not found.")
try:
SegmentService.delete_child_chunk(child_chunk, dataset, db.session())
SegmentService.delete_child_chunk(child_chunk, dataset, session)
except ChildChunkDeleteIndexServiceError as e:
raise ChildChunkDeleteIndexError(str(e))
return "", 204
@ -858,8 +917,10 @@ class ChildChunkUpdateApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_session
def patch(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
@ -869,34 +930,34 @@ class ChildChunkUpdateApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
if not current_user.is_dataset_editor:
raise Forbidden()
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
segment_id_str = str(segment_id)
segment_ref, segment = _get_segment_for_document(dataset, document, segment_id_str)
segment_ref, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
child_chunk_id_str = str(child_chunk_id)
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session())
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, session=session)
if not child_chunk:
raise NotFound("Child chunk not found.")
# validate args
try:
payload = ChildChunkUpdatePayload.model_validate(console_ns.payload or {})
child_chunk = SegmentService.update_child_chunk(
payload.content, child_chunk, segment, document, dataset, db.session()
payload.content, child_chunk, segment, document, dataset, session
)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))

51
api/controllers/console/datasets/external.py
View File

@ -1,3 +1,4 @@
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from uuid import UUID
@ -10,9 +11,13 @@ from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
import services
from controllers.common.fields import UsageCountResponse
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.schema import (
query_params_from_model,
register_response_schema_models,
register_schema_models,
)
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.app.wraps import with_session
from controllers.console.datasets.error import DatasetNameDuplicateError
from controllers.console.wraps import (
RBACPermission,
@ -25,10 +30,11 @@ from controllers.console.wraps import (
with_current_user,
)
from fields.base import ResponseModel
from fields.dataset_fields import DatasetDetailResponse
from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from libs.helper import dump_response
from libs.login import login_required
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 services.external_knowledge_service import ExternalDatasetService
@ -90,6 +96,28 @@ 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) # noqa: no-new-getattr response adapter 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
@ -162,14 +190,15 @@ class ExternalApiTemplateListApi(Resource):
@login_required
@with_current_tenant_id
@account_initialization_required
def get(self, current_tenant_id: str):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str):
query = ExternalApiTemplateListQuery.model_validate(request.args.to_dict())
external_knowledge_apis, total = ExternalDatasetService.get_external_knowledge_apis(
query.page, query.limit, current_tenant_id, query.keyword
query.page, query.limit, current_tenant_id, query.keyword, session=session
)
return ExternalKnowledgeApiListResponse(
data=[ExternalKnowledgeApiResponse.model_validate(item) for item in external_knowledge_apis],
data=[external_knowledge_api_response(item, session=session) for item in external_knowledge_apis],
has_more=len(external_knowledge_apis) == query.limit,
limit=query.limit,
total=total,
@ -210,7 +239,7 @@ class ExternalApiTemplateListApi(Resource):
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
return dump_response(ExternalKnowledgeApiResponse, external_knowledge_api), 201
return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 201
@console_ns.route("/datasets/external-knowledge-api/<uuid:external_knowledge_api_id>")
@ -237,7 +266,7 @@ class ExternalApiTemplateApi(Resource):
if external_knowledge_api is None:
raise NotFound("API template not found.")
return dump_response(ExternalKnowledgeApiResponse, external_knowledge_api), 200
return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 200
@console_ns.doc("update_external_api_template")
@console_ns.doc(description="Update external knowledge API template")
@ -269,7 +298,7 @@ class ExternalApiTemplateApi(Resource):
session=session,
)
return dump_response(ExternalKnowledgeApiResponse, external_knowledge_api), 200
return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 200
@setup_required
@login_required
@ -354,7 +383,9 @@ class ExternalDatasetCreateApi(Resource):
[dataset_id_str],
session=session,
)
data = DatasetDetailResponse.model_validate(dataset).model_dump(mode="json")
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

2
api/controllers/console/datasets/hit_testing.py
View File

@ -53,7 +53,7 @@ class HitTestingApi(Resource, DatasetsHitTestingBase):
) -> dict[str, object]:
dataset_id_str = str(dataset_id)
dataset = self.get_and_validate_dataset(dataset_id_str, current_user, current_tenant_id)
dataset = self.get_and_validate_dataset(session, dataset_id_str, current_user, current_tenant_id)
args = self.parse_args(console_ns.payload)
self.hit_testing_args_check(args)

10
api/controllers/console/datasets/hit_testing_base.py
View File

@ -19,7 +19,6 @@ from core.errors.error import (
ProviderTokenNotInitError,
QuotaExceededError,
)
from extensions.ext_database import db
from graphon.model_runtime.errors.invoke import InvokeError
from libs.login import resolve_account_fallback
from models.account import Account
@ -83,15 +82,18 @@ class DatasetsHitTestingBase:
@staticmethod
def get_and_validate_dataset(
dataset_id: str, current_user: Account | None = None, current_tenant_id: str | None = None
session: Session,
dataset_id: str,
current_user: Account | None = None,
current_tenant_id: str | None = None,
) -> Dataset:
current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
dataset = DatasetService.get_dataset(dataset_id, db.session())
dataset = DatasetService.get_dataset(dataset_id, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))

64
api/controllers/console/datasets/metadata.py
View File

@ -2,10 +2,12 @@ from typing import Literal
from uuid import UUID
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.controller_schemas import MetadataUpdatePayload
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,
@ -17,7 +19,6 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from extensions.ext_database import db
from fields.dataset_fields import (
DatasetMetadataBuiltInFieldsResponse,
DatasetMetadataListResponse,
@ -57,17 +58,18 @@ class DatasetMetadataCreateApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID):
@with_session
def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID):
metadata_args = MetadataArgs.model_validate(console_ns.payload or {})
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.create_metadata(
dataset_id_str, metadata_args, current_user, current_tenant_id, session=db.session()
dataset_id_str, metadata_args, current_user, current_tenant_id, session=session
)
return dump_response(DatasetMetadataResponse, metadata), 201
@ -79,12 +81,13 @@ class DatasetMetadataCreateApi(Resource):
200, "Metadata retrieved successfully", console_ns.models[DatasetMetadataListResponse.__name__]
)
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
metadata = MetadataService.get_dataset_metadatas(dataset, session=db.session())
metadata = MetadataService.get_dataset_metadatas(dataset, session)
return dump_response(DatasetMetadataListResponse, metadata), 200
@ -99,19 +102,27 @@ class DatasetMetadataApi(Resource):
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def patch(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, metadata_id: UUID):
@with_session
def patch(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
dataset_id: UUID,
metadata_id: UUID,
):
payload = MetadataUpdatePayload.model_validate(console_ns.payload or {})
name = payload.name
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.update_metadata_name(
dataset_id_str, metadata_id_str, name, current_user, current_tenant_id, session=db.session()
dataset_id_str, metadata_id_str, name, current_user, current_tenant_id, session=session
)
return dump_response(DatasetMetadataResponse, metadata), 200
@ -122,15 +133,16 @@ class DatasetMetadataApi(Resource):
@console_ns.response(204, "Metadata deleted successfully")
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def delete(self, current_user: Account, dataset_id: UUID, metadata_id: UUID):
@with_session
def delete(self, session: Session, 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(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session=db.session())
MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session)
# Frontend callers only await success and invalidate metadata caches; no response body is consumed.
return "", 204
@ -160,18 +172,19 @@ class DatasetMetadataBuiltInFieldActionApi(Resource):
@console_ns.response(204, "Action completed successfully")
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable"]):
@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, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
match action:
case "enable":
MetadataService.enable_built_in_field(dataset, session=db.session())
MetadataService.enable_built_in_field(dataset, session)
case "disable":
MetadataService.disable_built_in_field(dataset, session=db.session())
MetadataService.disable_built_in_field(dataset, session)
# Frontend callers only await success and invalidate metadata caches; no response body is consumed.
return "", 204
@ -189,16 +202,17 @@ class DocumentMetadataEditApi(Resource):
)
@with_current_user
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID):
@with_session
def post(self, session: Session, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata_args = MetadataOperationData.model_validate(console_ns.payload or {})
MetadataService.update_documents_metadata(dataset, metadata_args, current_user, session=db.session())
MetadataService.update_documents_metadata(dataset, metadata_args, current_user, session=session)
# Frontend callers only await success and invalidate caches; no response body is consumed.
return "", 204

10
api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py
View File

@ -27,7 +27,7 @@ from controllers.console.app.workflow import (
WorkflowResponse,
)
from controllers.console.app.wraps import with_session
from controllers.console.datasets.wraps import get_rag_pipeline
from controllers.console.datasets.wraps import get_rag_pipeline, load_rag_pipeline
from controllers.console.wraps import (
RBACPermission,
RBACResourceScope,
@ -344,11 +344,11 @@ class DraftRagPipelineRunApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_current_user
@with_session
@get_rag_pipeline
def post(self, session: Session, current_user: Account, pipeline: Pipeline):
def post(self, session: Session, current_user: Account, pipeline_id: UUID):
"""
Run draft workflow
"""
pipeline = load_rag_pipeline(session, str(pipeline_id))
payload = DraftWorkflowRunPayload.model_validate(console_ns.payload or {})
args = payload.model_dump()
@ -378,11 +378,11 @@ class PublishedRagPipelineRunApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
@with_current_user
@with_session
@get_rag_pipeline
def post(self, session: Session, current_user: Account, pipeline: Pipeline):
def post(self, session: Session, current_user: Account, pipeline_id: UUID):
"""
Run published workflow
"""
pipeline = load_rag_pipeline(session, str(pipeline_id))
payload = PublishedWorkflowRunPayload.model_validate(console_ns.payload or {})
args = payload.model_dump(exclude_none=True)
streaming = payload.response_mode == "streaming"

23
api/controllers/console/datasets/wraps.py
View File

@ -1,13 +1,21 @@
from collections.abc import Callable
from functools import wraps
from sqlalchemy import select
from sqlalchemy.orm import Session
from controllers.console.datasets.error import PipelineNotFoundError
from extensions.ext_database import db
from libs.login import current_account_with_tenant
from models.dataset import Pipeline
from services.rag_pipeline.rag_pipeline import RagPipelineService
def load_rag_pipeline(session: Session, pipeline_id: str) -> Pipeline:
_, current_tenant_id = current_account_with_tenant()
pipeline = RagPipelineService.get_pipeline_by_id(pipeline_id, current_tenant_id, session=session)
if not pipeline:
raise PipelineNotFoundError()
return pipeline
def get_rag_pipeline[**P, R](view_func: Callable[P, R]) -> Callable[P, R]:
@ -16,22 +24,11 @@ def get_rag_pipeline[**P, R](view_func: Callable[P, R]) -> Callable[P, R]:
if not kwargs.get("pipeline_id"):
raise ValueError("missing pipeline_id in path parameters")
_, current_tenant_id = current_account_with_tenant()
pipeline_id = kwargs.get("pipeline_id")
pipeline_id = str(pipeline_id)
del kwargs["pipeline_id"]
stmt = select(Pipeline).where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_tenant_id).limit(1)
# Migrated handlers pass the request Session as args[1]; legacy handlers still use db.session.
session = args[1] if len(args) > 1 and isinstance(args[1], Session) else db.session
pipeline = session.scalar(stmt)
if not pipeline:
raise PipelineNotFoundError()
kwargs["pipeline"] = pipeline
kwargs["pipeline"] = load_rag_pipeline(db.session(), pipeline_id)
return view_func(*args, **kwargs)

11
api/controllers/console/explore/audio.py
View File

@ -50,14 +50,19 @@ register_response_schema_models(console_ns, AudioBinaryResponse, AudioTranscript
class ChatAudioApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[AudioTranscriptResponse.__name__])
def post(self, installed_app: InstalledApp):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
file = request.files["file"]
try:
response = AudioService.transcript_asr(app_model=app_model, file=file, end_user=None)
response = AudioService.transcript_asr(
app_model=app_model,
file=file,
session=db.session(),
end_user=None,
)
return response
except services.errors.app_model_config.AppModelConfigBrokenError:
@ -96,7 +101,7 @@ class ChatTextApi(InstalledAppResource):
@console_ns.expect(console_ns.models[TextToAudioPayload.__name__])
@console_ns.response(200, "Success", console_ns.models[AudioBinaryResponse.__name__])
def post(self, installed_app: InstalledApp):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
try:

16
api/controllers/console/explore/completion.py
View File

@ -90,7 +90,7 @@ class CompletionApi(InstalledAppResource):
@with_current_user
@with_session
def post(self, session: Session, current_user: Account, installed_app: InstalledApp):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
if app_model.mode != AppMode.COMPLETION:
@ -146,8 +146,9 @@ class CompletionApi(InstalledAppResource):
class CompletionStopApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@with_current_user_id
def post(self, current_user_id: str, installed_app: InstalledApp, task_id: str):
app_model = installed_app.app
@with_session(write=False)
def post(self, session: Session, current_user_id: str, installed_app: InstalledApp, task_id: str):
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
if app_model.mode != AppMode.COMPLETION:
@ -173,7 +174,7 @@ class ChatApi(InstalledAppResource):
@with_current_user
@with_session
def post(self, session: Session, current_user: Account, installed_app: InstalledApp):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)
@ -195,7 +196,7 @@ class ChatApi(InstalledAppResource):
app_model=app_model,
conversation_id=payload.conversation_id,
user=current_user,
session=db.session(),
session=session,
)
response = AppGenerateService.generate(
@ -240,8 +241,9 @@ class ChatApi(InstalledAppResource):
class ChatStopApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
@with_current_user_id
def post(self, current_user_id: str, installed_app: InstalledApp, task_id: str):
app_model = installed_app.app
@with_session(write=False)
def post(self, session: Session, current_user_id: str, installed_app: InstalledApp, task_id: str):
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)

24
api/controllers/console/explore/conversation.py
View File

@ -16,6 +16,7 @@ from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from fields.conversation_fields import (
ConversationInfiniteScrollPagination,
ConversationResponseSource,
ResultResponse,
SimpleConversation,
)
@ -53,7 +54,7 @@ class ConversationListApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[ConversationInfiniteScrollPagination.__name__])
@with_current_user
def get(self, current_user: Account, installed_app: InstalledApp):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)
@ -84,7 +85,13 @@ class ConversationListApi(InstalledAppResource):
pinned=args.pinned,
)
adapter = TypeAdapter(SimpleConversation)
conversations = [adapter.validate_python(item, from_attributes=True) for item in pagination.data]
conversations = [
adapter.validate_python(
ConversationResponseSource(item, session=session),
from_attributes=True,
)
for item in pagination.data
]
return ConversationInfiniteScrollPagination(
limit=pagination.limit,
has_more=pagination.has_more,
@ -102,7 +109,7 @@ class ConversationApi(InstalledAppResource):
@console_ns.response(204, "Conversation deleted successfully")
@with_current_user
def delete(self, current_user: Account, installed_app: InstalledApp, c_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)
@ -127,7 +134,7 @@ class ConversationRenameApi(InstalledAppResource):
@console_ns.response(200, "Conversation renamed successfully", console_ns.models[SimpleConversation.__name__])
@with_current_user
def post(self, current_user: Account, installed_app: InstalledApp, c_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)
@ -139,12 +146,13 @@ class ConversationRenameApi(InstalledAppResource):
payload = ConversationRenamePayload.model_validate(console_ns.payload or {})
try:
session = db.session()
conversation = ConversationService.rename(
app_model, conversation_id, current_user, payload.name, payload.auto_generate, session=db.session()
app_model, conversation_id, current_user, payload.name, payload.auto_generate, session=session
)
return (
TypeAdapter(SimpleConversation)
.validate_python(conversation, from_attributes=True)
.validate_python(ConversationResponseSource(conversation, session=session), from_attributes=True)
.model_dump(mode="json")
)
except ConversationNotExistsError:
@ -159,7 +167,7 @@ class ConversationPinApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[ResultResponse.__name__])
@with_current_user
def patch(self, current_user: Account, installed_app: InstalledApp, c_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)
@ -184,7 +192,7 @@ class ConversationUnPinApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[ResultResponse.__name__])
@with_current_user
def patch(self, current_user: Account, installed_app: InstalledApp, c_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)

18
api/controllers/console/explore/message.py
View File

@ -28,7 +28,7 @@ from controllers.console.wraps import with_current_user
from core.app.entities.app_invoke_entities import InvokeFrom
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from extensions.ext_database import db
from fields.conversation_fields import ResultResponse
from fields.conversation_fields import MessageResponseSource, ResultResponse
from fields.message_fields import (
ExploreMessageInfiniteScrollPagination,
ExploreMessageListItem,
@ -76,7 +76,8 @@ class MessageListApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[ExploreMessageInfiniteScrollPagination.__name__])
@with_current_user
def get(self, current_user: Account, installed_app: InstalledApp):
app_model = installed_app.app
session = db.session()
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
@ -92,10 +93,13 @@ class MessageListApi(InstalledAppResource):
args.conversation_id,
args.first_id or None,
args.limit,
session=db.session(),
session=session,
)
adapter = TypeAdapter(ExploreMessageListItem)
items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
items = [
adapter.validate_python(MessageResponseSource(message, session=session), from_attributes=True)
for message in pagination.data
]
return ExploreMessageInfiniteScrollPagination(
limit=pagination.limit,
has_more=pagination.has_more,
@ -116,7 +120,7 @@ class MessageFeedbackApi(InstalledAppResource):
@console_ns.response(200, "Feedback submitted successfully", console_ns.models[ResultResponse.__name__])
@with_current_user
def post(self, current_user: Account, installed_app: InstalledApp, message_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
@ -149,7 +153,7 @@ class MessageMoreLikeThisApi(InstalledAppResource):
@with_current_user
@with_session
def get(self, session: Session, current_user: Account, installed_app: InstalledApp, message_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
if app_model.mode != "completion":
@ -199,7 +203,7 @@ class MessageSuggestedQuestionApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[SuggestedQuestionsResponse.__name__])
@with_current_user
def get(self, current_user: Account, installed_app: InstalledApp, message_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
app_mode = AppMode.value_of(app_model.mode)

17
api/controllers/console/explore/parameter.py
View File

@ -9,7 +9,7 @@ from controllers.console.app.error import AppUnavailableError
from controllers.console.explore.wraps import InstalledAppResource
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
from extensions.ext_database import db
from models.model import AppMode, InstalledApp
from models.model import AppMode, InstalledApp, load_annotation_reply_config
from services.app_service import AppService
@ -32,24 +32,29 @@ class AppParameterApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[fields.Parameters.__name__])
def get(self, installed_app: InstalledApp):
"""Retrieve app parameters."""
app_model = installed_app.app
session = db.session()
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
workflow = app_model.workflow
workflow = app_model.workflow_with_session(session=session)
if workflow is None:
raise AppUnavailableError()
features_dict: dict[str, Any] = workflow.features_dict
user_input_form = workflow.user_input_form(to_old_structure=True)
else:
app_model_config = app_model.app_model_config
app_model_config = app_model.app_model_config_with_session(session=session)
if app_model_config is None:
raise AppUnavailableError()
features_dict = cast(dict[str, Any], app_model_config.to_dict())
annotation_reply = load_annotation_reply_config(session, app_model.id)
features_dict = cast(
dict[str, Any],
app_model_config.to_dict(annotation_reply=annotation_reply),
)
user_input_form = features_dict.get("user_input_form", [])
@ -62,7 +67,7 @@ class ExploreAppMetaApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[ExploreAppMetaResponse.__name__])
def get(self, installed_app: InstalledApp):
"""Get app meta"""
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if not app_model:
raise ValueError("App not found")
return AppService().get_app_meta(app_model, session=db.session())

16
api/controllers/console/explore/saved_message.py
View File

@ -12,7 +12,7 @@ from controllers.console.explore.error import NotCompletionAppError
from controllers.console.explore.wraps import InstalledAppResource
from controllers.console.wraps import with_current_user
from extensions.ext_database import db
from fields.conversation_fields import ResultResponse
from fields.conversation_fields import MessageResponseSource, ResultResponse
from fields.message_fields import SavedMessageInfiniteScrollPagination, SavedMessageItem
from models import Account
from models.model import InstalledApp
@ -29,7 +29,8 @@ class SavedMessageListApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[SavedMessageInfiniteScrollPagination.__name__])
@with_current_user
def get(self, current_user: Account, installed_app: InstalledApp):
app_model = installed_app.app
session = db.session()
app_model = installed_app.app_with_session(session=session)
if app_model is None:
raise AppUnavailableError()
if app_model.mode != "completion":
@ -38,10 +39,13 @@ class SavedMessageListApi(InstalledAppResource):
args = SavedMessageListQuery.model_validate(request.args.to_dict())
pagination = SavedMessageService.pagination_by_last_id(
app_model, current_user, str(args.last_id) if args.last_id else None, args.limit, session=db.session()
app_model, current_user, str(args.last_id) if args.last_id else None, args.limit, session=session
)
adapter = TypeAdapter(SavedMessageItem)
items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
items = [
adapter.validate_python(MessageResponseSource(message, session=session), from_attributes=True)
for message in pagination.data
]
return SavedMessageInfiniteScrollPagination(
limit=pagination.limit,
has_more=pagination.has_more,
@ -52,7 +56,7 @@ class SavedMessageListApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[ResultResponse.__name__])
@with_current_user
def post(self, current_user: Account, installed_app: InstalledApp):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()
if app_model.mode != "completion":
@ -75,7 +79,7 @@ class SavedMessageApi(InstalledAppResource):
@console_ns.response(204, "Saved message deleted successfully")
@with_current_user
def delete(self, current_user: Account, installed_app: InstalledApp, message_id: UUID):
app_model = installed_app.app
app_model = installed_app.app_with_session(session=db.session())
if app_model is None:
raise AppUnavailableError()

81
api/controllers/console/explore/trial.py
View File

@ -1,4 +1,6 @@
import logging
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Literal
@ -65,11 +67,12 @@ from libs import helper
from libs.helper import dump_response, to_timestamp, uuid_value
from models import Account
from models.account import TenantStatus
from models.model import AppMode, Site
from models.model import AppMode, Site, load_annotation_reply_config
from models.workflow import Workflow
from services.account_service import TenantService
from services.app_generate_service import AppGenerateService
from services.app_ref_service import AppRefService
from services.app_service import AppService
from services.app_service import AppResponseView, AppService
from services.audio_service import AudioService
from services.dataset_service import DatasetService
from services.errors.audio import (
@ -379,6 +382,27 @@ class TrialWorkflowResponse(ResponseModel):
return to_timestamp(value)
@dataclass(frozen=True)
class TrialWorkflowResponseSource:
workflow: Workflow
session: Session
@property
def created_by_account(self) -> Account | None:
return self.workflow.get_created_by_account(session=self.session)
@property
def updated_by_account(self) -> Account | None:
return self.workflow.get_updated_by_account(session=self.session)
@property
def tool_published(self) -> bool:
return self.workflow.get_tool_published(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.workflow, name) # noqa: no-new-getattr response adapter delegates model fields
register_schema_models(
console_ns,
WorkflowRunRequest,
@ -594,7 +618,12 @@ class TrialChatAudioApi(TrialAppResource):
app_id = app_model.id
user_id = current_user.id
response = AudioService.transcript_asr(app_model=app_model, file=file, end_user=None)
response = AudioService.transcript_asr(
app_model=app_model,
file=file,
session=db.session(),
end_user=None,
)
RecommendedAppService.add_trial_app_record(app_id, user_id, session=db.session())
return response
except services.errors.app_model_config.AppModelConfigBrokenError:
@ -746,19 +775,21 @@ class TrialSitApi(Resource):
"""Resource for trial app sites."""
@console_ns.response(200, "Success", console_ns.models[SiteResponse.__name__])
@with_session(write=False)
@get_app_model_with_trial(None)
def get(self, app_model):
def get(self, session: Session, app_model):
"""Retrieve app site info.
Returns the site configuration for the application including theme, icons, and text.
"""
site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
site = session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
if not site:
raise Forbidden()
assert app_model.tenant
if app_model.tenant.status == TenantStatus.ARCHIVE:
tenant = TenantService.get_tenant_by_id(app_model.tenant_id, session=session)
assert tenant
if tenant.status == TenantStatus.ARCHIVE:
raise Forbidden()
return SiteResponse.model_validate(site).model_dump(mode="json")
@ -768,26 +799,29 @@ class TrialAppParameterApi(Resource):
"""Resource for app variables."""
@console_ns.response(200, "Success", console_ns.models[ParametersResponse.__name__])
@with_session(write=False)
@get_app_model_with_trial(None)
def get(self, app_model):
def get(self, session: Session, app_model):
"""Retrieve app parameters."""
if app_model is None:
raise AppUnavailableError()
features_dict: Mapping[str, Any]
if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
workflow = app_model.workflow
workflow = app_model.workflow_with_session(session=session)
if workflow is None:
raise AppUnavailableError()
features_dict = workflow.features_dict
user_input_form = workflow.user_input_form(to_old_structure=True)
else:
app_model_config = app_model.app_model_config
app_model_config = app_model.app_model_config_with_session(session=session)
if app_model_config is None:
raise AppUnavailableError()
features_dict = app_model_config.to_dict()
annotation_reply = load_annotation_reply_config(session, app_model_config.app_id)
features_dict = app_model_config.to_dict(annotation_reply=annotation_reply)
user_input_form = features_dict.get("user_input_form", [])
@ -797,43 +831,52 @@ class TrialAppParameterApi(Resource):
class AppApi(Resource):
@console_ns.response(200, "Success", console_ns.models[TrialAppDetailResponse.__name__])
@with_session(write=False)
@get_app_model_with_trial(None)
def get(self, app_model):
def get(self, session: Session, app_model):
"""Get app detail"""
app_service = AppService()
app_model = app_service.get_app(app_model)
app_model = app_service.get_app(app_model, session=session)
return dump_response(TrialAppDetailResponse, app_model)
return TrialAppDetailResponse.model_validate(
AppResponseView(app_model, session=session),
from_attributes=True,
).model_dump(mode="json")
class AppWorkflowApi(Resource):
@console_ns.response(200, "Success", console_ns.models[TrialWorkflowResponse.__name__])
@with_session(write=False)
@get_app_model_with_trial(None)
def get(self, app_model):
def get(self, session: Session, app_model):
"""Get workflow detail"""
if not app_model.workflow_id:
raise AppUnavailableError()
workflow = db.session.get(Workflow, app_model.workflow_id)
workflow = app_model.workflow_with_session(session=session)
if workflow is None:
raise AppUnavailableError()
return dump_response(TrialWorkflowResponse, workflow)
return TrialWorkflowResponse.model_validate(
TrialWorkflowResponseSource(workflow=workflow, session=session),
from_attributes=True,
).model_dump(mode="json")
class DatasetListApi(Resource):
@console_ns.doc(params=query_params_from_model(TrialDatasetListQuery))
@console_ns.response(200, "Success", console_ns.models[TrialDatasetListResponse.__name__])
@with_session(write=False)
@get_app_model_with_trial(None)
def get(self, app_model):
def get(self, session: Session, app_model):
page = request.args.get("page", default=1, type=int)
limit = request.args.get("limit", default=20, type=int)
ids = request.args.getlist("ids")
tenant_id = app_model.tenant_id
if ids:
datasets, total = DatasetService.get_datasets_by_ids(ids, tenant_id)
datasets, total = DatasetService.get_datasets_by_ids(ids, tenant_id, session=session)
else:
raise NeedAddIdsError()

7
api/controllers/console/explore/workflow.py
View File

@ -51,7 +51,7 @@ class InstalledAppWorkflowRunApi(InstalledAppResource):
"""
Run workflow
"""
app_model = installed_app.app
app_model = installed_app.app_with_session(session=session)
if not app_model:
raise NotWorkflowAppError()
app_mode = AppMode.value_of(app_model.mode)
@ -92,11 +92,12 @@ class InstalledAppWorkflowRunApi(InstalledAppResource):
@console_ns.route("/installed-apps/<uuid:installed_app_id>/workflows/tasks/<string:task_id>/stop")
class InstalledAppWorkflowTaskStopApi(InstalledAppResource):
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
def post(self, installed_app: InstalledApp, task_id: str):
@with_session(write=False)
def post(self, session: Session, installed_app: InstalledApp, task_id: str):
"""
Stop workflow task
"""
app_model = installed_app.app
app_model = installed_app.app_with_session(session=session)
if not app_model:
raise NotWorkflowAppError()
app_mode = AppMode.value_of(app_model.mode)

9
api/controllers/console/explore/wraps.py
View File

@ -30,7 +30,7 @@ def installed_app_required[**P, R](view: Callable[Concatenate[InstalledApp, P],
if installed_app is None:
raise NotFound("Installed app not found")
if not installed_app.app:
if not installed_app.app_with_session(session=db.session()):
db.session.delete(installed_app)
db.session.commit()
@ -74,17 +74,18 @@ def trial_app_required[**P, R](view: Callable[Concatenate[App, P], R] | None = N
@wraps(view)
def decorated(app_id: str, *args: P.args, **kwargs: P.kwargs):
current_user, _ = current_account_with_tenant()
session = db.session()
trial_app = db.session.scalar(select(TrialApp).where(TrialApp.app_id == str(app_id)).limit(1))
trial_app = session.scalar(select(TrialApp).where(TrialApp.app_id == str(app_id)).limit(1))
if trial_app is None:
raise TrialAppNotAllowed()
app = trial_app.app
app = trial_app.app_with_session(session=session)
if app is None:
raise TrialAppNotAllowed()
account_trial_app_record = db.session.scalar(
account_trial_app_record = session.scalar(
select(AccountTrialAppRecord)
.where(AccountTrialAppRecord.account_id == current_user.id, AccountTrialAppRecord.app_id == app_id)
.limit(1)

10
api/controllers/console/socketio/workflow.py
View File

@ -4,7 +4,7 @@ from typing import cast
from flask import Request as FlaskRequest
from extensions.ext_database import db
from core.db.session_factory import session_factory
from extensions.ext_socketio import sio
from libs.passport import PassportService
from libs.token import extract_access_token
@ -43,8 +43,8 @@ def socket_connect(sid, environ, auth):
logging.warning("Socket connect rejected: missing user_id (sid=%s)", sid)
return False
with sio.app.app_context():
user = AccountService.load_logged_in_account(account_id=user_id, session=db.session())
with sio.app.app_context(), session_factory.create_session() as session:
user = AccountService.load_logged_in_account(account_id=user_id, session=session)
if not user:
logging.warning("Socket connect rejected: user not found (user_id=%s, sid=%s)", user_id, sid)
return False
@ -69,8 +69,8 @@ def handle_user_connect(sid, data):
if not workflow_id:
return {"msg": "workflow_id is required"}, 400
with sio.app.app_context():
result = collaboration_service.authorize_and_join_workflow_room(workflow_id, sid, session=db.session())
with sio.app.app_context(), session_factory.create_session() as session:
result = collaboration_service.authorize_and_join_workflow_room(workflow_id, sid, session=session)
if not result:
return {"msg": "unauthorized"}, 401

54
api/controllers/console/workspace/workspace.py
View File

@ -6,7 +6,8 @@ from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import select
from werkzeug.exceptions import Unauthorized
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound, Unauthorized
import services
from configs import dify_config
@ -23,6 +24,7 @@ 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.admin import admin_required
from controllers.console.error import AccountNotLinkTenantError
@ -220,10 +222,11 @@ class TenantListApi(Resource):
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account):
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, current_user: Account):
tenant_rows: list[tuple[Tenant, TenantAccountJoin]] = [
(tenant, membership)
for tenant, membership in TenantService.get_workspaces_for_account(current_user.id, session=db.session())
for tenant, membership in TenantService.get_workspaces_for_account(current_user.id, session=session)
if tenant.status == TenantStatus.NORMAL
]
tenants = [tenant for tenant, _ in tenant_rows]
@ -274,11 +277,12 @@ class WorkspaceListApi(Resource):
@console_ns.response(HTTPStatus.OK, "Success", console_ns.models[WorkspacePaginationResponse.__name__])
@setup_required
@admin_required
def get(self):
@with_session(write=False)
def get(self, session: Session):
args = query_params_from_request(WorkspaceListQuery)
stmt = select(Tenant).order_by(Tenant.created_at.desc())
tenants = paginate_query(stmt, page=args.page, per_page=args.limit)
tenants = paginate_query(stmt, session=session, page=args.page, per_page=args.limit)
has_more = False
if tenants.has_next:
@ -297,7 +301,8 @@ class TenantApi(Resource):
@account_initialization_required
@console_ns.response(HTTPStatus.OK, "Success", console_ns.models[TenantInfoResponse.__name__])
@with_current_user
def post(self, current_user: Account):
@with_session
def post(self, session: Session, current_user: Account):
if request.path == "/info":
logger.warning("Deprecated URL /info was used.")
@ -306,17 +311,17 @@ class TenantApi(Resource):
raise ValueError("No current tenant")
if tenant.status == TenantStatus.ARCHIVE:
tenants = TenantService.get_join_tenants(current_user, session=db.session())
tenants = TenantService.get_join_tenants(current_user, session=session)
# if there is any tenant, switch to the first one
if len(tenants) > 0:
TenantService.switch_tenant(current_user, tenants[0].id, session=db.session())
TenantService.switch_tenant(current_user, tenants[0].id, session=session)
tenant = tenants[0]
# else, raise Unauthorized
else:
raise Unauthorized("workspace is archived")
return (
dump_response(TenantInfoResponse, WorkspaceService.get_tenant_info(tenant, session=db.session())),
dump_response(TenantInfoResponse, WorkspaceService.get_tenant_info(tenant, session=session)),
HTTPStatus.OK,
)
@ -329,22 +334,23 @@ class SwitchWorkspaceApi(Resource):
@login_required
@account_initialization_required
@with_current_user
def post(self, current_user: Account):
@with_session
def post(self, session: Session, current_user: Account):
payload = console_ns.payload or {}
args = SwitchWorkspacePayload.model_validate(payload)
# Check whether the tenant_id belongs to the current account.
try:
TenantService.switch_tenant(current_user, args.tenant_id, session=db.session())
TenantService.switch_tenant(current_user, args.tenant_id, session=session)
except Exception:
raise AccountNotLinkTenantError("Account not link tenant")
new_tenant = db.session.get(Tenant, args.tenant_id) # Get new tenant
new_tenant = TenantService.get_tenant_by_id(args.tenant_id, session=session)
if new_tenant is None:
raise ValueError("Tenant not found")
return SwitchWorkspaceResponse(
result="success", new_tenant=WorkspaceService.get_tenant_info(new_tenant, session=db.session())
result="success", new_tenant=WorkspaceService.get_tenant_info(new_tenant, session=session)
).model_dump(mode="json")
@ -357,10 +363,13 @@ class CustomConfigWorkspaceApi(Resource):
@account_initialization_required
@cloud_edition_billing_resource_check("workspace_custom")
@with_current_tenant_id
def post(self, current_tenant_id: str):
@with_session
def post(self, session: Session, current_tenant_id: str):
payload = console_ns.payload or {}
args = WorkspaceCustomConfigPayload.model_validate(payload)
tenant = db.get_or_404(Tenant, current_tenant_id)
tenant = TenantService.get_tenant_by_id(current_tenant_id, session=session)
if tenant is None:
raise NotFound()
custom_config_dict: TenantCustomConfigDict = {
"remove_webapp_brand": args.remove_webapp_brand
@ -372,10 +381,10 @@ class CustomConfigWorkspaceApi(Resource):
}
tenant.custom_config_dict = custom_config_dict
db.session.commit()
session.commit()
return WorkspaceTenantResultResponse(
result="success", tenant=WorkspaceService.get_tenant_info(tenant, session=db.session())
result="success", tenant=WorkspaceService.get_tenant_info(tenant, session=session)
).model_dump(mode="json")
@ -430,18 +439,21 @@ class WorkspaceInfoApi(Resource):
@account_initialization_required
# Change workspace name
@with_current_tenant_id
def post(self, current_tenant_id: str):
@with_session
def post(self, session: Session, current_tenant_id: str):
payload = console_ns.payload or {}
args = WorkspaceInfoPayload.model_validate(payload)
if not current_tenant_id:
raise ValueError("No current tenant")
tenant = db.get_or_404(Tenant, current_tenant_id)
tenant = TenantService.get_tenant_by_id(current_tenant_id, session=session)
if tenant is None:
raise NotFound()
tenant.name = args.name
db.session.commit()
session.commit()
return WorkspaceTenantResultResponse(
result="success", tenant=WorkspaceService.get_tenant_info(tenant, session=db.session())
result="success", tenant=WorkspaceService.get_tenant_info(tenant, session=session)
).model_dump(mode="json")

3
api/controllers/inner_api/app/dsl.py
View File

@ -54,9 +54,8 @@ class EnterpriseAppDSLImport(Resource):
if account is None:
return {"message": f"account '{args.creator_email}' not found or inactive"}, 404
account.set_tenant_id(workspace_id)
with Session(db.engine, expire_on_commit=False) as session:
account.set_tenant_id_with_session(workspace_id, session=session)
dsl_service = AppDslService(session)
result = dsl_service.import_app(
account=account,

17
api/controllers/openapi/_input_schema.py
View File

@ -4,9 +4,11 @@ from __future__ import annotations
from typing import Any, cast
from sqlalchemy.orm import Session
from controllers.service_api.app.error import AppUnavailableError
from models import App
from models.model import AppMode
from models.model import AppMode, load_annotation_reply_config
JSON_SCHEMA_DRAFT = "https://json-schema.org/draft/2020-12/schema"
@ -89,13 +91,13 @@ def _form_to_jsonschema(form: list[dict[str, Any]]) -> tuple[dict[str, Any], lis
return properties, required
def resolve_app_config(app: App) -> tuple[dict[str, Any], list[dict[str, Any]]]:
def resolve_app_config(app: App, *, session: Session) -> tuple[dict[str, Any], list[dict[str, Any]]]:
"""Resolve `(features_dict, user_input_form)` for parameters / schema derivation.
Raises `AppUnavailableError` on misconfigured apps.
"""
if app.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
workflow = app.workflow
workflow = app.workflow_with_session(session=session)
if workflow is None:
raise AppUnavailableError()
return (
@ -103,21 +105,22 @@ def resolve_app_config(app: App) -> tuple[dict[str, Any], list[dict[str, Any]]]:
cast(list[dict[str, Any]], workflow.user_input_form(to_old_structure=True)),
)
app_model_config = app.app_model_config
app_model_config = app.app_model_config_with_session(session=session)
if app_model_config is None:
raise AppUnavailableError()
features_dict = cast(dict[str, Any], app_model_config.to_dict())
annotation_reply = load_annotation_reply_config(session, app_model_config.app_id)
features_dict = cast(dict[str, Any], app_model_config.to_dict(annotation_reply=annotation_reply))
return features_dict, cast(list[dict[str, Any]], features_dict.get("user_input_form", []))
def build_input_schema(app: App) -> dict[str, Any]:
def build_input_schema(app: App, *, session: Session) -> dict[str, Any]:
"""Derive Draft 2020-12 JSON Schema from `user_input_form` + app mode.
chat / agent-chat / advanced-chat: top-level `query` (required, minLength=1) + `inputs` object.
completion / workflow: `inputs` object only.
Raises `AppUnavailableError` on misconfigured apps.
"""
_, user_input_form = resolve_app_config(app)
_, user_input_form = resolve_app_config(app, session=session)
inputs_props, inputs_required = _form_to_jsonschema(user_input_form)
properties: dict[str, Any] = {}

29
api/controllers/openapi/account.py
View File

@ -3,8 +3,10 @@ from __future__ import annotations
from datetime import UTC, datetime
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.session import with_session
from controllers.openapi import openapi_ns
from controllers.openapi._contract import accepts, returns
from controllers.openapi._models import (
@ -18,7 +20,6 @@ from controllers.openapi._models import (
)
from controllers.openapi.auth.composition import auth_router
from controllers.openapi.auth.data import AuthData
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from libs.oauth_bearer import (
Scope,
@ -41,14 +42,13 @@ from services.oauth_device_flow import (
class AccountApi(Resource):
@auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, AccountResponse, description="Account info")
def get(self, *, auth_data: AuthData):
@with_session(write=False)
def get(self, session: Session, *, auth_data: AuthData):
enforce(LIMIT_ME_PER_ACCOUNT, key=f"account:{auth_data.account_id}")
account_id_str = str(auth_data.account_id) if auth_data.account_id else None
account = AccountService.get_account_by_id(account_id_str, session=db.session()) if account_id_str else None
memberships = (
TenantService.get_account_memberships(account_id_str, session=db.session()) if account_id_str else []
)
account = AccountService.get_account_by_id(account_id_str, session=session) if account_id_str else None
memberships = TenantService.get_account_memberships(account_id_str, session=session) if account_id_str else []
default_ws_id = _pick_default_workspace(memberships)
return AccountResponse(
@ -64,8 +64,9 @@ class AccountApi(Resource):
class AccountSessionsSelfApi(Resource):
@auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, RevokeResponse, description="Session revoked")
def delete(self, *, auth_data: AuthData):
revoke_oauth_token(redis_client, str(auth_data.token_id), session=db.session())
@with_session
def delete(self, session: Session, *, auth_data: AuthData):
revoke_oauth_token(redis_client, str(auth_data.token_id), session=session)
return RevokeResponse(status="revoked")
@ -74,7 +75,8 @@ class AccountSessionsApi(Resource):
@auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, SessionListResponse, description="Session list")
@accepts(query=SessionListQuery)
def get(self, *, auth_data: AuthData, query: SessionListQuery):
@with_session(write=False)
def get(self, session: Session, *, auth_data: AuthData, query: SessionListQuery):
# SessionListQuery enforces the advertised bounds (extra='forbid', page>=1,
# 1<=limit<=MAX_PAGE_LIMIT) so the server rejects out-of-range paging rather
# than silently coercing (e.g. page=0 -> empty slice).
@ -83,7 +85,7 @@ class AccountSessionsApi(Resource):
page = query.page
limit = query.limit
all_rows = list_active_sessions(ctx, now, session=db.session())
all_rows = list_active_sessions(ctx, now, session=session)
total = len(all_rows)
sliced = all_rows[(page - 1) * limit : page * limit]
@ -114,15 +116,16 @@ class AccountSessionsApi(Resource):
class AccountSessionByIdApi(Resource):
@auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, RevokeResponse, description="Session revoked")
def delete(self, session_id: str, *, auth_data: AuthData):
@with_session
def delete(self, session: Session, session_id: str, *, auth_data: AuthData):
ctx = get_auth_ctx()
# 404 (not 403) on cross-subject so the endpoint doesn't leak
# token IDs that belong to other subjects.
if not token_belongs_to_subject(session_id, ctx, session=db.session()):
if not token_belongs_to_subject(session_id, ctx, session=session):
raise NotFound("session not found")
revoke_oauth_token(redis_client, session_id, session=db.session())
revoke_oauth_token(redis_client, session_id, session=session)
return RevokeResponse(status="revoked")

43
api/controllers/openapi/apps.py
View File

@ -6,11 +6,13 @@ import uuid as _uuid
from typing import Any, cast
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import Conflict, NotFound, UnprocessableEntity
from configs import dify_config
from controllers.common.app_access import AppAccessFilter, resolve_app_access_filter
from controllers.common.fields import Parameters
from controllers.common.session import with_session
from controllers.common.wraps import RBACPermission, RBACResourceScope
from controllers.openapi import openapi_ns
from controllers.openapi._contract import accepts, returns
@ -28,7 +30,6 @@ from controllers.openapi.auth.composition import auth_router
from controllers.openapi.auth.data import AuthData, CallerKind, RBACRequirement
from controllers.service_api.app.error import AppUnavailableError
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
from extensions.ext_database import db
from libs.oauth_bearer import Scope, TokenType
from models import App
from models.enums import AppStatus
@ -56,7 +57,7 @@ _EMPTY_PARAMETERS: dict[str, Any] = {
class AppReadResource(Resource):
"""Base for per-app read endpoints; subclasses call `_load()` for membership/exists checks."""
def _load(self, app_id: str, workspace_id: str | None = None) -> App:
def _load(self, session: Session, app_id: str, workspace_id: str | None = None) -> App:
try:
parsed_uuid = _uuid.UUID(app_id)
is_uuid = True
@ -66,13 +67,13 @@ class AppReadResource(Resource):
if is_uuid:
# ``str(parsed_uuid)`` normalises to the canonical dashed form.
app = AppService.get_visible_app_by_id(str(parsed_uuid), session=db.session())
app = AppService.get_visible_app_by_id(str(parsed_uuid), session)
if app is None:
raise NotFound("app not found")
else:
if not workspace_id:
raise UnprocessableEntity("workspace_id is required for name-based lookup")
matches = AppService.find_visible_apps_by_name(name=app_id, tenant_id=workspace_id, session=db.session())
matches = AppService.find_visible_apps_by_name(session, name=app_id, tenant_id=workspace_id)
if len(matches) == 0:
raise NotFound("app not found")
if len(matches) > 1:
@ -86,14 +87,14 @@ class AppReadResource(Resource):
return app
def parameters_payload(app: App) -> dict:
def parameters_payload(app: App, *, session: Session) -> dict:
"""Mirrors service_api/app/app.py::AppParameterApi response body."""
features_dict, user_input_form = resolve_app_config(app)
features_dict, user_input_form = resolve_app_config(app, session=session)
parameters = get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form)
return Parameters.model_validate(parameters).model_dump(mode="json")
def build_app_describe_response(app: App, fields: set[str] | None) -> AppDescribeResponse:
def build_app_describe_response(app: App, fields: set[str] | None, *, session: Session) -> AppDescribeResponse:
"""Public projection of an app (name / params / input schema) — never internal config."""
want_info = fields is None or "info" in fields
want_params = fields is None or "parameters" in fields
@ -117,12 +118,12 @@ def build_app_describe_response(app: App, fields: set[str] | None) -> AppDescrib
input_schema: dict[str, Any] | None = None
if want_params:
try:
parameters = parameters_payload(app)
parameters = parameters_payload(app, session=session)
except AppUnavailableError:
parameters = dict(_EMPTY_PARAMETERS)
if want_schema:
try:
input_schema = build_input_schema(app)
input_schema = build_input_schema(app, session=session)
except AppUnavailableError:
input_schema = dict(EMPTY_INPUT_SCHEMA)
@ -138,10 +139,11 @@ class AppDescribeApi(AppReadResource):
)
@returns(200, AppDescribeResponse, description="App description")
@accepts(query=AppDescribeQuery)
def get(self, app_id: str, *, auth_data: AuthData, query: AppDescribeQuery):
@with_session(write=False)
def get(self, session: Session, app_id: str, *, auth_data: AuthData, query: AppDescribeQuery):
# describe is UUID-only (workspace_id query param dropped in #37212).
app = self._load(app_id)
return build_app_describe_response(app, query.fields)
app = self._load(session, app_id)
return build_app_describe_response(app, query.fields, session=session)
@openapi_ns.route("/apps")
@ -149,7 +151,8 @@ class AppListApi(Resource):
@auth_router.guard_workspace(scope=Scope.APPS_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, AppListResponse, description="App list")
@accepts(query=AppListQuery)
def get(self, *, auth_data: AuthData, query: AppListQuery):
@with_session(write=False)
def get(self, session: Session, *, auth_data: AuthData, query: AppListQuery):
workspace_id = query.workspace_id
empty = AppListResponse(page=query.page, limit=query.limit, total=0, has_more=False, data=[])
@ -173,11 +176,15 @@ class AppListApi(Resource):
)
access_filter = AppAccessFilter.unrestricted()
if apply_rbac_filter:
access_filter = resolve_app_access_filter(workspace_id, str(auth_data.account_id))
access_filter = resolve_app_access_filter(
workspace_id,
str(auth_data.account_id),
session=session,
)
tenant_name: str | None = None
if parsed_uuid is not None:
app: App | None = AppService.get_visible_app_by_id(str(parsed_uuid), session=db.session())
app: App | None = AppService.get_visible_app_by_id(str(parsed_uuid), session)
if app is None or str(app.tenant_id) != workspace_id:
return empty
if not _is_listable(app):
@ -188,7 +195,7 @@ class AppListApi(Resource):
str(app.id), str(app.maintainer) if app.maintainer else None, str(auth_data.account_id)
):
return empty
tenant_name = TenantService.get_tenant_name(workspace_id, session=db.session())
tenant_name = TenantService.get_tenant_name(workspace_id, session=session)
item = AppListRow(
id=str(app.id),
name=app.name,
@ -215,13 +222,13 @@ class AppListApi(Resource):
if apply_rbac_filter:
access_filter.apply_to_params(params)
pagination = AppService().get_paginate_apps(str(auth_data.account_id), workspace_id, params, db.session())
pagination = AppService().get_paginate_apps(str(auth_data.account_id), workspace_id, params, session)
if pagination is None:
return empty
tenant_name = None
if pagination.items:
tenant_name = TenantService.get_tenant_name(workspace_id, session=db.session())
tenant_name = TenantService.get_tenant_name(workspace_id, session=session)
items = [
AppListRow(

15
api/controllers/openapi/apps_permitted_external.py
View File

@ -8,8 +8,10 @@ EE blueprint chain so this module is unreachable there.
from __future__ import annotations
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.session import with_session
from controllers.openapi import openapi_ns
from controllers.openapi._contract import accepts, returns
from controllers.openapi._models import (
@ -22,7 +24,6 @@ from controllers.openapi._models import (
from controllers.openapi.apps import build_app_describe_response
from controllers.openapi.auth.composition import auth_router
from controllers.openapi.auth.data import AuthData, Edition
from extensions.ext_database import db
from libs.oauth_bearer import Scope, TokenType
from models import App
from models.enums import AppStatus
@ -40,7 +41,8 @@ class PermittedExternalAppsListApi(Resource):
)
@returns(200, PermittedExternalAppsListResponse, description="Permitted external apps list")
@accepts(query=PermittedExternalAppsListQuery)
def get(self, *, auth_data: AuthData, query: PermittedExternalAppsListQuery):
@with_session(write=False)
def get(self, session: Session, *, auth_data: AuthData, query: PermittedExternalAppsListQuery):
page_result = list_permitted_apps(
page=query.page,
limit=query.limit,
@ -55,10 +57,10 @@ class PermittedExternalAppsListApi(Resource):
return env
apps_by_id: dict[str, App] = {
str(a.id): a for a in AppService.find_visible_apps_by_ids(page_result.app_ids, session=db.session())
str(a.id): a for a in AppService.find_visible_apps_by_ids(page_result.app_ids, session)
}
tenant_ids = list({str(a.tenant_id) for a in apps_by_id.values()})
tenants_by_id = {str(t.id): t for t in TenantService.get_tenants_by_ids(tenant_ids, session=db.session())}
tenants_by_id = {str(t.id): t for t in TenantService.get_tenants_by_ids(tenant_ids, session=session)}
items: list[AppListRow] = []
for app_id in page_result.app_ids:
@ -96,9 +98,10 @@ class PermittedExternalAppDescribeApi(Resource):
)
@returns(200, AppDescribeResponse, description="Permitted external app description")
@accepts(query=AppDescribeQuery)
def get(self, app_id: str, *, auth_data: AuthData, query: AppDescribeQuery):
@with_session(write=False)
def get(self, session: Session, app_id: str, *, auth_data: AuthData, query: AppDescribeQuery):
# App already loaded and ACL-checked by the external_sso pipeline; project it.
app = auth_data.app
if app is None:
raise NotFound("app not found")
return build_app_describe_response(app, query.fields)
return build_app_describe_response(app, query.fields, session=session)

25
api/controllers/openapi/auth/prepare.py
View File

@ -6,7 +6,7 @@ from flask import request
from werkzeug.exceptions import Forbidden, InternalServerError, NotFound, Unauthorized
from controllers.openapi.auth.data import AuthData, CallerKind
from extensions.ext_database import db
from core.db.session_factory import session_factory
from models.account import AccountStatus, TenantStatus
from models.enums import AppStatus, EndUserType
from services.account_service import AccountService, TenantService
@ -23,7 +23,8 @@ def load_app(data: AuthData) -> None:
uuid.UUID(app_id)
except ValueError:
raise NotFound("app not found")
app = AppService.get_app_by_id(app_id, session=db.session())
with session_factory.create_session() as session:
app = AppService.get_app_by_id(app_id, session)
if not app or app.status != AppStatus.NORMAL:
raise NotFound("app not found")
data.app = app
@ -34,7 +35,8 @@ def load_tenant(data: AuthData) -> None:
return
if data.app is None:
raise InternalServerError("pipeline_invariant_violated: app not loaded before load_tenant")
tenant = TenantService.get_tenant_by_id(str(data.app.tenant_id), session=db.session())
with session_factory.create_session() as session:
tenant = TenantService.get_tenant_by_id(str(data.app.tenant_id), session=session)
if tenant is None or tenant.status == TenantStatus.ARCHIVE:
raise Forbidden("workspace unavailable")
data.tenant = tenant
@ -50,7 +52,8 @@ def load_tenant_from_request(data: AuthData) -> None:
uuid.UUID(workspace_id)
except ValueError:
raise NotFound("workspace not found")
tenant = TenantService.get_tenant_by_id(workspace_id, session=db.session())
with session_factory.create_session() as session:
tenant = TenantService.get_tenant_by_id(workspace_id, session=session)
if tenant is None or tenant.status == TenantStatus.ARCHIVE:
raise NotFound("workspace not found")
data.tenant = tenant
@ -59,11 +62,12 @@ def load_tenant_from_request(data: AuthData) -> None:
def load_account(data: AuthData) -> None:
if data.caller is not None:
return
account = AccountService.get_account_by_id(str(data.account_id), session=db.session())
if account is None:
raise Unauthorized("account not found")
if data.tenant:
account.current_tenant = data.tenant
with session_factory.create_session() as session:
account = AccountService.get_account_by_id(str(data.account_id), session=session)
if account is None:
raise Unauthorized("account not found")
if data.tenant:
account.set_current_tenant_with_session(data.tenant, session=session)
data.caller = account
data.caller_kind = CallerKind.ACCOUNT
@ -75,7 +79,8 @@ def load_workspace_role(data: AuthData) -> None:
return
if data.caller is not None and getattr(data.caller, "status", None) != AccountStatus.ACTIVE:
return
role = TenantService.get_account_role_in_tenant(str(data.account_id), str(data.tenant.id), session=db.session())
with session_factory.create_session() as session:
role = TenantService.get_account_role_in_tenant(str(data.account_id), str(data.tenant.id), session=session)
if role is None:
return
data.tenant_role = role

78
api/controllers/openapi/workspaces.py
View File

@ -15,9 +15,11 @@ from itertools import starmap
from urllib import parse
from flask_restx import Resource
from sqlalchemy.orm import Session
from werkzeug.exceptions import BadRequest, NotFound
from configs import dify_config
from controllers.common.session import with_session
from controllers.openapi import openapi_ns
from controllers.openapi._contract import accepts, returns
from controllers.openapi._errors import MemberLicenseExceeded, MemberLimitExceeded
@ -35,7 +37,6 @@ from controllers.openapi._models import (
)
from controllers.openapi.auth.composition import auth_router
from controllers.openapi.auth.data import AuthData
from extensions.ext_database import db
from libs.oauth_bearer import Scope, TokenType
from models import Account, Tenant, TenantAccountJoin
from models.account import TenantAccountRole, TenantStatus
@ -64,15 +65,15 @@ def _member_response(account: Account) -> MemberResponse:
)
def _load_tenant(workspace_id: str) -> Tenant:
tenant = TenantService.get_tenant_by_id(workspace_id, session=db.session())
def _load_tenant(session: Session, workspace_id: str) -> Tenant:
tenant = TenantService.get_tenant_by_id(workspace_id, session=session)
if tenant is None or tenant.status != TenantStatus.NORMAL:
raise NotFound("workspace not found")
return tenant
def _load_account(account_id: object) -> Account:
account = AccountService.get_account_by_id(str(account_id), session=db.session()) if account_id else None
def _load_account(session: Session, account_id: object) -> Account:
account = AccountService.get_account_by_id(str(account_id), session=session) if account_id else None
if account is None:
raise RuntimeError("authenticated account_id has no Account row")
return account
@ -94,8 +95,9 @@ def _check_member_invite_quota(tenant_id: str) -> None:
class WorkspacesApi(Resource):
@auth_router.guard(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, WorkspaceListResponse, description="Workspace list")
def get(self, *, auth_data: AuthData):
rows = TenantService.get_workspaces_for_account(str(auth_data.account_id), session=db.session())
@with_session(write=False)
def get(self, session: Session, *, auth_data: AuthData):
rows = TenantService.get_workspaces_for_account(str(auth_data.account_id), session=session)
return WorkspaceListResponse(workspaces=list(starmap(_workspace_summary, rows)))
@ -104,8 +106,9 @@ class WorkspacesApi(Resource):
class WorkspaceByIdApi(Resource):
@auth_router.guard(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, WorkspaceDetailResponse, description="Workspace detail")
def get(self, workspace_id: str, *, auth_data: AuthData):
row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=db.session())
@with_session(write=False)
def get(self, session: Session, workspace_id: str, *, auth_data: AuthData):
row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=session)
# 404 (not 403) on non-member so workspace IDs don't leak across tenants.
if row is None:
raise NotFound("workspace not found")
@ -125,15 +128,16 @@ class WorkspaceSwitchApi(Resource):
@auth_router.guard_workspace(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, WorkspaceDetailResponse, description="Workspace detail")
def post(self, workspace_id: str, *, auth_data: AuthData):
account = _load_account(auth_data.account_id)
@with_session
def post(self, session: Session, workspace_id: str, *, auth_data: AuthData):
account = _load_account(session, auth_data.account_id)
try:
TenantService.switch_tenant(account, workspace_id, session=db.session())
TenantService.switch_tenant(account, workspace_id, session=session)
except AccountNotLinkTenantError:
raise NotFound("workspace not found")
row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=db.session())
row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=session)
if row is None:
raise NotFound("workspace not found")
tenant, membership = row
@ -151,9 +155,10 @@ class WorkspaceMembersApi(Resource):
@auth_router.guard_workspace(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
@returns(200, MemberListResponse, description="Member list")
@accepts(query=MemberListQuery)
def get(self, workspace_id: str, *, auth_data: AuthData, query: MemberListQuery):
tenant = _load_tenant(workspace_id)
members = TenantService.get_tenant_members(tenant, session=db.session())
@with_session(write=False)
def get(self, session: Session, workspace_id: str, *, auth_data: AuthData, query: MemberListQuery):
tenant = _load_tenant(session, workspace_id)
members = TenantService.get_tenant_members(tenant, session=session)
total = len(members)
start = (query.page - 1) * query.limit
page_items = members[start : start + query.limit]
@ -172,9 +177,10 @@ class WorkspaceMembersApi(Resource):
)
@returns(201, MemberInviteResponse, description="Member invited")
@accepts(body=MemberInvitePayload)
def post(self, workspace_id: str, *, auth_data: AuthData, body: MemberInvitePayload):
inviter = _load_account(auth_data.account_id)
tenant = _load_tenant(workspace_id)
@with_session
def post(self, session: Session, workspace_id: str, *, auth_data: AuthData, body: MemberInvitePayload):
inviter = _load_account(session, auth_data.account_id)
tenant = _load_tenant(session, workspace_id)
_check_member_invite_quota(str(tenant.id))
@ -185,7 +191,7 @@ class WorkspaceMembersApi(Resource):
language=None,
role=body.role,
inviter=inviter,
session=db.session(),
session=session,
)
except AccountAlreadyInTenantError as exc:
raise BadRequest(str(exc))
@ -197,7 +203,7 @@ class WorkspaceMembersApi(Resource):
raise BadRequest(str(exc))
normalized_email = body.email.lower()
member = AccountService.get_account_by_email_with_case_fallback(normalized_email, session=db.session())
member = AccountService.get_account_by_email_with_case_fallback(normalized_email, session=session)
if member is None:
# invite_new_member just created or fetched this account.
raise RuntimeError("invited member missing from DB after invite")
@ -229,15 +235,16 @@ class WorkspaceMemberApi(Resource):
allowed_roles=frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN}),
)
@returns(200, MemberActionResponse, description="Member removed")
def delete(self, workspace_id: str, member_id: str, *, auth_data: AuthData):
operator = _load_account(auth_data.account_id)
tenant = _load_tenant(workspace_id)
member = AccountService.get_account_by_id(member_id, session=db.session())
@with_session
def delete(self, session: Session, workspace_id: str, member_id: str, *, auth_data: AuthData):
operator = _load_account(session, auth_data.account_id)
tenant = _load_tenant(session, workspace_id)
member = AccountService.get_account_by_id(member_id, session=session)
if member is None:
raise NotFound("member not found")
try:
TenantService.remove_member_from_tenant(tenant, member, operator, session=db.session())
TenantService.remove_member_from_tenant(tenant, member, operator, session=session)
except CannotOperateSelfError as exc:
raise BadRequest(str(exc))
except NoPermissionError as exc:
@ -254,15 +261,24 @@ class WorkspaceMemberApi(Resource):
)
@returns(200, MemberActionResponse, description="Role updated")
@accepts(body=MemberRoleUpdatePayload)
def patch(self, workspace_id: str, member_id: str, *, auth_data: AuthData, body: MemberRoleUpdatePayload):
operator = _load_account(auth_data.account_id)
tenant = _load_tenant(workspace_id)
member = AccountService.get_account_by_id(member_id, session=db.session())
@with_session
def patch(
self,
session: Session,
workspace_id: str,
member_id: str,
*,
auth_data: AuthData,
body: MemberRoleUpdatePayload,
):
operator = _load_account(session, auth_data.account_id)
tenant = _load_tenant(session, workspace_id)
member = AccountService.get_account_by_id(member_id, session=session)
if member is None:
raise NotFound("member not found")
try:
TenantService.update_member_role(tenant, member, body.role, operator, session=db.session())
TenantService.update_member_role(tenant, member, body.role, operator, session=session)
except CannotOperateSelfError as exc:
raise BadRequest(str(exc))
except NoPermissionError as exc:

25
api/controllers/service_api/app/annotation.py
View File

@ -5,12 +5,13 @@ from flask import request
from flask_restx import Resource
from flask_restx.api import HTTPStatus
from pydantic import BaseModel, Field, TypeAdapter
from sqlalchemy.orm import Session
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console.wraps import edit_permission_required
from controllers.service_api import service_api_ns
from controllers.service_api.wraps import validate_app_token
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from fields.annotation_fields import (
Annotation,
@ -205,12 +206,13 @@ class AnnotationListApi(Resource):
service_api_ns.models[AnnotationList.__name__],
)
@validate_app_token
def get(self, app_model: App):
@with_session(write=False)
def get(self, session: Session, app_model: App):
"""List annotations for the application."""
query = AnnotationListQuery.model_validate(request.args.to_dict(flat=True))
annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id(
app_model.id, query.page, query.limit, query.keyword, session=db.session()
app_model.id, query.page, query.limit, query.keyword, session
)
annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True)
return AnnotationList(
@ -247,13 +249,12 @@ class AnnotationListApi(Resource):
service_api_ns.models[Annotation.__name__],
)
@validate_app_token
def post(self, app_model: App):
@with_session
def post(self, session: Session, app_model: App):
"""Create a new annotation."""
payload = AnnotationCreatePayload.model_validate(service_api_ns.payload or {})
insert_args: InsertAnnotationArgs = {"question": payload.question, "answer": payload.answer}
annotation = AppAnnotationService.insert_app_annotation_directly(
insert_args, app_model.id, session=db.session()
)
annotation = AppAnnotationService.insert_app_annotation_directly(insert_args, app_model.id, session)
return dump_response(Annotation, annotation), HTTPStatus.CREATED
@ -287,14 +288,15 @@ class AnnotationUpdateDeleteApi(Resource):
service_api_ns.models[Annotation.__name__],
)
@validate_app_token
@with_session
@edit_permission_required
def put(self, app_model: App, annotation_id: UUID):
def put(self, session: Session, app_model: App, annotation_id: UUID):
"""Update an existing annotation."""
payload = AnnotationCreatePayload.model_validate(service_api_ns.payload or {})
update_args: UpdateAnnotationArgs = {"question": payload.question, "answer": payload.answer}
app_ref = AppRefService.create_app_ref(app_model)
annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id))
annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, db.session())
annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, session)
return dump_response(Annotation, annotation)
@service_api_ns.doc(
@ -319,10 +321,11 @@ class AnnotationUpdateDeleteApi(Resource):
}
)
@validate_app_token
@with_session
@edit_permission_required
def delete(self, app_model: App, annotation_id: UUID):
def delete(self, session: Session, app_model: App, annotation_id: UUID):
"""Delete an annotation."""
app_ref = AppRefService.create_app_ref(app_model)
annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id))
AppAnnotationService.delete_app_annotation(annotation_ref, db.session())
AppAnnotationService.delete_app_annotation(annotation_ref, session)
return "", 204

27
api/controllers/service_api/app/app.py
View File

@ -2,6 +2,7 @@ from typing import Any, cast
from flask_restx import Resource
from pydantic import Field
from sqlalchemy.orm import Session
from controllers.common.agent_app_parameters import get_published_agent_app_feature_dict_and_user_input_form
from controllers.common.fields import Parameters
@ -13,7 +14,7 @@ from core.app.app_config.common.parameters_mapping import get_parameters_from_fe
from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError
from extensions.ext_database import db
from fields.base import ResponseModel
from models.model import App, AppMode
from models.model import App, AppMode, load_annotation_reply_config
from services.app_service import AppService
@ -32,9 +33,13 @@ class AppMetaResponse(ResponseModel):
register_response_schema_models(service_api_ns, Parameters, AppMetaResponse, AppInfoResponse)
def _get_agent_app_feature_dict_and_user_input_form(app_model: App) -> tuple[dict[str, Any], list[dict[str, Any]]]:
def _get_agent_app_feature_dict_and_user_input_form(
app_model: App,
*,
session: Session,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
try:
return get_published_agent_app_feature_dict_and_user_input_form(app_model)
return get_published_agent_app_feature_dict_and_user_input_form(app_model, session=session)
except AgentAppNotPublishedError:
raise AgentNotPublishedError()
except AgentAppGeneratorError:
@ -73,23 +78,31 @@ class AppParameterApi(Resource):
Returns the input form parameters and configuration for the application.
"""
session = db.session()
features_dict: dict[str, Any]
user_input_form: list[dict[str, Any]]
if app_model.mode == AppMode.AGENT:
features_dict, user_input_form = _get_agent_app_feature_dict_and_user_input_form(app_model)
features_dict, user_input_form = _get_agent_app_feature_dict_and_user_input_form(
app_model,
session=session,
)
elif app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
workflow = app_model.workflow
workflow = app_model.workflow_with_session(session=session)
if workflow is None:
raise AppUnavailableError()
features_dict = workflow.features_dict
user_input_form = workflow.user_input_form(to_old_structure=True)
else:
app_model_config = app_model.app_model_config
app_model_config = app_model.app_model_config_with_session(session=session)
if app_model_config is None:
raise AppUnavailableError()
features_dict = cast(dict[str, Any], app_model_config.to_dict())
annotation_reply = load_annotation_reply_config(session, app_model.id)
features_dict = cast(
dict[str, Any],
app_model_config.to_dict(annotation_reply=annotation_reply),
)
user_input_form = features_dict.get("user_input_form", [])

7
api/controllers/service_api/app/audio.py
View File

@ -104,7 +104,12 @@ class AudioApi(Resource):
file = request.files["file"]
try:
response = AudioService.transcript_asr(app_model=app_model, file=file, end_user=end_user.id)
response = AudioService.transcript_asr(
app_model=app_model,
file=file,
session=db.session(),
end_user=end_user.id,
)
return dump_response(AudioTranscriptResponse, response)
except services.errors.app_model_config.AppModelConfigBrokenError:

3
api/controllers/service_api/app/completion.py
View File

@ -43,7 +43,6 @@ from core.errors.error import (
)
from core.helper.trace_id_helper import get_external_trace_id, get_trace_session_id, omit_trace_session_id_from_payload
from enums.cloud_plan import CloudPlan
from extensions.ext_database import db
from graphon.model_runtime.errors.invoke import InvokeError
from libs import helper
from libs.helper import UUIDStrOrEmpty
@ -399,7 +398,7 @@ class ChatApi(Resource):
app_model=app_model,
conversation_id=payload.conversation_id,
user=end_user,
session=db.session(),
session=session,
)
response = AppGenerateService.generate(

14
api/controllers/service_api/app/conversation.py
View File

@ -21,6 +21,7 @@ from fields._value_type_serializer import serialize_value_type
from fields.base import ResponseModel
from fields.conversation_fields import (
ConversationInfiniteScrollPagination,
ConversationResponseSource,
SimpleConversation,
)
from graphon.variables.types import SegmentType
@ -204,7 +205,13 @@ class ConversationApi(Resource):
sort_by=query_args.sort_by,
)
adapter = TypeAdapter(SimpleConversation)
conversations = [adapter.validate_python(item, from_attributes=True) for item in pagination.data]
conversations = [
adapter.validate_python(
ConversationResponseSource(item, session=session),
from_attributes=True,
)
for item in pagination.data
]
return ConversationInfiniteScrollPagination(
limit=pagination.limit, has_more=pagination.has_more, data=conversations
).model_dump(mode="json")
@ -294,10 +301,11 @@ class ConversationRenameApi(Resource):
payload = ConversationRenamePayload.model_validate(service_api_ns.payload or {})
try:
session = db.session()
conversation = ConversationService.rename(
app_model, conversation_id, end_user, payload.name, payload.auto_generate, session=db.session()
app_model, conversation_id, end_user, payload.name, payload.auto_generate, session=session
)
return dump_response(SimpleConversation, conversation)
return dump_response(SimpleConversation, ConversationResponseSource(conversation, session=session))
except services.errors.conversation.ConversationNotExistsError:
raise NotFound("Conversation Not Exists.")

10
api/controllers/service_api/app/message.py
View File

@ -17,7 +17,7 @@ from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate
from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from fields.base import ResponseModel
from fields.conversation_fields import ResultResponse
from fields.conversation_fields import MessageResponseSource, ResultResponse
from fields.message_fields import MessageInfiniteScrollPagination, MessageListItem
from models.enums import FeedbackRating
from models.model import App, AppMode, EndUser
@ -111,11 +111,15 @@ class MessageListApi(Resource):
first_id = query_args.first_id or None
try:
session = db.session()
pagination = MessageService.pagination_by_first_id(
app_model, end_user, conversation_id, first_id, query_args.limit, session=db.session()
app_model, end_user, conversation_id, first_id, query_args.limit, session=session
)
adapter = TypeAdapter(MessageListItem)
items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
items = [
adapter.validate_python(MessageResponseSource(message, session=session), from_attributes=True)
for message in pagination.data
]
return MessageInfiniteScrollPagination(
limit=pagination.limit, has_more=pagination.has_more, data=items
).model_dump(mode="json")

104
api/controllers/service_api/dataset/dataset.py
View File

@ -24,7 +24,7 @@ from controllers.common.schema import (
register_response_schema_models,
register_schema_models,
)
from controllers.console.app.wraps import with_session
from controllers.common.session import with_session
from controllers.console.wraps import edit_permission_required
from controllers.service_api import service_api_ns
from controllers.service_api.dataset.error import DatasetInUseError, DatasetNameDuplicateError, InvalidActionError
@ -34,9 +34,8 @@ from controllers.service_api.wraps import (
)
from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from extensions.ext_database import db
from fields.base import ResponseModel
from fields.dataset_fields import DatasetDetailResponse
from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from graphon.model_runtime.entities.model_entities import ModelType
from libs.helper import dump_response
from libs.login import current_user
@ -93,8 +92,10 @@ _SERVICE_DATASET_DETAIL_EXCLUDE = {"permission_keys"}
_SERVICE_DATASET_LIST_EXCLUDE = {"data": {"__all__": _SERVICE_DATASET_DETAIL_EXCLUDE}}
def _dump_service_dataset_detail(dataset: Any) -> dict[str, Any]:
return DatasetDetailResponse.model_validate(dataset, from_attributes=True).model_dump(
def _dump_service_dataset_detail(dataset: Any, *, session: Session) -> dict[str, Any]:
return DatasetDetailResponse.model_validate(
dataset_detail_response_source(dataset, session=session), from_attributes=True
).model_dump(
mode="json",
exclude=_SERVICE_DATASET_DETAIL_EXCLUDE,
)
@ -405,7 +406,8 @@ class DatasetListApi(DatasetApiResource):
"Datasets retrieved successfully",
service_api_ns.models[DatasetListResponse.__name__],
)
def get(self, tenant_id):
@with_session(write=False)
def get(self, session: Session, tenant_id):
"""Resource for getting datasets."""
query_params: dict[str, str | list[str]] = dict(request.args.to_dict())
if "tag_ids" in request.args:
@ -416,7 +418,7 @@ class DatasetListApi(DatasetApiResource):
datasets, total = DatasetService.get_datasets(
query.page,
query.limit,
db.session(),
session,
tenant_id,
current_user,
query.keyword,
@ -436,7 +438,7 @@ class DatasetListApi(DatasetApiResource):
for embedding_model in embedding_models:
model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
data = [_dump_service_dataset_detail(dataset) for dataset in datasets]
data = [_dump_service_dataset_detail(dataset, session=session) for dataset in datasets]
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"]))
@ -538,7 +540,7 @@ class DatasetListApi(DatasetApiResource):
)
initialize_created_app_rbac_access_task.delay(tenant_id, current_user.id, dataset_id=dataset.id)
return _dump_service_dataset_detail(dataset), 200
return _dump_service_dataset_detail(dataset, session=session), 200
@service_api_ns.route("/datasets/<uuid:dataset_id>")
@ -574,16 +576,17 @@ class DatasetApi(DatasetApiResource):
"Dataset retrieved successfully",
service_api_ns.models[DatasetDetailWithPartialMembersResponse.__name__],
)
def get(self, _, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, _, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
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)
data = _dump_service_dataset_detail(dataset, session=session)
# check embedding setting
assert isinstance(current_user, Account)
cid = current_user.current_tenant_id
@ -612,7 +615,7 @@ class DatasetApi(DatasetApiResource):
retrieval_model_dict["search_method"] = "keyword_search"
if data.get("permission") == "partial_members":
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session())
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session)
data.update({"partial_member_list": part_users_list})
return _dump_service_dataset_with_partial_members(data), 200
@ -651,7 +654,7 @@ class DatasetApi(DatasetApiResource):
@with_session
def patch(self, session: Session, _, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
@ -692,7 +695,7 @@ class DatasetApi(DatasetApiResource):
dataset,
str(payload.permission) if payload.permission else None,
payload.partial_member_list,
session=db.session(),
session=session,
)
dataset = DatasetService.update_dataset(dataset_id_str, update_data, current_user, session=session)
@ -700,19 +703,19 @@ class DatasetApi(DatasetApiResource):
if dataset is None:
raise NotFound("Dataset not found.")
result_data = _dump_service_dataset_detail(dataset)
result_data = _dump_service_dataset_detail(dataset, session=session)
assert isinstance(current_user, Account)
tenant_id = current_user.current_tenant_id
if payload.partial_member_list and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM:
DatasetPermissionService.update_partial_member_list(
tenant_id, dataset_id_str, payload.partial_member_list, db.session()
tenant_id, dataset_id_str, payload.partial_member_list, session
)
# clear partial member list when permission is only_me or all_team_members
elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}:
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session())
DatasetPermissionService.clear_partial_member_list(dataset_id_str, session)
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session())
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session)
result_data.update({"partial_member_list": partial_member_list})
return _dump_service_dataset_with_partial_members(result_data), 200
@ -745,7 +748,8 @@ class DatasetApi(DatasetApiResource):
}
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def delete(self, _, dataset_id: UUID):
@with_session
def delete(self, session: Session, _, dataset_id: UUID):
"""
Deletes a dataset given its ID.
@ -765,8 +769,8 @@ class DatasetApi(DatasetApiResource):
dataset_id_str = str(dataset_id)
try:
if DatasetService.delete_dataset(dataset_id_str, current_user, db.session()):
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session())
if DatasetService.delete_dataset(dataset_id_str, current_user, session):
DatasetPermissionService.clear_partial_member_list(dataset_id_str, session)
return "", 204
else:
raise NotFound("Dataset not found.")
@ -812,7 +816,14 @@ class DocumentStatusApi(DatasetApiResource):
}
)
@service_api_ns.expect(service_api_ns.models[DocumentStatusPayload.__name__])
def patch(self, tenant_id, dataset_id: UUID, action: Literal["enable", "disable", "archive", "un_archive"]):
@with_session
def patch(
self,
session: Session,
tenant_id,
dataset_id: UUID,
action: Literal["enable", "disable", "archive", "un_archive"],
):
"""
Batch update document status.
@ -831,14 +842,14 @@ class DocumentStatusApi(DatasetApiResource):
InvalidActionError: If the action is invalid or cannot be performed.
"""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
# Check user's permission
try:
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
@ -850,7 +861,7 @@ class DocumentStatusApi(DatasetApiResource):
document_ids = data.get("document_ids", [])
try:
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session())
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, session)
except services.errors.document.DocumentIndexingError as e:
raise InvalidActionError(str(e))
except ValueError as e:
@ -882,12 +893,13 @@ class DatasetTagsApi(DatasetApiResource):
"Tags retrieved successfully",
service_api_ns.models[KnowledgeTagListResponse.__name__],
)
def get(self, _):
@with_session(write=False)
def get(self, session: Session, _):
"""Get all knowledge type tags."""
assert isinstance(current_user, Account)
cid = current_user.current_tenant_id
assert cid is not None
tags = TagService.get_tags("knowledge", cid, session=db.session())
tags = TagService.get_tags("knowledge", cid, session=session)
return dump_response(KnowledgeTagListResponse, tags), 200
@service_api_ns.doc(
@ -913,14 +925,15 @@ class DatasetTagsApi(DatasetApiResource):
"Tag created successfully",
service_api_ns.models[KnowledgeTagResponse.__name__],
)
def post(self, _):
@with_session
def post(self, session: Session, _):
"""Add a knowledge type tag."""
assert isinstance(current_user, Account)
if not (current_user.has_edit_permission or current_user.is_dataset_editor):
raise Forbidden()
payload = TagCreatePayload.model_validate(service_api_ns.payload or {})
tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=TagType.KNOWLEDGE), db.session())
tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=TagType.KNOWLEDGE), session)
response = KnowledgeTagResponse(id=tag.id, name=tag.name, type=tag.type, binding_count="0")
return response.model_dump(mode="json"), 200
@ -948,7 +961,8 @@ class DatasetTagsApi(DatasetApiResource):
"Tag updated successfully",
service_api_ns.models[KnowledgeTagResponse.__name__],
)
def patch(self, _):
@with_session
def patch(self, session: Session, _):
assert isinstance(current_user, Account)
if not (current_user.has_edit_permission or current_user.is_dataset_editor):
raise Forbidden()
@ -956,10 +970,10 @@ class DatasetTagsApi(DatasetApiResource):
payload = TagUpdatePayload.model_validate(service_api_ns.payload or {})
tag_id = payload.tag_id
tag = TagService.update_tags(
UpdateTagServicePayload(name=payload.name), tag_id, db.session(), tag_type=TagType.KNOWLEDGE
UpdateTagServicePayload(name=payload.name), tag_id, session, tag_type=TagType.KNOWLEDGE
)
binding_count = TagService.get_tag_binding_count(tag_id, db.session(), tag_type=TagType.KNOWLEDGE)
binding_count = TagService.get_tag_binding_count(tag_id, session, tag_type=TagType.KNOWLEDGE)
response = KnowledgeTagResponse(id=tag.id, name=tag.name, type=tag.type, binding_count=str(binding_count))
return response.model_dump(mode="json"), 200
@ -983,10 +997,11 @@ class DatasetTagsApi(DatasetApiResource):
}
)
@edit_permission_required
def delete(self, _):
@with_session
def delete(self, session: Session, _):
"""Delete a knowledge type tag."""
payload = TagDeletePayload.model_validate(service_api_ns.payload or {})
TagService.delete_tag(payload.tag_id, db.session(), tag_type=TagType.KNOWLEDGE)
TagService.delete_tag(payload.tag_id, session, tag_type=TagType.KNOWLEDGE)
return "", 204
@ -1011,7 +1026,8 @@ class DatasetTagBindingApi(DatasetApiResource):
403: "Forbidden - insufficient permissions",
}
)
def post(self, _):
@with_session
def post(self, session: Session, _):
# The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
assert isinstance(current_user, Account)
if not (current_user.has_edit_permission or current_user.is_dataset_editor):
@ -1020,7 +1036,7 @@ class DatasetTagBindingApi(DatasetApiResource):
payload = TagBindingPayload.model_validate(service_api_ns.payload or {})
TagService.save_tag_binding(
TagBindingCreatePayload(tag_ids=payload.tag_ids, target_id=payload.target_id, type=TagType.KNOWLEDGE),
db.session(),
session,
)
return "", 204
@ -1046,7 +1062,8 @@ class DatasetTagUnbindingApi(DatasetApiResource):
403: "Forbidden - insufficient permissions",
}
)
def post(self, _):
@with_session
def post(self, session: Session, _):
# The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
assert isinstance(current_user, Account)
if not (current_user.has_edit_permission or current_user.is_dataset_editor):
@ -1055,7 +1072,7 @@ class DatasetTagUnbindingApi(DatasetApiResource):
payload = TagUnbindingPayload.model_validate(service_api_ns.payload or {})
TagService.delete_tag_binding(
TagBindingDeletePayload(tag_ids=payload.tag_ids, target_id=payload.target_id, type=TagType.KNOWLEDGE),
db.session(),
session,
)
return "", 204
@ -1085,14 +1102,13 @@ class DatasetTagsBindingStatusApi(DatasetApiResource):
"Tags retrieved successfully",
service_api_ns.models[DatasetBoundTagListResponse.__name__],
)
def get(self, _, *args, **kwargs):
@with_session(write=False)
def get(self, session: Session, _, *args, **kwargs):
"""Get all knowledge type tags."""
dataset_id = kwargs.get("dataset_id")
assert isinstance(current_user, Account)
assert current_user.current_tenant_id is not None
tags = TagService.get_tags_by_target_id(
"knowledge", current_user.current_tenant_id, str(dataset_id), db.session()
)
tags = TagService.get_tags_by_target_id("knowledge", current_user.current_tenant_id, str(dataset_id), session)
response = DatasetBoundTagListResponse(
data=[DatasetBoundTagResponse(id=tag.id, name=tag.name) for tag in tags],
total=len(tags),

212
api/controllers/service_api/dataset/document.py
View File

@ -6,6 +6,7 @@ deprecated in generated API docs so clients migrate toward the canonical paths.
"""
import json
from collections.abc import Mapping
from contextlib import ExitStack
from copy import deepcopy
from typing import Annotated, Any, Literal, Self, override
@ -23,6 +24,7 @@ from pydantic import (
)
from pydantic.json_schema import SkipJsonSchema
from sqlalchemy import desc, func, select
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, NotFound
import services
@ -42,6 +44,7 @@ from controllers.common.schema import (
register_response_schema_models,
register_schema_models,
)
from controllers.common.session import with_session
from controllers.service_api import service_api_ns
from controllers.service_api.app.error import ProviderNotInitializeError
from controllers.service_api.dataset.error import (
@ -65,6 +68,8 @@ from fields.document_fields import (
DocumentMetadataResponse,
DocumentResponse,
DocumentStatusListResponse,
document_response,
document_responses,
normalize_enum,
)
from libs.helper import dump_response
@ -291,6 +296,13 @@ class DocumentAndBatchResponse(ResponseModel):
batch: str
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},
)
# Use SkipJsonSchema to support 3 metadata modes
class DocumentDetailResponse(ResponseModel):
id: str
@ -357,14 +369,14 @@ register_response_schema_models(
)
def _create_document_by_text(tenant_id: str, dataset_id: UUID) -> tuple[Document, str]:
def _create_document_by_text(session: Session, tenant_id: str, dataset_id: UUID) -> tuple[Document, str]:
"""Create a document from text for both canonical and legacy routes."""
payload = DocumentTextCreatePayload.model_validate(service_api_ns.payload or {})
args = payload.model_dump(exclude_none=True)
dataset_id_str = str(dataset_id)
tenant_id_str = tenant_id
dataset = db.session.scalar(
tenant_id_str = str(tenant_id)
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id_str, Dataset.id == dataset_id_str).limit(1)
)
@ -414,9 +426,11 @@ def _create_document_by_text(tenant_id: str, dataset_id: UUID) -> tuple[Document
dataset=dataset,
knowledge_config=knowledge_config,
account=current_user,
dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
dataset_process_rule=dataset.get_latest_process_rule(session=session)
if "process_rule" not in args
else None,
created_from="api",
session=db.session(),
session=session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@ -425,10 +439,12 @@ def _create_document_by_text(tenant_id: str, dataset_id: UUID) -> tuple[Document
return document, batch
def _update_document_by_text(tenant_id: str, dataset_id: UUID, document_id: UUID) -> tuple[Document, str]:
def _update_document_by_text(
session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID
) -> tuple[Document, str]:
"""Update a document from text for both canonical and legacy routes."""
payload = DocumentTextUpdate.model_validate(service_api_ns.payload or {})
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == str(dataset_id)).limit(1)
)
args = payload.model_dump(exclude_none=True)
@ -474,9 +490,11 @@ def _update_document_by_text(tenant_id: str, dataset_id: UUID, document_id: UUID
dataset=dataset,
knowledge_config=knowledge_config,
account=current_user,
dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
dataset_process_rule=dataset.get_latest_process_rule(session=session)
if "process_rule" not in args
else None,
created_from="api",
session=db.session(),
session=session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@ -524,10 +542,11 @@ class DocumentAddByTextApi(DatasetApiResource):
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_resource_check("documents", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID):
"""Create document by text."""
document, batch = _create_document_by_text(tenant_id=tenant_id, dataset_id=dataset_id)
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
document, batch = _create_document_by_text(session=session, tenant_id=tenant_id, dataset_id=dataset_id)
return _document_and_batch_response(document, batch, session=session), 200
@service_api_ns.route("/datasets/<uuid:dataset_id>/document/create_by_text")
@ -557,10 +576,11 @@ class DeprecatedDocumentAddByTextApi(DatasetApiResource):
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_resource_check("documents", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID):
"""Create document by text through the deprecated underscore alias."""
document, batch = _create_document_by_text(tenant_id=tenant_id, dataset_id=dataset_id)
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
document, batch = _create_document_by_text(session=session, tenant_id=tenant_id, dataset_id=dataset_id)
return _document_and_batch_response(document, batch, session=session), 200
@service_api_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update-by-text")
@ -602,10 +622,13 @@ class DocumentUpdateByTextApi(DatasetApiResource):
)
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID):
"""Update document by text."""
document, batch = _update_document_by_text(tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id)
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
document, batch = _update_document_by_text(
session=session, tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id
)
return _document_and_batch_response(document, batch, session=session), 200
@service_api_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_text")
@ -634,10 +657,13 @@ class DeprecatedDocumentUpdateByTextApi(DatasetApiResource):
)
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID):
"""Update document by text through the deprecated underscore alias."""
document, batch = _update_document_by_text(tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id)
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
document, batch = _update_document_by_text(
session=session, tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id
)
return _document_and_batch_response(document, batch, session=session), 200
@service_api_ns.route(
@ -694,9 +720,10 @@ class DocumentAddByFileApi(DatasetApiResource):
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_resource_check("documents", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id, dataset_id: UUID):
@with_session
def post(self, session: Session, tenant_id, dataset_id: UUID):
"""Create document by upload file."""
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).limit(1)
)
@ -767,7 +794,7 @@ class DocumentAddByFileApi(DatasetApiResource):
knowledge_config = KnowledgeConfig.model_validate(args)
DocumentService.document_create_args_validate(knowledge_config)
dataset_process_rule = dataset.latest_process_rule if "process_rule" not in args else None
dataset_process_rule = dataset.get_latest_process_rule(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.")
@ -775,23 +802,24 @@ class DocumentAddByFileApi(DatasetApiResource):
documents, batch = DocumentService.save_document_with_dataset_id(
dataset=dataset,
knowledge_config=knowledge_config,
account=dataset.created_by_account,
account=dataset.get_created_by_account(session=session),
dataset_process_rule=dataset_process_rule,
created_from="api",
session=db.session(),
session=session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
document = documents[0]
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
return _document_and_batch_response(document, batch, session=session), 200
def _update_document_by_file(tenant_id: str, dataset_id: UUID, document_id: UUID) -> tuple[Document, str]:
def _update_document_by_file(
session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID
) -> tuple[Document, str]:
"""Update a document from an uploaded file for canonical and deprecated routes."""
dataset_id_str = str(dataset_id)
tenant_id_str = tenant_id
dataset = db.session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id_str, Dataset.id == dataset_id_str).limit(1)
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -852,10 +880,12 @@ def _update_document_by_file(tenant_id: str, dataset_id: UUID, document_id: UUID
documents, _ = DocumentService.save_document_with_dataset_id(
dataset=dataset,
knowledge_config=knowledge_config,
account=dataset.created_by_account,
dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
account=dataset.get_created_by_account(session=session),
dataset_process_rule=dataset.get_latest_process_rule(session=session)
if "process_rule" not in args
else None,
created_from="api",
session=db.session(),
session=session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@ -912,10 +942,13 @@ class DeprecatedDocumentUpdateByFileApi(DatasetApiResource):
)
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID):
"""Update document by file through the deprecated file-update aliases."""
document, batch = _update_document_by_file(tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id)
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
document, batch = _update_document_by_file(
session=session, tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id
)
return _document_and_batch_response(document, batch, session=session), 200
@service_api_ns.route("/datasets/<uuid:dataset_id>/documents")
@ -945,11 +978,12 @@ class DocumentListApi(DatasetApiResource):
@service_api_ns.response(
200, "Documents retrieved successfully", service_api_ns.models[DocumentListResponse.__name__]
)
def get(self, tenant_id, dataset_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id, dataset_id: UUID):
dataset_id_str = str(dataset_id)
tenant_id = str(tenant_id)
query_params = query_params_from_request(DocumentListQuery)
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -967,7 +1001,7 @@ class DocumentListApi(DatasetApiResource):
query = query.order_by(desc(Document.created_at), desc(Document.position))
paginated_documents = paginate_query(
query, page=query_params.page, per_page=query_params.limit, max_per_page=100
query, session=session, page=query_params.page, per_page=query_params.limit, max_per_page=100
)
documents = paginated_documents.items
@ -975,11 +1009,11 @@ class DocumentListApi(DatasetApiResource):
documents=documents,
dataset=dataset,
tenant_id=tenant_id,
session=db.session(),
session=session,
)
response = {
"data": documents,
"data": document_responses(documents, session=session),
"has_more": len(documents) == query_params.limit,
"limit": query_params.limit,
"total": paginated_documents.total,
@ -1020,7 +1054,8 @@ class DocumentBatchDownloadZipApi(DatasetApiResource):
)
@service_api_ns.response(200, "ZIP archive generated successfully")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id, dataset_id: UUID):
@with_session(write=False)
def post(self, session: Session, tenant_id, dataset_id: UUID):
payload = DocumentBatchDownloadZipPayload.model_validate(service_api_ns.payload or {})
upload_files, download_name = DocumentService.prepare_document_batch_download_zip(
@ -1028,7 +1063,7 @@ class DocumentBatchDownloadZipApi(DatasetApiResource):
document_ids=[str(document_id) for document_id in payload.document_ids],
tenant_id=str(tenant_id),
current_user=current_user,
session=db.session(),
session=session,
)
with ExitStack() as stack:
@ -1076,23 +1111,24 @@ class DocumentIndexingStatusApi(DatasetApiResource):
"Indexing status retrieved successfully",
service_api_ns.models[DocumentStatusListResponse.__name__],
)
def get(self, tenant_id, dataset_id: UUID, batch: str):
@with_session(write=False)
def get(self, session: Session, tenant_id, dataset_id: UUID, batch: str):
dataset_id_str = str(dataset_id)
tenant_id = str(tenant_id)
# get dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
raise NotFound("Dataset not found.")
# get documents
documents = DocumentService.get_batch_documents(dataset_id_str, batch, db.session())
documents = DocumentService.get_batch_documents(dataset_id_str, batch, session)
if not documents:
raise NotFound("Documents not found.")
documents_status = []
for document in documents:
completed_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.completed_at.isnot(None),
DocumentSegment.document_id == str(document.id),
@ -1102,7 +1138,7 @@ class DocumentIndexingStatusApi(DatasetApiResource):
or 0
)
total_segments = (
db.session.scalar(
session.scalar(
select(func.count(DocumentSegment.id)).where(
DocumentSegment.document_id == str(document.id),
DocumentSegment.status != SegmentStatus.RE_SEGMENT,
@ -1160,9 +1196,12 @@ class DocumentDownloadApi(DatasetApiResource):
service_api_ns.models[UrlResponse.__name__],
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def get(self, tenant_id, dataset_id: UUID, document_id: UUID):
dataset = self.get_dataset(str(dataset_id), str(tenant_id))
document = DocumentService.get_document(dataset.id, str(document_id), session=db.session())
@with_session(write=False)
def get(self, session: Session, tenant_id, dataset_id: UUID, document_id: UUID):
dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), str(tenant_id), session=session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, str(document_id), session=session)
if not document:
raise NotFound("Document not found.")
@ -1170,9 +1209,7 @@ class DocumentDownloadApi(DatasetApiResource):
if document.tenant_id != str(tenant_id):
raise Forbidden("No permission.")
return UrlResponse(url=DocumentService.get_document_download_url(document, db.session())).model_dump(
mode="json"
)
return UrlResponse(url=DocumentService.get_document_download_url(document, session)).model_dump(mode="json")
@service_api_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
@ -1219,13 +1256,16 @@ class DocumentApi(DatasetApiResource):
"Document retrieved successfully",
service_api_ns.models[DocumentDetailResponse.__name__],
)
def get(self, tenant_id, dataset_id: UUID, document_id: UUID):
@with_session(write=False)
def get(self, session: Session, tenant_id, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = self.get_dataset(dataset_id_str, tenant_id)
dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
@ -1250,17 +1290,23 @@ class DocumentApi(DatasetApiResource):
document_id=document_id_str,
dataset_id=dataset_id_str,
tenant_id=tenant_id,
session=db.session(),
session=session,
)
if metadata == "only":
response_include = {"id", "doc_type", "doc_metadata"}
response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details}
response = {
"id": document.id,
"doc_type": document.doc_type,
"doc_metadata": document.get_doc_metadata_details(session=session),
}
elif metadata == "without":
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session)
response_exclude = {"doc_type", "doc_metadata"}
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session())
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
data_source_info = document.data_source_detail_dict
document_process_rule = document.get_dataset_process_rule(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)
response = {
"id": document.id,
"position": document.position,
@ -1283,9 +1329,9 @@ class DocumentApi(DatasetApiResource):
"disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None,
"disabled_by": document.disabled_by,
"archived": document.archived,
"segment_count": document.segment_count,
"average_segment_length": document.average_segment_length,
"hit_count": document.hit_count,
"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),
"display_status": document.display_status,
"doc_form": document.doc_form,
"doc_language": document.doc_language,
@ -1293,9 +1339,11 @@ class DocumentApi(DatasetApiResource):
"need_summary": document.need_summary if document.need_summary is not None else False,
}
else:
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session())
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
data_source_info = document.data_source_detail_dict
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session)
document_process_rule = document.get_dataset_process_rule(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)
response = {
"id": document.id,
"position": document.position,
@ -1319,10 +1367,10 @@ class DocumentApi(DatasetApiResource):
"disabled_by": document.disabled_by,
"archived": document.archived,
"doc_type": document.doc_type,
"doc_metadata": document.doc_metadata_details,
"segment_count": document.segment_count,
"average_segment_length": document.average_segment_length,
"hit_count": document.hit_count,
"doc_metadata": document.get_doc_metadata_details(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),
"display_status": document.display_status,
"doc_form": document.doc_form,
"doc_language": document.doc_language,
@ -1372,10 +1420,13 @@ class DocumentApi(DatasetApiResource):
)
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def patch(self, tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def patch(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID):
"""Update document by file on the canonical document resource."""
document, batch = _update_document_by_file(tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id)
return dump_response(DocumentAndBatchResponse, {"document": document, "batch": batch}), 200
document, batch = _update_document_by_file(
session=session, tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id
)
return _document_and_batch_response(document, batch, session=session), 200
@service_api_ns.doc(
summary="Delete Document",
@ -1400,21 +1451,22 @@ class DocumentApi(DatasetApiResource):
}
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def delete(self, tenant_id, dataset_id: UUID, document_id: UUID):
@with_session
def delete(self, session: Session, tenant_id, dataset_id: UUID, document_id: UUID):
"""Delete document."""
document_id_str = str(document_id)
dataset_id_str = str(dataset_id)
tenant_id = str(tenant_id)
# get dataset info
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
raise ValueError("Dataset does not exist.")
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
# 404 if document not found
if document is None:
@ -1426,7 +1478,7 @@ class DocumentApi(DatasetApiResource):
try:
# delete document
DocumentService.delete_document(document, db.session())
DocumentService.delete_document(document, session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")

2
api/controllers/service_api/dataset/hit_testing.py
View File

@ -61,7 +61,7 @@ class HitTestingApi(DatasetApiResource, DatasetsHitTestingBase):
Tests retrieval performance for the specified dataset.
"""
dataset_id_str = str(dataset_id)
dataset = self.get_and_validate_dataset(dataset_id_str)
dataset = self.get_and_validate_dataset(session, dataset_id_str)
args = self.parse_args(service_api_ns.payload)
self.hit_testing_args_check(args)

59
api/controllers/service_api/dataset/metadata.py
View File

@ -2,13 +2,14 @@ from typing import Literal
from uuid import UUID
from flask_login import current_user
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.common.controller_schemas import MetadataUpdatePayload
from controllers.common.schema import register_response_schema_models, register_schema_model, register_schema_models
from controllers.common.session import with_session
from controllers.service_api import service_api_ns
from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_rate_limit_check
from extensions.ext_database import db
from fields.dataset_fields import (
DatasetMetadataActionResponse,
DatasetMetadataBuiltInFieldsResponse,
@ -76,17 +77,18 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource):
201, "Metadata created successfully", service_api_ns.models[DatasetMetadataResponse.__name__]
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id, dataset_id: UUID):
@with_session
def post(self, session: Session, tenant_id, dataset_id: UUID):
"""Create metadata for a dataset."""
metadata_args = MetadataArgs.model_validate(service_api_ns.payload or {})
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.create_metadata(dataset_id_str, metadata_args, session=db.session())
metadata = MetadataService.create_metadata(dataset_id_str, metadata_args, session=session)
return dump_response(DatasetMetadataResponse, metadata), 201
@service_api_ns.doc(
@ -113,13 +115,14 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource):
@service_api_ns.response(
200, "Metadata retrieved successfully", service_api_ns.models[DatasetMetadataListResponse.__name__]
)
def get(self, tenant_id, dataset_id: UUID):
@with_session(write=False)
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, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
metadata = MetadataService.get_dataset_metadatas(dataset, session=db.session())
metadata = MetadataService.get_dataset_metadatas(dataset, session)
return dump_response(DatasetMetadataListResponse, metadata), 200
@ -148,20 +151,19 @@ class DatasetMetadataServiceApi(DatasetApiResource):
200, "Metadata updated successfully", service_api_ns.models[DatasetMetadataResponse.__name__]
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def patch(self, tenant_id, dataset_id: UUID, metadata_id: UUID):
@with_session
def patch(self, session: Session, tenant_id, dataset_id: UUID, metadata_id: UUID):
"""Update metadata name."""
payload = MetadataUpdatePayload.model_validate(service_api_ns.payload or {})
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata = MetadataService.update_metadata_name(
dataset_id_str, metadata_id_str, payload.name, session=db.session()
)
metadata = MetadataService.update_metadata_name(dataset_id_str, metadata_id_str, payload.name, session=session)
return dump_response(DatasetMetadataResponse, metadata), 200
@service_api_ns.doc(
@ -187,16 +189,17 @@ class DatasetMetadataServiceApi(DatasetApiResource):
)
@service_api_ns.response(204, "Metadata deleted successfully")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def delete(self, tenant_id, dataset_id: UUID, metadata_id: UUID):
@with_session
def delete(self, session: Session, tenant_id, dataset_id: UUID, metadata_id: UUID):
"""Delete metadata."""
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session=db.session())
MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session)
return "", 204
@ -256,19 +259,20 @@ class DatasetMetadataBuiltInFieldActionServiceApi(DatasetApiResource):
200, "Action completed successfully", service_api_ns.models[DatasetMetadataActionResponse.__name__]
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id, dataset_id: UUID, action: Literal["enable", "disable"]):
@with_session
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, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
match action:
case "enable":
MetadataService.enable_built_in_field(dataset, session=db.session())
MetadataService.enable_built_in_field(dataset, session)
case "disable":
MetadataService.disable_built_in_field(dataset, session=db.session())
MetadataService.disable_built_in_field(dataset, session)
return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200
@ -302,16 +306,17 @@ class DocumentMetadataEditServiceApi(DatasetApiResource):
service_api_ns.models[DatasetMetadataActionResponse.__name__],
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id, dataset_id: UUID):
@with_session
def post(self, session: Session, tenant_id, dataset_id: UUID):
"""Update metadata for multiple documents."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str, db.session())
dataset = DatasetService.get_dataset(dataset_id_str, session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session())
DatasetService.check_dataset_permission(dataset, current_user, session)
metadata_args = MetadataOperationData.model_validate(service_api_ns.payload or {})
MetadataService.update_documents_metadata(dataset, metadata_args, session=db.session())
MetadataService.update_documents_metadata(dataset, metadata_args, session=session)
return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200

145
api/controllers/service_api/dataset/segment.py
View File

@ -3,6 +3,7 @@ from uuid import UUID
from pydantic import BaseModel, Field, ValidationError, field_validator
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from configs import dify_config
@ -13,6 +14,7 @@ from controllers.common.schema import (
register_response_schema_models,
register_schema_models,
)
from controllers.common.session import with_session
from controllers.service_api import service_api_ns
from controllers.service_api.app.error import ProviderNotInitializeError
from controllers.service_api.wraps import (
@ -24,7 +26,6 @@ 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_database import db
from fields.base import ResponseModel
from fields.segment_fields import (
ChildChunkDetailResponse,
@ -129,7 +130,7 @@ register_response_schema_models(
def _get_segment_for_document(
dataset: Dataset, document: Document, segment_id: str
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)
@ -137,7 +138,7 @@ def _get_segment_for_document(
raise NotFound("Document not found.")
segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id)
segment = SegmentService.get_segment_by_ref(segment_ref, db.session())
segment = SegmentService.get_segment_by_ref(segment_ref, session=session)
if not segment:
raise NotFound("Segment not found.")
return segment_ref, segment
@ -179,19 +180,20 @@ class SegmentApi(DatasetApiResource):
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_knowledge_limit_check("add_segment", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID):
_, current_tenant_id = current_account_with_tenant()
"""Create single segment."""
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
if document.indexing_status != "completed":
@ -227,17 +229,17 @@ 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, db.session())
list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset, session)
)
segment_ids = [segment.id for segment in segments]
summaries: dict[str, str | None] = {}
if segment_ids:
summary_records = SummaryIndexService.get_segments_summaries(
segment_ids=segment_ids, dataset_id=dataset_id_str, session=db.session()
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),
"data": segment_responses_with_summaries(segments, summaries, session=session),
"doc_form": document.doc_form,
}
return dump_response(SegmentCreateListResponse, response), 200
@ -266,7 +268,8 @@ class SegmentApi(DatasetApiResource):
"Segments retrieved successfully",
service_api_ns.models[SegmentListResponse.__name__],
)
def get(self, tenant_id: str, dataset_id: UUID, document_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID):
_, current_tenant_id = current_account_with_tenant()
"""Get segments."""
# check dataset
@ -278,14 +281,14 @@ class SegmentApi(DatasetApiResource):
page = args.page
limit = args.limit
dataset_id_str = str(dataset_id)
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
# check embedding model setting
@ -306,6 +309,7 @@ class SegmentApi(DatasetApiResource):
raise ProviderNotInitializeError(ex.description)
segments, total = SegmentService.get_segments(
session=session,
document_id=document_id_str,
tenant_id=current_tenant_id,
status_list=args.status,
@ -317,12 +321,12 @@ class SegmentApi(DatasetApiResource):
summaries: dict[str, str | None] = {}
if segment_ids:
summary_records = SummaryIndexService.get_segments_summaries(
segment_ids=segment_ids, dataset_id=dataset_id_str, session=db.session()
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),
"data": segment_responses_with_summaries(segments, summaries, session=session),
"doc_form": document.doc_form,
"total": total,
"has_more": len(segments) == limit,
@ -354,11 +358,12 @@ class DatasetSegmentApi(DatasetApiResource):
}
)
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def delete(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
@with_session
def delete(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
current_account_with_tenant()
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -367,12 +372,12 @@ class DatasetSegmentApi(DatasetApiResource):
DatasetService.check_dataset_model_setting(dataset)
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
SegmentService.delete_segment(segment, document, dataset, db.session())
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
SegmentService.delete_segment(segment, document, dataset, session)
return "", 204
@service_api_ns.doc(
@ -397,11 +402,12 @@ class DatasetSegmentApi(DatasetApiResource):
@service_api_ns.response(200, "Segment updated successfully", service_api_ns.models[SegmentDetailResponse.__name__])
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
_, current_tenant_id = current_account_with_tenant()
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -410,7 +416,7 @@ class DatasetSegmentApi(DatasetApiResource):
DatasetService.check_dataset_model_setting(dataset)
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
@ -430,16 +436,18 @@ class DatasetSegmentApi(DatasetApiResource):
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
payload = SegmentUpdatePayload.model_validate(service_api_ns.payload or {})
updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, db.session())
updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, session)
summary = SummaryIndexService.get_segment_summary(
segment_id=updated_segment.id, dataset_id=dataset_id_str, session=db.session()
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),
"data": segment_response_with_summary(
updated_segment, summary.summary_content if summary else None, session=session
),
"doc_form": document.doc_form,
}
return dump_response(SegmentDetailResponse, response), 200
@ -470,11 +478,12 @@ class DatasetSegmentApi(DatasetApiResource):
"Segment retrieved successfully",
service_api_ns.models[SegmentDetailResponse.__name__],
)
def get(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
current_account_with_tenant()
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -483,17 +492,19 @@ class DatasetSegmentApi(DatasetApiResource):
DatasetService.check_dataset_model_setting(dataset)
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
summary = SummaryIndexService.get_segment_summary(
segment_id=segment.id, dataset_id=dataset_id_str, session=db.session()
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),
"data": segment_response_with_summary(
segment, summary.summary_content if summary else None, session=session
),
"doc_form": document.doc_form,
}
return dump_response(SegmentDetailResponse, response), 200
@ -533,12 +544,13 @@ class ChildChunkApi(DatasetApiResource):
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_knowledge_limit_check("add_segment", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def post(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
@with_session
def post(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
_, current_tenant_id = current_account_with_tenant()
"""Create child chunk."""
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -546,12 +558,12 @@ class ChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
_, segment = _get_segment_for_document(dataset, document, segment_id_str)
_, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
# check embedding model setting
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
@ -574,7 +586,7 @@ class ChildChunkApi(DatasetApiResource):
payload = ChildChunkCreatePayload.model_validate(service_api_ns.payload or {})
try:
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session())
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, session)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
@ -604,12 +616,13 @@ class ChildChunkApi(DatasetApiResource):
"Child chunks retrieved successfully",
service_api_ns.models[ChildChunkListResponse.__name__],
)
def get(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
@with_session
def get(self, session: Session, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
current_account_with_tenant()
"""Get child chunks."""
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -617,12 +630,12 @@ class ChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
_get_segment_for_document(dataset, document, segment_id_str)
_get_segment_for_document(session, dataset, document, segment_id_str)
args = query_params_from_request(ChildChunkListQuery, use_defaults_for_malformed_ints=True)
@ -631,7 +644,13 @@ class ChildChunkApi(DatasetApiResource):
keyword = args.keyword
child_chunks = SegmentService.get_child_chunks(
segment_id_str, document_id_str, dataset_id_str, page, limit, keyword
segment_id_str,
document_id_str,
dataset_id_str,
page,
limit,
keyword,
session=session,
)
response = {
@ -671,12 +690,21 @@ class DatasetChildChunkApi(DatasetApiResource):
)
@cloud_edition_billing_knowledge_limit_check("add_segment", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def delete(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID, child_chunk_id: UUID):
@with_session
def delete(
self,
session: Session,
tenant_id: str,
dataset_id: UUID,
document_id: UUID,
segment_id: UUID,
child_chunk_id: UUID,
):
current_account_with_tenant()
"""Delete child chunk."""
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -684,21 +712,21 @@ class DatasetChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session())
document = DocumentService.get_document(dataset.id, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
segment_ref, _ = _get_segment_for_document(dataset, document, segment_id_str)
segment_ref, _ = _get_segment_for_document(session, dataset, document, segment_id_str)
child_chunk_id_str = str(child_chunk_id)
# check child chunk
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session())
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, session=session)
if not child_chunk:
raise NotFound("Child chunk not found.")
try:
SegmentService.delete_child_chunk(child_chunk, dataset, db.session())
SegmentService.delete_child_chunk(child_chunk, dataset, session)
except ChildChunkDeleteIndexServiceError as e:
raise ChildChunkDeleteIndexError(str(e))
@ -732,12 +760,21 @@ class DatasetChildChunkApi(DatasetApiResource):
@cloud_edition_billing_resource_check("vector_space", "dataset")
@cloud_edition_billing_knowledge_limit_check("add_segment", "dataset")
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def patch(self, tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID, child_chunk_id: UUID):
@with_session
def patch(
self,
session: Session,
tenant_id: str,
dataset_id: UUID,
document_id: UUID,
segment_id: UUID,
child_chunk_id: UUID,
):
current_account_with_tenant()
"""Update child chunk."""
dataset_id_str = str(dataset_id)
# check dataset
dataset = db.session.scalar(
dataset = session.scalar(
select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str).limit(1)
)
if not dataset:
@ -745,16 +782,16 @@ class DatasetChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# get document
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session())
document = DocumentService.get_document(dataset_id_str, document_id_str, session=session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
segment_ref, segment = _get_segment_for_document(dataset, document, segment_id_str)
segment_ref, segment = _get_segment_for_document(session, dataset, document, segment_id_str)
child_chunk_id_str = str(child_chunk_id)
# get child chunk
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session())
child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, session=session)
if not child_chunk:
raise NotFound("Child chunk not found.")
@ -763,7 +800,7 @@ class DatasetChildChunkApi(DatasetApiResource):
try:
child_chunk = SegmentService.update_child_chunk(
payload.content, child_chunk, segment, document, dataset, db.session()
payload.content, child_chunk, segment, document, dataset, session
)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))

4
api/controllers/service_api/wraps.py
View File

@ -161,7 +161,7 @@ def validate_app_token[**P, R](
if tenant_owner_info:
tenant_model, account = tenant_owner_info
account.current_tenant = tenant_model
account.set_current_tenant_with_session(tenant_model, session=db.session())
current_app.login_manager._update_request_context_with_user(account) # type: ignore
user_logged_in.send(current_app._get_current_object(), user=current_user) # type: ignore
else:
@ -333,7 +333,7 @@ def validate_dataset_token[R](view: Callable[..., R]) -> Callable[..., R]:
account = db.session.get(Account, ta.account_id)
# Login admin
if account:
account.current_tenant = tenant
account.set_current_tenant_with_session(tenant, session=db.session())
current_app.login_manager._update_request_context_with_user(account) # type: ignore
user_logged_in.send(current_app._get_current_object(), user=current_user) # type: ignore
else:

18
api/controllers/web/app.py
View File

@ -15,7 +15,7 @@ from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPu
from extensions.ext_database import db
from libs.passport import PassportService
from libs.token import extract_webapp_passport
from models.model import App, AppMode, EndUser
from models.model import App, AppMode, EndUser, load_annotation_reply_config
from services.app_service import AppService
from services.enterprise.enterprise_service import EnterpriseService
from services.feature_service import FeatureService
@ -77,28 +77,36 @@ class AppParameterApi(WebApiResource):
@web_ns.response(200, "Success", web_ns.models[fields.Parameters.__name__])
def get(self, app_model: App, end_user: EndUser):
"""Retrieve app parameters."""
session = db.session()
features_dict: dict[str, Any]
user_input_form: list[dict[str, Any]]
if app_model.mode == AppMode.AGENT:
try:
features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(app_model)
features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(
app_model,
session=session,
)
except AgentAppNotPublishedError:
raise AgentNotPublishedError()
except AgentAppGeneratorError:
raise AppUnavailableError()
elif app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
workflow = app_model.workflow
workflow = app_model.workflow_with_session(session=session)
if workflow is None:
raise AppUnavailableError()
features_dict = workflow.features_dict
user_input_form = workflow.user_input_form(to_old_structure=True)
else:
app_model_config = app_model.app_model_config
app_model_config = app_model.app_model_config_with_session(session=session)
if app_model_config is None:
raise AppUnavailableError()
features_dict = cast(dict[str, Any], app_model_config.to_dict())
annotation_reply = load_annotation_reply_config(session, app_model.id)
features_dict = cast(
dict[str, Any],
app_model_config.to_dict(annotation_reply=annotation_reply),
)
user_input_form = features_dict.get("user_input_form", [])

7
api/controllers/web/audio.py
View File

@ -79,7 +79,12 @@ class AudioApi(WebApiResource):
file = request.files["file"]
try:
response = AudioService.transcript_asr(app_model=app_model, file=file, end_user=end_user.external_user_id)
response = AudioService.transcript_asr(
app_model=app_model,
file=file,
session=db.session(),
end_user=end_user.external_user_id,
)
return dump_response(AudioToTextResponse, response)
except services.errors.app_model_config.AppModelConfigBrokenError:

3
api/controllers/web/completion.py
View File

@ -30,7 +30,6 @@ from core.errors.error import (
ProviderTokenNotInitError,
QuotaExceededError,
)
from extensions.ext_database import db
from graphon.model_runtime.errors.invoke import InvokeError
from libs import helper
from libs.helper import uuid_value
@ -224,7 +223,7 @@ class ChatApi(WebApiResource):
app_model=app_model,
conversation_id=payload.conversation_id,
user=end_user,
session=db.session(),
session=session,
)
response = AppGenerateService.generate(

14
api/controllers/web/conversation.py
View File

@ -15,6 +15,7 @@ from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from fields.conversation_fields import (
ConversationInfiniteScrollPagination,
ConversationResponseSource,
ResultResponse,
SimpleConversation,
)
@ -80,7 +81,13 @@ class ConversationListApi(WebApiResource):
sort_by=query.sort_by,
)
adapter = TypeAdapter(SimpleConversation)
conversations = [adapter.validate_python(item, from_attributes=True) for item in pagination.data]
conversations = [
adapter.validate_python(
ConversationResponseSource(item, session=session),
from_attributes=True,
)
for item in pagination.data
]
return ConversationInfiniteScrollPagination(
limit=pagination.limit,
has_more=pagination.has_more,
@ -156,12 +163,13 @@ class ConversationRenameApi(WebApiResource):
payload = ConversationRenamePayload.model_validate(web_ns.payload or {})
try:
session = db.session()
conversation = ConversationService.rename(
app_model, conversation_id, end_user, payload.name, payload.auto_generate, session=db.session()
app_model, conversation_id, end_user, payload.name, payload.auto_generate, session=session
)
return (
TypeAdapter(SimpleConversation)
.validate_python(conversation, from_attributes=True)
.validate_python(ConversationResponseSource(conversation, session=session), from_attributes=True)
.model_dump(mode="json")
)
except ConversationNotExistsError:

10
api/controllers/web/message.py
View File

@ -26,7 +26,7 @@ from controllers.web.wraps import WebApiResource
from core.app.entities.app_invoke_entities import InvokeFrom
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from extensions.ext_database import db
from fields.conversation_fields import ResultResponse
from fields.conversation_fields import MessageResponseSource, ResultResponse
from fields.message_fields import SuggestedQuestionsResponse, WebMessageInfiniteScrollPagination, WebMessageListItem
from graphon.model_runtime.errors.invoke import InvokeError
from libs import helper
@ -86,11 +86,15 @@ class MessageListApi(WebApiResource):
query = MessageListQuery.model_validate(raw_args)
try:
session = db.session()
pagination = MessageService.pagination_by_first_id(
app_model, end_user, query.conversation_id, query.first_id, query.limit, session=db.session()
app_model, end_user, query.conversation_id, query.first_id, query.limit, session=session
)
adapter = TypeAdapter(WebMessageListItem)
items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
items = [
adapter.validate_python(MessageResponseSource(message, session=session), from_attributes=True)
for message in pagination.data
]
return WebMessageInfiniteScrollPagination(
limit=pagination.limit,
has_more=pagination.has_more,

10
api/controllers/web/saved_message.py
View File

@ -10,7 +10,7 @@ from controllers.web import web_ns
from controllers.web.error import NotCompletionAppError
from controllers.web.wraps import WebApiResource
from extensions.ext_database import db
from fields.conversation_fields import ResultResponse
from fields.conversation_fields import MessageResponseSource, ResultResponse
from fields.message_fields import SavedMessageInfiniteScrollPagination, SavedMessageItem
from models.model import App, EndUser
from services.errors.message import MessageNotExistsError
@ -43,11 +43,15 @@ class SavedMessageListApi(WebApiResource):
raw_args = request.args.to_dict()
query = SavedMessageListQuery.model_validate(raw_args)
session = db.session()
pagination = SavedMessageService.pagination_by_last_id(
app_model, end_user, query.last_id, query.limit, session=db.session()
app_model, end_user, query.last_id, query.limit, session=session
)
adapter = TypeAdapter(SavedMessageItem)
items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
items = [
adapter.validate_python(MessageResponseSource(message, session=session), from_attributes=True)
for message in pagination.data
]
return SavedMessageInfiniteScrollPagination(
limit=pagination.limit, has_more=pagination.has_more, data=items
).model_dump(mode="json")

32
api/core/agent/base_agent_runner.py
View File

@ -42,7 +42,7 @@ from graphon.model_runtime.entities.message_entities import ImagePromptMessageCo
from graphon.model_runtime.entities.model_entities import ModelFeature
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
from models.enums import CreatorUserRole
from models.model import Conversation, Message, MessageAgentThought, MessageFile
from models.model import Conversation, Message, MessageAgentThought, MessageFile, load_annotation_reply_config
logger = logging.getLogger(__name__)
_file_access_controller = DatabaseFileAccessController()
@ -76,7 +76,9 @@ class BaseAgentRunner(AppRunner):
self.message = message
self.user_id = user_id
self.memory = memory
self.history_prompt_messages = self.organize_agent_history(prompt_messages=prompt_messages or [])
self.history_prompt_messages = self.organize_agent_history(
session=session, prompt_messages=prompt_messages or []
)
self.model_instance = model_instance
# init callback
@ -104,7 +106,7 @@ class BaseAgentRunner(AppRunner):
)
# get how many agent thoughts have been created
self.agent_thought_count = (
db.session.scalar(
session.scalar(
select(func.count())
.select_from(MessageAgentThought)
.where(
@ -113,7 +115,7 @@ class BaseAgentRunner(AppRunner):
)
or 0
)
db.session.close()
session.close()
# check if model supports stream tool call
llm_model = cast(LargeLanguageModel, model_instance.model_type_instance)
@ -350,7 +352,7 @@ class BaseAgentRunner(AppRunner):
db.session.commit()
db.session.close()
def organize_agent_history(self, prompt_messages: list[PromptMessage]) -> list[PromptMessage]:
def organize_agent_history(self, prompt_messages: list[PromptMessage], *, session: Session) -> list[PromptMessage]:
"""
Organize agent history
"""
@ -362,7 +364,7 @@ class BaseAgentRunner(AppRunner):
messages = (
(
db.session.execute(
session.execute(
select(Message)
.where(Message.conversation_id == self.message.conversation_id)
.order_by(Message.created_at.desc())
@ -378,8 +380,8 @@ class BaseAgentRunner(AppRunner):
if message.id == self.message.id:
continue
result.append(self.organize_agent_user_prompt(message))
agent_thoughts = message.agent_thoughts
result.append(self.organize_agent_user_prompt(message, session=session))
agent_thoughts = message.agent_thoughts_with_session(session=session)
if agent_thoughts:
for agent_thought in agent_thoughts:
tool_names_raw = agent_thought.tool
@ -441,17 +443,21 @@ class BaseAgentRunner(AppRunner):
if message.answer:
result.append(AssistantPromptMessage(content=message.answer))
db.session.close()
session.close()
return result
def organize_agent_user_prompt(self, message: Message) -> UserPromptMessage:
def organize_agent_user_prompt(self, message: Message, *, session: Session) -> UserPromptMessage:
stmt = select(MessageFile).where(MessageFile.message_id == message.id)
files = db.session.scalars(stmt).all()
files = session.scalars(stmt).all()
if not files:
return UserPromptMessage(content=message.query)
if message.app_model_config:
file_extra_config = FileUploadConfigManager.convert(message.app_model_config.to_dict())
app_model_config = message.app_model_config_with_session(session=session)
if app_model_config:
annotation_reply = load_annotation_reply_config(session, app_model_config.app_id)
file_extra_config = FileUploadConfigManager.convert(
app_model_config.to_dict(annotation_reply=annotation_reply)
)
else:
file_extra_config = None

13
api/core/agent/cot_agent_runner.py
View File

@ -114,7 +114,11 @@ class CotAgentRunner(BaseAgentRunner, ABC):
message_file_ids: list[str] = []
agent_thought_id = self.create_agent_thought(
message_id=message.id, message="", tool_name="", tool_input="", messages_ids=message_file_ids
message_id=message.id,
message="",
tool_name="",
tool_input="",
messages_ids=message_file_ids,
)
if iteration_step > 1:
@ -125,6 +129,11 @@ class CotAgentRunner(BaseAgentRunner, ABC):
# recalc llm max tokens
prompt_messages = self._organize_prompt_messages()
self.recalc_llm_max_tokens(self.model_config, prompt_messages)
# Release any setup/tool transaction before waiting on the provider stream.
session.commit()
session.close()
# invoke model
chunks = model_instance.invoke_llm(
prompt_messages=prompt_messages,
@ -333,6 +342,8 @@ class CotAgentRunner(BaseAgentRunner, ABC):
agent_tool_callback=self.agent_callback,
trace_manager=trace_manager,
)
session.commit()
session.close()
# publish files
for message_file_id in message_files:

13
api/core/agent/fc_agent_runner.py
View File

@ -87,12 +87,21 @@ class FunctionCallAgentRunner(BaseAgentRunner):
message_file_ids: list[str] = []
agent_thought_id = self.create_agent_thought(
message_id=message.id, message="", tool_name="", tool_input="", messages_ids=message_file_ids
message_id=message.id,
message="",
tool_name="",
tool_input="",
messages_ids=message_file_ids,
)
# recalc llm max tokens
prompt_messages = self._organize_prompt_messages()
self.recalc_llm_max_tokens(self.model_config, prompt_messages)
# Release any setup/tool transaction before waiting on the provider stream.
session.commit()
session.close()
# invoke model
chunks: Union[Generator[LLMResultChunk, None, None], LLMResult] = model_instance.invoke_llm(
prompt_messages=prompt_messages,
@ -256,6 +265,8 @@ class FunctionCallAgentRunner(BaseAgentRunner):
message_id=self.message.id,
conversation_id=self.conversation.id,
)
session.commit()
session.close()
# publish files
for message_file_id in message_files:
# publish message file

17
api/core/app/app_config/easy_ui_based_app/dataset/manager.py
View File

@ -1,6 +1,8 @@
import uuid
from typing import Any, Literal, cast
from sqlalchemy.orm import Session
from core.app.app_config.entities import (
DatasetEntity,
DatasetRetrieveConfigEntity,
@ -9,7 +11,6 @@ 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 extensions.ext_database import db
from models.model import AppMode, AppModelConfigDict
from services.dataset_service import DatasetService
@ -140,7 +141,7 @@ class DatasetConfigManager:
@classmethod
def validate_and_set_defaults(
cls, tenant_id: str, app_mode: AppMode, config: dict[str, Any]
cls, tenant_id: str, app_mode: AppMode, config: dict[str, Any], session: Session
) -> tuple[dict[str, Any], list[str]]:
"""
Validate and set defaults for dataset feature
@ -150,7 +151,7 @@ class DatasetConfigManager:
:param config: app model config args
"""
# Extract dataset config for legacy compatibility
config = cls.extract_dataset_config_for_legacy_compatibility(tenant_id, app_mode, config)
config = cls.extract_dataset_config_for_legacy_compatibility(tenant_id, app_mode, config, session)
# dataset_configs
if "dataset_configs" not in config or not config.get("dataset_configs"):
@ -175,7 +176,9 @@ class DatasetConfigManager:
return config, ["agent_mode", "dataset_configs", "dataset_query_variable"]
@classmethod
def extract_dataset_config_for_legacy_compatibility(cls, tenant_id: str, app_mode: AppMode, config: dict[str, Any]):
def extract_dataset_config_for_legacy_compatibility(
cls, tenant_id: str, app_mode: AppMode, config: dict[str, Any], session: Session
):
"""
Extract dataset config for legacy compatibility
@ -238,7 +241,7 @@ class DatasetConfigManager:
except ValueError:
raise ValueError("id in dataset must be of UUID type")
if not cls.is_dataset_exists(tenant_id, tool_item["id"]):
if not cls.is_dataset_exists(tenant_id, tool_item["id"], session):
raise ValueError("Dataset ID does not exist, please check your permission.")
has_datasets = True
@ -255,9 +258,9 @@ class DatasetConfigManager:
return config
@classmethod
def is_dataset_exists(cls, tenant_id: str, dataset_id: str) -> bool:
def is_dataset_exists(cls, tenant_id: str, dataset_id: str, session: Session) -> bool:
# verify if the dataset ID exists
dataset = DatasetService.get_dataset(dataset_id, db.session())
dataset = DatasetService.get_dataset(dataset_id, session)
if not dataset:
return False

42
api/core/app/apps/advanced_chat/app_generator.py
View File

@ -86,6 +86,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
workflow_run_id: str,
streaming: Literal[False],
pause_state_config: PauseStateLayerConfig | None = None,
*,
session: Session,
) -> Mapping[str, Any]: ...
@overload
@ -99,6 +101,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
workflow_run_id: str,
streaming: Literal[True],
pause_state_config: PauseStateLayerConfig | None = None,
*,
session: Session,
) -> Generator[Mapping | str, None, None]: ...
@overload
@ -112,6 +116,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
workflow_run_id: str,
streaming: bool,
pause_state_config: PauseStateLayerConfig | None = None,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping, None, None]: ...
def generate(
@ -124,6 +130,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
workflow_run_id: str,
streaming: bool = True,
pause_state_config: PauseStateLayerConfig | None = None,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping, None, None]:
"""
Generate App response.
@ -134,6 +142,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
:param args: request args
:param invoke_from: invoke from source
:param streaming: is stream
:param session: database session supplied by the caller
"""
if not args.get("query"):
raise ValueError("query is required")
@ -157,7 +166,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
if conversation_id:
try:
conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user, session=db.session()
app_model=app_model, conversation_id=conversation_id, user=user, session=session
)
except ConversationNotExistsError:
if invoke_from == InvokeFrom.SERVICE_API:
@ -255,6 +264,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
conversation=conversation,
stream=streaming,
pause_state_config=pause_state_config,
session=session,
)
def resume(
@ -265,6 +275,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
user: Account | EndUser,
conversation: Conversation,
message: Message,
session: Session,
application_generate_entity: AdvancedChatAppGenerateEntity,
workflow_execution_repository: WorkflowExecutionRepository,
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
@ -301,6 +312,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
pause_state_config=pause_state_config,
graph_runtime_state=graph_runtime_state,
response_stream_filter=response_stream_filter,
session=session,
)
def single_iteration_generate(
@ -311,6 +323,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
user: Account | EndUser,
args: Mapping[str, Any],
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -321,6 +335,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
:param user: account or end user
:param args: request args
:param streaming: is streamed
:param session: database session supplied by the caller
"""
if not node_id:
raise ValueError("node_id is required")
@ -377,7 +392,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
tenant_id=application_generate_entity.app_config.tenant_id,
user_id=user.id,
)
draft_var_srv = WorkflowDraftVariableService(db.session())
draft_var_srv = WorkflowDraftVariableService(session)
draft_var_srv.prefill_conversation_variable_default_values(workflow, user_id=user.id)
return self._generate(
@ -390,6 +405,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
conversation=None,
stream=streaming,
variable_loader=var_loader,
session=session,
)
def single_loop_generate(
@ -400,6 +416,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
user: Account | EndUser,
args: LoopNodeRunPayload,
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -410,6 +428,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
:param user: account or end user
:param args: request args
:param streaming: is stream
:param session: database session supplied by the caller
"""
if not node_id:
raise ValueError("node_id is required")
@ -464,7 +483,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
tenant_id=application_generate_entity.app_config.tenant_id,
user_id=user.id,
)
draft_var_srv = WorkflowDraftVariableService(db.session())
draft_var_srv = WorkflowDraftVariableService(session)
draft_var_srv.prefill_conversation_variable_default_values(workflow, user_id=user.id)
return self._generate(
@ -477,6 +496,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
conversation=None,
stream=streaming,
variable_loader=var_loader,
session=session,
)
def _generate(
@ -486,6 +506,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
user: Account | EndUser,
invoke_from: InvokeFrom,
application_generate_entity: AdvancedChatAppGenerateEntity,
session: Session,
workflow_execution_repository: WorkflowExecutionRepository,
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
conversation: Conversation | None = None,
@ -504,6 +525,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
:param user: account or end user
:param invoke_from: invoke from source
:param application_generate_entity: application generate entity
:param session: database session supplied by the caller
:param workflow_execution_repository: repository for workflow execution
:param workflow_node_execution_repository: repository for workflow node execution
:param conversation: conversation
@ -519,18 +541,22 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
if conversation is not None and message is not None:
pass
else:
conversation, message = self._init_generate_records(application_generate_entity, conversation)
conversation, message = self._init_generate_records(
application_generate_entity,
conversation,
session=session,
)
if is_first_conversation:
# update conversation features
conversation.override_model_configs = workflow.features
db.session.commit()
db.session.refresh(conversation)
session.commit()
session.refresh(conversation)
# get conversation dialogue count
# NOTE: dialogue_count should not start from 0,
# because during the first conversation, dialogue_count should be 1.
self._dialogue_count = get_thread_messages_length(conversation.id) + 1
self._dialogue_count = get_thread_messages_length(conversation.id, session=session) + 1
# init queue manager
queue_manager = MessageBasedAppQueueManager(
@ -582,7 +608,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
workflow_snapshot = WorkflowSnapshot.from_workflow(workflow)
conversation_snapshot = ConversationSnapshot.from_conversation(conversation)
message_snapshot = MessageSnapshot.from_message(message)
db.session.close()
session.close()
# return response or stream generator
response = self._handle_advanced_chat_response(

67
api/core/app/apps/advanced_chat/app_runner.py
View File

@ -39,7 +39,6 @@ from core.workflow.system_variables import (
)
from core.workflow.variable_pool_initializer import add_node_inputs_to_pool, add_variables_to_pool
from core.workflow.workflow_entry import WorkflowEntry
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from extensions.otel import WorkflowAppRunnerHandler, trace_span
from extensions.workflow_warm_shutdown import WORKFLOW_WARM_SHUTDOWN_ABORT_REASON, celery_warm_shutdown_started
@ -173,12 +172,21 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
)
# annotation reply
if self.handle_annotation_reply(
app_record=self._app,
message=self.message,
query=new_query,
app_generate_entity=self.application_generate_entity,
):
with create_session() as session:
annotation_reply = self.handle_annotation_reply(
app_record=self._app,
message=self.message,
query=new_query,
app_generate_entity=self.application_generate_entity,
session=session,
)
session.commit()
if annotation_reply:
self._publish_event(QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id))
self._complete_with_stream_output(
text=annotation_reply.content,
stopped_by=QueueStopEvent.StopBy.ANNOTATION_REPLY,
)
return
# Initialize conversation variables
@ -212,10 +220,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
trace_session_id=self.application_generate_entity.extras.get("trace_session_id"),
)
# Release the Flask scoped session before workflow execution so a checked-out DB connection
# is not held for the lifetime of the graph run.
db.session.close()
# RUN WORKFLOW
# Create Redis command channel for this workflow execution
task_id = self.application_generate_entity.task_id
@ -300,26 +304,22 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
return False, new_inputs, new_query
def handle_annotation_reply(
self, app_record: App, message: Message, query: str, app_generate_entity: AdvancedChatAppGenerateEntity
) -> bool:
annotation_reply = self.query_app_annotations_to_reply(
self,
app_record: App,
message: Message,
query: str,
app_generate_entity: AdvancedChatAppGenerateEntity,
session: Session,
) -> MessageAnnotation | None:
return self.query_app_annotations_to_reply(
app_record=app_record,
message=message,
query=query,
user_id=app_generate_entity.user_id,
invoke_from=app_generate_entity.invoke_from,
session=session,
)
if annotation_reply:
self._publish_event(QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id))
self._complete_with_stream_output(
text=annotation_reply.content, stopped_by=QueueStopEvent.StopBy.ANNOTATION_REPLY
)
return True
return False
def _complete_with_stream_output(self, text: str, stopped_by: QueueStopEvent.StopBy):
"""
Direct output
@ -329,7 +329,13 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
self._publish_event(QueueStopEvent(stopped_by=stopped_by))
def query_app_annotations_to_reply(
self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom
self,
app_record: App,
message: Message,
query: str,
user_id: str,
invoke_from: InvokeFrom,
session: Session,
) -> MessageAnnotation | None:
"""
Query app annotations to reply
@ -342,7 +348,12 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
"""
annotation_reply_feature = AnnotationReplyFeature()
return annotation_reply_feature.query(
app_record=app_record, message=message, query=query, user_id=user_id, invoke_from=invoke_from
app_record=app_record,
message=message,
query=query,
user_id=user_id,
invoke_from=invoke_from,
session=session,
)
def moderation_for_inputs(
@ -395,7 +406,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
existing_variables = self._create_all_conversation_variables(session)
else:
# Check and add any missing variables from the workflow
existing_variables = self._sync_missing_conversation_variables(session, existing_variables)
existing_variables = self._sync_missing_conversation_variables(existing_variables, session)
# Convert to Variable objects for use in the workflow
conversation_variables = [var.to_variable() for var in existing_variables]
@ -435,7 +446,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
return new_variables
def _sync_missing_conversation_variables(
self, session: Session, existing_variables: list[ConversationVariable]
self, existing_variables: list[ConversationVariable], session: Session
) -> list[ConversationVariable]:
"""
Sync missing conversation variables from the workflow definition.

16
api/core/app/apps/advanced_chat/generate_task_pipeline.py
View File

@ -9,7 +9,7 @@ from threading import Thread
from typing import Any, Union
from sqlalchemy import select, update
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.orm import Session
from constants.tts_auto_play_timeout import TTS_AUTO_PLAY_TIMEOUT, TTS_AUTO_PLAY_YIELD_CPU_TIME
from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
@ -71,12 +71,12 @@ from core.app.entities.task_entities import (
from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk
from core.db.session_factory import session_factory
from core.ops.ops_trace_manager import TraceQueueManager
from core.repositories.human_input_repository import HumanInputFormRepositoryImpl
from core.workflow.file_reference import resolve_file_record_id
from core.workflow.nodes.human_input.pause_reason import HumanInputRequired
from core.workflow.system_variables import build_system_variables
from extensions.ext_database import db
from graphon.enums import WorkflowExecutionStatus
from graphon.model_runtime.entities.llm_entities import LLMUsage
from graphon.model_runtime.utils.encoders import jsonable_encoder
@ -399,8 +399,13 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
@contextmanager
def _database_session(self):
"""Context manager for database sessions."""
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
yield session
with session_factory.create_session() as session:
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
def _ensure_workflow_initialized(self):
"""Fluent validation for workflow state."""
@ -825,7 +830,8 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
self, event: QueueAnnotationReplyEvent, **kwargs
) -> Generator[StreamResponse, None, None]:
"""Handle annotation reply events."""
self._message_cycle_manager.handle_annotation_reply(event)
with self._database_session() as session:
self._message_cycle_manager.handle_annotation_reply(event, session)
yield from ()
def _handle_message_replace_event(

17
api/core/app/apps/agent_app/app_config_manager.py
View File

@ -24,7 +24,7 @@ from core.app.app_config.entities import (
from core.app.apps.agent_app.app_feature_projection import merge_agent_app_features
from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form
from models.agent_config_entities import AgentSoulConfig
from models.model import App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
from models.model import AnnotationReplyConfig, App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
class AgentAppConfig(EasyUIBasedAppConfig):
@ -43,11 +43,16 @@ class AgentAppConfigManager(BaseAppConfigManager):
*,
app_model: App,
agent_soul: AgentSoulConfig,
annotation_reply: AnnotationReplyConfig | None,
app_model_config: AppModelConfig | None = None,
conversation: Conversation | None = None,
) -> AgentAppConfig:
"""Build the Agent App config from the Agent Soul (+ optional feature flags)."""
config_dict = cls._synthesize_config_dict(agent_soul, app_model_config)
config_dict = cls._synthesize_config_dict(
agent_soul,
app_model_config,
annotation_reply=annotation_reply,
)
# The synthesized dict is shaped like an app_model_config; the EasyUI
# sub-managers type their param as AppModelConfigDict (a TypedDict).
typed_config = cast(AppModelConfigDict, config_dict)
@ -77,6 +82,8 @@ class AgentAppConfigManager(BaseAppConfigManager):
def _synthesize_config_dict(
agent_soul: AgentSoulConfig,
app_model_config: AppModelConfig | None,
*,
annotation_reply: AnnotationReplyConfig | None,
) -> dict[str, Any]:
"""Shape a Soul + feature flags into an ``app_model_config``-style dict.
@ -84,7 +91,11 @@ class AgentAppConfigManager(BaseAppConfigManager):
``app_model_config`` when one exists; model + prompt always come from
the Agent Soul (the single source of truth for those).
"""
base = merge_agent_app_features(agent_soul=agent_soul, app_model_config=app_model_config)
base = merge_agent_app_features(
agent_soul=agent_soul,
app_model_config=app_model_config,
annotation_reply=annotation_reply,
)
model = agent_soul.model
if model is not None:

9
api/core/app/apps/agent_app/app_feature_projection.py
View File

@ -1,12 +1,14 @@
from typing import Any
from models.agent_config_entities import AgentSoulConfig
from models.model import AnnotationReplyConfig
def merge_agent_app_features(
*,
agent_soul: AgentSoulConfig,
app_model_config: Any | None,
annotation_reply: AnnotationReplyConfig | None,
) -> dict[str, Any]:
"""Project public Agent App features from legacy config plus Agent Soul.
@ -14,7 +16,12 @@ def merge_agent_app_features(
opening statements. Agent Soul is the source of truth for Agent-owned
features like file upload, so Soul fields override same-named legacy keys.
"""
features: dict[str, Any] = dict(app_model_config.to_dict()) if app_model_config else {}
if app_model_config is None:
features: dict[str, Any] = {}
else:
if annotation_reply is None:
raise ValueError("Annotation reply config is required")
features = dict(app_model_config.to_dict(annotation_reply=annotation_reply))
soul_features = agent_soul.app_features.model_dump(mode="json", exclude_none=True)
features.update(soul_features)
return features

146
api/core/app/apps/agent_app/app_generator.py
View File

@ -21,6 +21,7 @@ from typing import Any, Literal
from flask import Flask, current_app
from pydantic import JsonValue
from sqlalchemy import and_, or_, select
from sqlalchemy.orm import Session
from clients.agent_backend import AgentBackendRunEventAdapter
from clients.agent_backend.factory import create_agent_backend_run_client
@ -46,10 +47,11 @@ from core.app.entities.app_invoke_entities import (
UserFrom,
)
from core.app.llm.model_access import build_dify_model_access
from core.db.session_factory import session_factory
from core.ops.ops_trace_manager import TraceQueueManager
from core.workflow.file_reference import build_file_reference, is_canonical_file_reference
from extensions.ext_database import db
from models import Account, App, EndUser, Message
from models import Account, App, AppModelConfig, EndUser, Message, MessageAnnotation
from models.agent import (
APP_BACKED_AGENT_SOURCES,
Agent,
@ -61,6 +63,7 @@ from models.agent import (
AgentStatus,
)
from models.agent_config_entities import AgentSoulConfig
from models.model import load_annotation_reply_config
from services.conversation_service import ConversationService
logger = logging.getLogger(__name__)
@ -137,6 +140,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
session: Session,
streaming: bool = True,
) -> Mapping[str, Any] | Generator[Mapping | str, None, None]:
if not streaming:
@ -152,6 +156,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
invoke_from=invoke_from,
draft_type=args.get("draft_type"),
user=user,
session=session,
)
runtime_session_snapshot_id = self._runtime_session_snapshot_id(
invoke_from=invoke_from,
@ -162,15 +167,19 @@ class AgentAppGenerator(MessageBasedAppGenerator):
conversation_id = args.get("conversation_id")
if conversation_id:
conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user, session=db.session()
app_model=app_model, conversation_id=conversation_id, user=user, session=session
)
# Build the EasyUI-shaped config from the Agent Soul so the chat pipeline
# can persist usage; the answer itself comes from the agent backend.
app_model_config = app_model.app_model_config
app_model_config = (
session.get(AppModelConfig, app_model.app_model_config_id) if app_model.app_model_config_id else None
)
annotation_reply = load_annotation_reply_config(session, app_model.id) if app_model_config else None
app_config = AgentAppConfigManager.get_app_config(
app_model=app_model,
agent_soul=agent_soul,
annotation_reply=annotation_reply,
app_model_config=app_model_config,
conversation=conversation,
)
@ -210,7 +219,11 @@ class AgentAppGenerator(MessageBasedAppGenerator):
agent_runtime_exit_intent=agent_runtime_exit_intent,
)
conversation, message = self._init_generate_records(application_generate_entity, conversation)
conversation, message = self._init_generate_records(
application_generate_entity,
conversation,
session=session,
)
queue_manager = MessageBasedAppQueueManager(
task_id=application_generate_entity.task_id,
@ -253,6 +266,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
user: Account | EndUser,
conversation_id: str,
invoke_from: InvokeFrom,
session: Session,
) -> None:
"""Resume an Agent App conversation after a submitted ask_human HITL form.
@ -263,19 +277,28 @@ class AgentAppGenerator(MessageBasedAppGenerator):
out of scope here — the message is persisted and can be re-fetched.
"""
conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user, session=db.session()
app_model=app_model, conversation_id=conversation_id, user=user, session=session
)
agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent(
app_model,
invoke_from=invoke_from,
draft_type=self._resume_draft_type(app_model=app_model, conversation=conversation, user=user),
draft_type=self._resume_draft_type(
app_model=app_model, conversation=conversation, user=user, session=session
),
user=user,
session=session,
)
app_model_config = (
session.get(AppModelConfig, app_model.app_model_config_id) if app_model.app_model_config_id else None
)
annotation_reply = load_annotation_reply_config(session, app_model.id) if app_model_config else None
app_config = AgentAppConfigManager.get_app_config(
app_model=app_model,
agent_soul=agent_soul,
app_model_config=app_model.app_model_config,
annotation_reply=annotation_reply,
app_model_config=app_model_config,
conversation=conversation,
)
model_conf = ModelConfigConverter.convert(app_config)
@ -287,7 +310,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
# turn's query); the continuation is driven by deferred_tool_results and
# the restored snapshot, not by re-processing this prompt. A blank prompt
# would drop the user-prompt layer and fail the snapshot match.
paused_message = db.session.scalar(
paused_message = session.scalar(
select(Message)
.where(Message.conversation_id == conversation.id, Message.query != "")
.order_by(Message.created_at.desc())
@ -318,7 +341,11 @@ class AgentAppGenerator(MessageBasedAppGenerator):
agent_config_version_kind=agent_config_version_kind,
)
conversation, message = self._init_generate_records(application_generate_entity, conversation)
conversation, message = self._init_generate_records(
application_generate_entity,
conversation,
session=session,
)
queue_manager = MessageBasedAppQueueManager(
task_id=application_generate_entity.task_id,
@ -357,7 +384,9 @@ class AgentAppGenerator(MessageBasedAppGenerator):
)
@staticmethod
def _resume_draft_type(*, app_model: App, conversation: Any, user: Account | EndUser) -> str | None:
def _resume_draft_type(
*, app_model: App, conversation: Any, user: Account | EndUser, session: Session
) -> str | None:
if conversation.invoke_from != InvokeFrom.DEBUGGER:
return None
active_session = AgentAppRuntimeSessionStore().load_active_session_for_conversation(
@ -367,7 +396,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
)
snapshot_id = active_session.scope.agent_config_snapshot_id if active_session is not None else None
if snapshot_id and isinstance(user, Account):
draft = db.session.scalar(
draft = session.scalar(
select(AgentConfigDraft).where(
AgentConfigDraft.tenant_id == app_model.tenant_id,
AgentConfigDraft.id == snapshot_id,
@ -413,15 +442,31 @@ class AgentAppGenerator(MessageBasedAppGenerator):
# Apply app-level input guards (content moderation + annotation
# reply) before reaching the Agent backend, mirroring the EasyUI
# chat / agent-chat runners. These can short-circuit the turn.
app_model = db.session.get(App, app_config.app_id)
if app_model is None:
raise AgentAppGeneratorError("App not found")
handled, query = self._run_input_guards(
application_generate_entity=application_generate_entity,
app_model=app_model,
message=message,
queue_manager=queue_manager,
)
with session_factory.get_session_maker().begin() as session:
app_model = session.get(App, app_config.app_id)
if app_model is None:
raise AgentAppGeneratorError("App not found")
handled, query, annotation_reply = self._run_input_guards(
session=session,
application_generate_entity=application_generate_entity,
app_model=app_model,
message=message,
queue_manager=queue_manager,
)
if annotation_reply:
from core.app.apps.agent_app.app_runner import publish_text_answer
from core.app.entities.queue_entities import QueueAnnotationReplyEvent
queue_manager.publish(
QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id),
PublishFrom.APPLICATION_MANAGER,
)
publish_text_answer(
queue_manager=queue_manager,
model_name=application_generate_entity.model_conf.model,
answer=annotation_reply.content,
user_query=query,
)
if handled:
return
query = _append_prompt_file_mappings(
@ -436,11 +481,13 @@ class AgentAppGenerator(MessageBasedAppGenerator):
user_from=user_from,
invoke_from=application_generate_entity.invoke_from,
)
_, _, agent_soul = self._resolve_agent_by_id(
tenant_id=app_config.tenant_id,
agent_id=application_generate_entity.agent_id,
snapshot_id=application_generate_entity.agent_config_snapshot_id,
)
with session_factory.create_session() as session:
_, _, agent_soul = self._resolve_agent_by_id(
tenant_id=app_config.tenant_id,
agent_id=application_generate_entity.agent_id,
snapshot_id=application_generate_entity.agent_config_snapshot_id,
session=session,
)
runner = self._build_runner(dify_context)
runner.run(
@ -502,20 +549,18 @@ class AgentAppGenerator(MessageBasedAppGenerator):
def _run_input_guards(
self,
*,
session: Session,
application_generate_entity: AgentAppGenerateEntity,
app_model: App,
message: Message,
queue_manager: AppQueueManager,
) -> tuple[bool, str]:
) -> tuple[bool, str, MessageAnnotation | None]:
"""Apply input moderation + annotation reply before the backend call.
Returns ``(handled, query)``: when ``handled`` is True a direct answer
has already been published (a blocked/preset moderation response or a
matched annotation) and the backend turn must be skipped. Otherwise
``query`` is the possibly moderation-overridden query to send onward.
Returns ``(handled, query, annotation_reply)``. Annotation output is
published by the caller only after this transaction commits.
"""
from core.app.apps.agent_app.app_runner import publish_text_answer
from core.app.entities.queue_entities import QueueAnnotationReplyEvent
from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
from core.moderation.base import ModerationError
from core.moderation.input_moderation import InputModeration
@ -538,7 +583,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
)
except ModerationError as e:
publish_text_answer(queue_manager=queue_manager, model_name=model_name, answer=str(e), user_query=query)
return True, query
return True, query, None
# annotation reply: a matching annotation answers the turn deterministically.
if query:
@ -548,21 +593,12 @@ class AgentAppGenerator(MessageBasedAppGenerator):
query=query,
user_id=application_generate_entity.user_id,
invoke_from=application_generate_entity.invoke_from,
session=session,
)
if annotation_reply:
queue_manager.publish(
QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id),
PublishFrom.APPLICATION_MANAGER,
)
publish_text_answer(
queue_manager=queue_manager,
model_name=model_name,
answer=annotation_reply.content,
user_query=query,
)
return True, query
return True, query, annotation_reply
return False, query
return False, query, None
def _resolve_agent(
self,
@ -571,8 +607,9 @@ class AgentAppGenerator(MessageBasedAppGenerator):
invoke_from: InvokeFrom,
draft_type: Any,
user: Account | EndUser,
session: Session,
) -> tuple[Agent, str, Literal["snapshot", "draft", "build_draft"], AgentSoulConfig]:
agent = db.session.scalar(
agent = session.scalar(
select(Agent)
.where(
Agent.tenant_id == app_model.tenant_id,
@ -603,6 +640,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
agent=agent,
draft_type=draft_type,
account_id=user.id if isinstance(user, Account) else None,
session=session,
)
agent_soul = AgentSoulConfig.model_validate(draft.config_snapshot_dict)
config_version_kind: Literal["snapshot", "draft", "build_draft"] = (
@ -617,6 +655,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
tenant_id=app_model.tenant_id,
agent_id=agent.id,
snapshot_id=agent.active_config_snapshot_id,
session=session,
)
return agent, snapshot.id, "snapshot", agent_soul
@ -633,7 +672,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
@staticmethod
def _resolve_debug_draft(
*, tenant_id: str, agent: Agent, draft_type: Any, account_id: str | None
*, tenant_id: str, agent: Agent, draft_type: Any, account_id: str | None, session: Session
) -> AgentConfigDraft:
effective_draft_type = (
AgentConfigDraftType.DEBUG_BUILD
@ -651,7 +690,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
stmt = stmt.where(AgentConfigDraft.account_id == account_id)
else:
stmt = stmt.where(AgentConfigDraft.account_id.is_(None))
draft = db.session.scalar(stmt.order_by(AgentConfigDraft.updated_at.desc()).limit(1))
draft = session.scalar(stmt.order_by(AgentConfigDraft.updated_at.desc()).limit(1))
if draft is not None:
return draft
if effective_draft_type == AgentConfigDraftType.DEBUG_BUILD:
@ -660,6 +699,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
tenant_id=tenant_id,
agent_id=agent.id,
snapshot_id=agent.active_config_snapshot_id,
session=session,
)
draft = AgentConfigDraft(
tenant_id=tenant_id,
@ -672,20 +712,20 @@ class AgentAppGenerator(MessageBasedAppGenerator):
created_by=agent.created_by,
updated_by=agent.updated_by,
)
db.session.add(draft)
db.session.flush()
session.add(draft)
session.flush()
return draft
@staticmethod
def _resolve_agent_by_id(
*, tenant_id: str, agent_id: str, snapshot_id: str | None
*, tenant_id: str, agent_id: str, snapshot_id: str | None, session: Session
) -> tuple[Agent, AgentConfigSnapshot | AgentConfigDraft, AgentSoulConfig]:
agent = db.session.scalar(select(Agent).where(Agent.id == agent_id, Agent.tenant_id == tenant_id))
agent = session.scalar(select(Agent).where(Agent.id == agent_id, Agent.tenant_id == tenant_id))
if agent is None:
raise AgentAppGeneratorError("Agent not found")
if not snapshot_id:
raise AgentAppGeneratorError("Agent has no published version")
snapshot = db.session.scalar(
snapshot = session.scalar(
select(AgentConfigSnapshot).where(
AgentConfigSnapshot.tenant_id == tenant_id,
AgentConfigSnapshot.agent_id == agent_id,
@ -695,7 +735,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
if snapshot is not None:
agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict)
return agent, snapshot, agent_soul
draft = db.session.scalar(
draft = session.scalar(
select(AgentConfigDraft).where(
AgentConfigDraft.tenant_id == tenant_id,
AgentConfigDraft.agent_id == agent_id,

18
api/core/app/apps/agent_chat/app_config_manager.py
View File

@ -2,6 +2,8 @@ import uuid
from collections.abc import Mapping
from typing import Any, cast
from sqlalchemy.orm import Session
from core.agent.entities import AgentEntity
from core.app.app_config.base_app_config_manager import BaseAppConfigManager
from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager
@ -20,7 +22,7 @@ from core.app.app_config.features.suggested_questions_after_answer.manager impor
)
from core.app.app_config.features.text_to_speech.manager import TextToSpeechConfigManager
from core.entities.agent_entities import PlanningStrategy
from models.model import App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
from models.model import AnnotationReplyConfig, App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
OLD_TOOLS = ["dataset", "google_search", "web_reader", "wikipedia", "current_datetime"]
@ -41,6 +43,8 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
app_model_config: AppModelConfig,
conversation: Conversation | None = None,
override_config_dict: AppModelConfigDict | None = None,
*,
annotation_reply: AnnotationReplyConfig | None,
) -> AgentChatAppConfig:
"""
Convert app model config to agent chat app config
@ -58,7 +62,7 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
config_from = EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG
if config_from != EasyUIBasedAppModelConfigFrom.ARGS:
app_model_config_dict = app_model_config.to_dict()
app_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply)
config_dict = app_model_config_dict.copy()
else:
if not override_config_dict:
@ -88,7 +92,7 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
return app_config
@classmethod
def config_validate(cls, tenant_id: str, config: Mapping[str, Any]) -> AppModelConfigDict:
def config_validate(cls, tenant_id: str, config: Mapping[str, Any], session: Session) -> AppModelConfigDict:
"""
Validate for agent chat app model config
@ -116,7 +120,7 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
related_config_keys.extend(current_related_config_keys)
# agent_mode
config, current_related_config_keys = cls.validate_agent_mode_and_set_defaults(tenant_id, config)
config, current_related_config_keys = cls.validate_agent_mode_and_set_defaults(tenant_id, config, session)
related_config_keys.extend(current_related_config_keys)
# opening_statement
@ -144,7 +148,7 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
# dataset configs
# dataset_query_variable
config, current_related_config_keys = DatasetConfigManager.validate_and_set_defaults(
tenant_id, app_mode, config
tenant_id, app_mode, config, session
)
related_config_keys.extend(current_related_config_keys)
@ -163,7 +167,7 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
@classmethod
def validate_agent_mode_and_set_defaults(
cls, tenant_id: str, config: dict[str, Any]
cls, tenant_id: str, config: dict[str, Any], session: Session
) -> tuple[dict[str, Any], list[str]]:
"""
Validate agent_mode and set defaults for agent feature
@ -220,7 +224,7 @@ class AgentChatAppConfigManager(BaseAppConfigManager):
except ValueError:
raise ValueError("id in dataset must be of UUID type")
if not DatasetConfigManager.is_dataset_exists(tenant_id, tool_item["id"]):
if not DatasetConfigManager.is_dataset_exists(tenant_id, tool_item["id"], session):
raise ValueError("Dataset ID does not exist, please check your permission.")
else:
# latest style, use key-value pair

48
api/core/app/apps/agent_chat/app_generator.py
View File

@ -21,13 +21,14 @@ from core.app.apps.exc import GenerateTaskStoppedError
from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, InvokeFrom
from core.db.session_factory import session_factory
from core.helper.trace_id_helper import extract_trace_session_id_from_args
from core.ops.ops_trace_manager import TraceQueueManager
from extensions.ext_database import db
from factories import file_factory
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
from libs.flask_utils import preserve_flask_contexts
from models import Account, App, EndUser
from models.model import load_annotation_reply_config
from services.conversation_service import ConversationService
logger = logging.getLogger(__name__)
@ -43,6 +44,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: Literal[False],
session: Session,
) -> Mapping[str, Any]: ...
@overload
@ -54,6 +56,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: Literal[True],
session: Session,
) -> Generator[Mapping | str, None, None]: ...
@overload
@ -65,6 +68,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: bool,
session: Session,
) -> Mapping | Generator[Mapping | str, None, None]: ...
def generate(
@ -75,6 +79,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: bool = True,
session: Session,
) -> Mapping | Generator[Mapping | str, None, None]:
"""
Generate App response.
@ -108,10 +113,14 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
conversation_id = args.get("conversation_id")
if conversation_id:
conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user, session=db.session()
app_model=app_model, conversation_id=conversation_id, user=user, session=session
)
# get app model config
app_model_config = self._get_app_model_config(app_model=app_model, conversation=conversation)
app_model_config = self._get_app_model_config(
app_model=app_model,
conversation=conversation,
session=session,
)
# validate override model config
override_model_config_dict = None
@ -123,11 +132,16 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
override_model_config_dict = AgentChatAppConfigManager.config_validate(
tenant_id=app_model.tenant_id,
config=args["model_config"],
session=session,
)
# always enable retriever resource in debugger mode
override_model_config_dict["retriever_resource"] = {"enabled": True}
annotation_reply = (
None if override_model_config_dict else load_annotation_reply_config(session, app_model_config.app_id)
)
# parse files
# TODO(QuantumGhost): Move file parsing logic to the API controller layer
# for better separation of concerns.
@ -137,7 +151,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from):
files = args.get("files") or []
file_extra_config = FileUploadConfigManager.convert(
override_model_config_dict or app_model_config.to_dict()
override_model_config_dict or app_model_config.to_dict(annotation_reply=annotation_reply)
)
if file_extra_config:
file_objs = file_factory.build_from_mappings(
@ -155,6 +169,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
app_model_config=app_model_config,
conversation=conversation,
override_config_dict=override_model_config_dict,
annotation_reply=annotation_reply,
)
# get tracing instance
@ -186,7 +201,11 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
)
# init generate records
(conversation, message) = self._init_generate_records(application_generate_entity, conversation)
(conversation, message) = self._init_generate_records(
application_generate_entity,
conversation,
session=session,
)
# init queue manager
queue_manager = MessageBasedAppQueueManager(
@ -205,7 +224,6 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
target=self._generate_worker,
kwargs={
"flask_app": current_app._get_current_object(), # type: ignore
"session": db.session(),
"context": context,
"application_generate_entity": application_generate_entity,
"queue_manager": queue_manager,
@ -230,7 +248,6 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
def _generate_worker(
self,
flask_app: Flask,
session: Session,
context: contextvars.Context,
application_generate_entity: AgentChatAppGenerateEntity,
queue_manager: AppQueueManager,
@ -255,13 +272,14 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
# chatbot app
runner = AgentChatAppRunner()
runner.run(
session=session,
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
conversation=conversation,
message=message,
)
with session_factory.create_session() as session:
runner.run(
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
conversation=conversation,
message=message,
session=session,
)
except GenerateTaskStoppedError:
pass
except InvokeAuthorizationError:
@ -278,5 +296,3 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
except Exception as e:
logger.exception("Unknown Error when generating")
queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER)
finally:
db.session.close()

31
api/core/app/apps/agent_chat/app_runner.py
View File

@ -32,14 +32,17 @@ class AgentChatAppRunner(AppRunner):
def run(
self,
session: Session,
application_generate_entity: AgentChatAppGenerateEntity,
queue_manager: AppQueueManager,
conversation: Conversation,
message: Message,
session: Session,
):
"""
Run assistant application
"""Run the assistant application with bounded explicit transactions.
The setup session is committed and released before the multi-step agent
runner begins model or tool I/O.
:param application_generate_entity: application generate entity
:param queue_manager: application queue manager
:param conversation: conversation
@ -49,10 +52,10 @@ class AgentChatAppRunner(AppRunner):
app_config = application_generate_entity.app_config
app_config = cast(AgentChatAppConfig, app_config)
app_stmt = select(App).where(App.id == app_config.app_id)
with create_session() as session:
app_record = session.scalar(app_stmt)
with create_session() as read_session:
app_record = read_session.scalar(app_stmt)
if app_record:
session.expunge(app_record)
read_session.expunge(app_record)
if not app_record:
raise ValueError("App not found")
@ -112,7 +115,10 @@ class AgentChatAppRunner(AppRunner):
query=query,
user_id=application_generate_entity.user_id,
invoke_from=application_generate_entity.invoke_from,
session=session,
)
session.commit()
session.close()
if annotation_reply:
queue_manager.publish(
@ -191,15 +197,15 @@ class AgentChatAppRunner(AppRunner):
agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING
conversation_stmt = select(Conversation).where(Conversation.id == conversation.id)
msg_stmt = select(Message).where(Message.id == message.id)
with create_session() as session:
conversation_result = session.scalar(conversation_stmt)
with create_session() as read_session:
conversation_result = read_session.scalar(conversation_stmt)
if conversation_result is None:
raise ValueError("Conversation not found")
message_result = session.scalar(msg_stmt)
message_result = read_session.scalar(msg_stmt)
if message_result is not None:
session.expunge(message_result)
session.expunge(conversation_result)
read_session.expunge(message_result)
read_session.expunge(conversation_result)
if message_result is None:
raise ValueError("Message not found")
@ -234,6 +240,9 @@ class AgentChatAppRunner(AppRunner):
model_instance=model_instance,
)
session.commit()
session.close()
invoke_result = runner.run(
session=session,
message=message,

63
api/core/app/apps/base_app_runner.py
View File

@ -5,7 +5,7 @@ from collections.abc import Generator, Mapping, Sequence
from mimetypes import guess_extension
from typing import TYPE_CHECKING, Any, Union
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session
from core.app.app_config.entities import ExternalDataVariableEntity, PromptTemplateEntity
from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
@ -24,6 +24,7 @@ from core.app.entities.queue_entities import (
)
from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
from core.app.features.hosting_moderation.hosting_moderation import HostingModerationFeature
from core.db.session_factory import session_factory
from core.external_data_tool.external_data_fetch import ExternalDataFetch
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_manager import ModelInstance
@ -32,7 +33,6 @@ from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
from core.prompt.simple_prompt_transform import ModelMode, SimplePromptTransform
from core.tools.tool_file_manager import ToolFileManager
from extensions.ext_database import db
from graphon.file import FileTransferMethod, FileType
from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
from graphon.model_runtime.entities.message_entities import (
@ -310,20 +310,30 @@ class AppRunner:
case list():
for content in message.content:
match content:
case str():
text += content
case TextPromptMessageContent():
text += content.data
case ImagePromptMessageContent():
if message_id and user_id and tenant_id:
try:
self._handle_multimodal_image_content(
content=content,
message_id=message_id,
user_id=user_id,
tenant_id=tenant_id,
queue_manager=queue_manager,
)
with session_factory.create_session() as session:
message_file_id = self._handle_multimodal_image_content(
session=session,
content=content,
message_id=message_id,
user_id=user_id,
tenant_id=tenant_id,
queue_manager=queue_manager,
)
session.commit()
if message_file_id:
queue_manager.publish(
QueueMessageFileEvent(message_file_id=message_file_id),
PublishFrom.APPLICATION_MANAGER,
)
_logger.info(
"QueueMessageFileEvent published for message_file_id: %s",
message_file_id,
)
except Exception:
_logger.exception("Failed to handle multimodal image output")
else:
@ -365,7 +375,8 @@ class AppRunner:
user_id: str,
tenant_id: str,
queue_manager: AppQueueManager,
):
session: Session,
) -> str | None:
"""
Handle multimodal image content from LLM response.
Save the image and create a MessageFile record.
@ -386,7 +397,7 @@ class AppRunner:
if not image_url and not base64_data:
_logger.warning("Image content has neither URL nor base64 data")
return
return None
tool_file_manager = ToolFileManager()
@ -420,10 +431,10 @@ class AppRunner:
)
_logger.info("Image saved successfully, tool_file_id: %s", tool_file.id)
else:
return
return None
except Exception:
_logger.exception("Failed to save image file")
return
return None
# Create MessageFile record.
# Use an independent session so this side-effect write does not
@ -441,16 +452,9 @@ class AppRunner:
created_by=user_id,
)
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
session.add(message_file)
# Publish QueueMessageFileEvent
queue_manager.publish(
QueueMessageFileEvent(message_file_id=message_file.id),
PublishFrom.APPLICATION_MANAGER,
)
_logger.info("QueueMessageFileEvent published for message_file_id: %s", message_file.id)
session.add(message_file)
session.flush()
return message_file.id
def moderation_for_inputs(
self,
@ -536,7 +540,7 @@ class AppRunner:
)
def query_app_annotations_to_reply(
self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom
self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom, session: Session
) -> MessageAnnotation | None:
"""
Query app annotations to reply
@ -549,5 +553,10 @@ class AppRunner:
"""
annotation_reply_feature = AnnotationReplyFeature()
return annotation_reply_feature.query(
app_record=app_record, message=message, query=query, user_id=user_id, invoke_from=invoke_from
app_record=app_record,
message=message,
query=query,
user_id=user_id,
invoke_from=invoke_from,
session=session,
)

12
api/core/app/apps/chat/app_config_manager.py
View File

@ -1,5 +1,7 @@
from typing import Any, cast
from sqlalchemy.orm import Session
from core.app.app_config.base_app_config_manager import BaseAppConfigManager
from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager
from core.app.app_config.easy_ui_based_app.dataset.manager import DatasetConfigManager
@ -15,7 +17,7 @@ from core.app.app_config.features.suggested_questions_after_answer.manager impor
SuggestedQuestionsAfterAnswerConfigManager,
)
from core.app.app_config.features.text_to_speech.manager import TextToSpeechConfigManager
from models.model import App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
from models.model import AnnotationReplyConfig, App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
class ChatAppConfig(EasyUIBasedAppConfig):
@ -34,6 +36,8 @@ class ChatAppConfigManager(BaseAppConfigManager):
app_model_config: AppModelConfig,
conversation: Conversation | None = None,
override_config_dict: AppModelConfigDict | None = None,
*,
annotation_reply: AnnotationReplyConfig | None,
) -> ChatAppConfig:
"""
Convert app model config to chat app config
@ -51,7 +55,7 @@ class ChatAppConfigManager(BaseAppConfigManager):
config_from = EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG
if config_from != EasyUIBasedAppModelConfigFrom.ARGS:
app_model_config_dict = app_model_config.to_dict()
app_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply)
config_dict = app_model_config_dict.copy()
else:
if not override_config_dict:
@ -81,7 +85,7 @@ class ChatAppConfigManager(BaseAppConfigManager):
return app_config
@classmethod
def config_validate(cls, tenant_id: str, config: dict[str, Any]) -> AppModelConfigDict:
def config_validate(cls, tenant_id: str, config: dict[str, Any], session: Session) -> AppModelConfigDict:
"""
Validate for chat app model config
@ -110,7 +114,7 @@ class ChatAppConfigManager(BaseAppConfigManager):
# dataset_query_variable
config, current_related_config_keys = DatasetConfigManager.validate_and_set_defaults(
tenant_id, app_mode, config
tenant_id, app_mode, config, session=session
)
related_config_keys.extend(current_related_config_keys)

55
api/core/app/apps/chat/app_generator.py
View File

@ -21,13 +21,14 @@ from core.app.apps.exc import GenerateTaskStoppedError
from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom
from core.db.session_factory import session_factory
from core.helper.trace_id_helper import extract_trace_session_id_from_args
from core.ops.ops_trace_manager import TraceQueueManager
from extensions.ext_database import db
from factories import file_factory
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
from models import Account
from models.model import App, EndUser
from models.model import App, EndUser, load_annotation_reply_config
from services.conversation_service import ConversationService
logger = logging.getLogger(__name__)
@ -37,44 +38,48 @@ class ChatAppGenerator(MessageBasedAppGenerator):
@overload
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: Literal[True],
*,
session: Session,
) -> Generator[Mapping | str, None, None]: ...
@overload
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: Literal[False],
*,
session: Session,
) -> Mapping[str, Any]: ...
@overload
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: bool,
*,
session: Session,
) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]: ...
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]:
"""
Generate App response.
@ -105,10 +110,14 @@ class ChatAppGenerator(MessageBasedAppGenerator):
conversation_id = args.get("conversation_id")
if conversation_id:
conversation = ConversationService.get_conversation(
app_model=app_model, conversation_id=conversation_id, user=user, session=db.session()
app_model=app_model, conversation_id=conversation_id, user=user, session=session
)
# get app model config
app_model_config = self._get_app_model_config(app_model=app_model, conversation=conversation)
app_model_config = self._get_app_model_config(
app_model=app_model,
conversation=conversation,
session=session,
)
# validate override model config
override_model_config_dict = None
@ -118,12 +127,16 @@ class ChatAppGenerator(MessageBasedAppGenerator):
# validate config
override_model_config_dict = ChatAppConfigManager.config_validate(
tenant_id=app_model.tenant_id, config=args.get("model_config", {})
tenant_id=app_model.tenant_id, config=args.get("model_config", {}), session=session
)
# always enable retriever resource in debugger mode
override_model_config_dict["retriever_resource"] = {"enabled": True}
annotation_reply = (
None if override_model_config_dict else load_annotation_reply_config(session, app_model_config.app_id)
)
# parse files
# TODO(QuantumGhost): Move file parsing logic to the API controller layer
# for better separation of concerns.
@ -133,7 +146,7 @@ class ChatAppGenerator(MessageBasedAppGenerator):
with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from):
files = args["files"] if args.get("files") else []
file_extra_config = FileUploadConfigManager.convert(
override_model_config_dict or app_model_config.to_dict()
override_model_config_dict or app_model_config.to_dict(annotation_reply=annotation_reply)
)
if file_extra_config:
file_objs = file_factory.build_from_mappings(
@ -151,6 +164,7 @@ class ChatAppGenerator(MessageBasedAppGenerator):
app_model_config=app_model_config,
conversation=conversation,
override_config_dict=override_model_config_dict,
annotation_reply=annotation_reply,
)
# get tracing instance
@ -183,7 +197,11 @@ class ChatAppGenerator(MessageBasedAppGenerator):
)
# init generate records
(conversation, message) = self._init_generate_records(application_generate_entity, conversation)
(conversation, message) = self._init_generate_records(
application_generate_entity,
conversation,
session=session,
)
# init queue manager
queue_manager = MessageBasedAppQueueManager(
@ -202,7 +220,6 @@ class ChatAppGenerator(MessageBasedAppGenerator):
def worker_with_context():
return context.run(
self._generate_worker,
session=session,
flask_app=current_app._get_current_object(), # type: ignore
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
@ -229,7 +246,6 @@ class ChatAppGenerator(MessageBasedAppGenerator):
def _generate_worker(
self,
flask_app: Flask,
session: Session,
application_generate_entity: ChatAppGenerateEntity,
queue_manager: AppQueueManager,
conversation_id: str,
@ -252,13 +268,14 @@ class ChatAppGenerator(MessageBasedAppGenerator):
# chatbot app
runner = ChatAppRunner()
runner.run(
session=session,
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
conversation=conversation,
message=message,
)
with session_factory.create_session() as session:
runner.run(
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
conversation=conversation,
message=message,
session=session,
)
except GenerateTaskStoppedError:
pass
except InvokeAuthorizationError:

26
api/core/app/apps/chat/app_runner.py
View File

@ -17,7 +17,6 @@ from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_manager import ModelInstance
from core.moderation.base import ModerationError
from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
from extensions.ext_database import db
from graphon.file import File
from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent
from models.model import App, Conversation, Message
@ -32,14 +31,17 @@ class ChatAppRunner(AppRunner):
def run(
self,
session: Session,
application_generate_entity: ChatAppGenerateEntity,
queue_manager: AppQueueManager,
conversation: Conversation,
message: Message,
session: Session,
):
"""
Run application
"""Run the application without retaining ``session`` during model I/O.
Database preparation is committed and the connection is released before
the provider response is requested or consumed.
:param application_generate_entity: application generate entity
:param queue_manager: application queue manager
:param conversation: conversation
@ -49,10 +51,10 @@ class ChatAppRunner(AppRunner):
app_config = application_generate_entity.app_config
app_config = cast(ChatAppConfig, app_config)
stmt = select(App).where(App.id == app_config.app_id)
with create_session() as session:
app_record = session.scalar(stmt)
with create_session() as read_session:
app_record = read_session.scalar(stmt)
if app_record:
session.expunge(app_record)
read_session.expunge(app_record)
if not app_record:
raise ValueError("App not found")
@ -123,7 +125,10 @@ class ChatAppRunner(AppRunner):
query=query,
user_id=application_generate_entity.user_id,
invoke_from=application_generate_entity.invoke_from,
session=session,
)
session.commit()
session.close()
if annotation_reply:
queue_manager.publish(
@ -188,6 +193,9 @@ class ChatAppRunner(AppRunner):
)
context_files = retrieved_files or []
session.commit()
session.close()
# reorganize all inputs and template to prompt messages
# Include: prompt template, inputs, query(optional), files(optional)
# memory(optional), external data, dataset context(optional)
@ -223,10 +231,6 @@ class ChatAppRunner(AppRunner):
model=application_generate_entity.model_conf.model,
)
# Release the Flask scoped session before LLM streaming so a checked-out DB connection
# is not held for the lifetime of the provider response.
db.session.close()
invoke_result = model_instance.invoke_llm(
prompt_messages=prompt_messages,
model_parameters=application_generate_entity.model_conf.parameters,

17
api/core/app/apps/completion/app_config_manager.py
View File

@ -1,5 +1,7 @@
from typing import Any, cast
from sqlalchemy.orm import Session
from core.app.app_config.base_app_config_manager import BaseAppConfigManager
from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager
from core.app.app_config.easy_ui_based_app.dataset.manager import DatasetConfigManager
@ -10,7 +12,7 @@ from core.app.app_config.entities import EasyUIBasedAppConfig, EasyUIBasedAppMod
from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
from core.app.app_config.features.more_like_this.manager import MoreLikeThisConfigManager
from core.app.app_config.features.text_to_speech.manager import TextToSpeechConfigManager
from models.model import App, AppMode, AppModelConfig, AppModelConfigDict
from models.model import AnnotationReplyConfig, App, AppMode, AppModelConfig, AppModelConfigDict
class CompletionAppConfig(EasyUIBasedAppConfig):
@ -24,7 +26,12 @@ class CompletionAppConfig(EasyUIBasedAppConfig):
class CompletionAppConfigManager(BaseAppConfigManager):
@classmethod
def get_app_config(
cls, app_model: App, app_model_config: AppModelConfig, override_config_dict: AppModelConfigDict | None = None
cls,
app_model: App,
app_model_config: AppModelConfig,
override_config_dict: AppModelConfigDict | None = None,
*,
annotation_reply: AnnotationReplyConfig | None,
) -> CompletionAppConfig:
"""
Convert app model config to completion app config
@ -39,7 +46,7 @@ class CompletionAppConfigManager(BaseAppConfigManager):
config_from = EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG
if config_from != EasyUIBasedAppModelConfigFrom.ARGS:
app_model_config_dict = app_model_config.to_dict()
app_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply)
config_dict = app_model_config_dict.copy()
else:
if not override_config_dict:
@ -68,7 +75,7 @@ class CompletionAppConfigManager(BaseAppConfigManager):
return app_config
@classmethod
def config_validate(cls, tenant_id: str, config: dict[str, Any]) -> AppModelConfigDict:
def config_validate(cls, tenant_id: str, config: dict[str, Any], session: Session) -> AppModelConfigDict:
"""
Validate for completion app model config
@ -97,7 +104,7 @@ class CompletionAppConfigManager(BaseAppConfigManager):
# dataset_query_variable
config, current_related_config_keys = DatasetConfigManager.validate_and_set_defaults(
tenant_id, app_mode, config
tenant_id, app_mode, config, session
)
related_config_keys.extend(current_related_config_keys)

89
api/core/app/apps/completion/app_generator.py
View File

@ -21,12 +21,14 @@ from core.app.apps.exc import GenerateTaskStoppedError
from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
from core.app.entities.app_invoke_entities import CompletionAppGenerateEntity, InvokeFrom
from core.db.session_factory import session_factory
from core.helper.trace_id_helper import extract_trace_session_id_from_args
from core.ops.ops_trace_manager import TraceQueueManager
from extensions.ext_database import db
from factories import file_factory
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
from models import Account, App, EndUser, Message
from models import Account, App, AppModelConfig, Conversation, EndUser, Message
from models.model import load_annotation_reply_config
from services.errors.app import MoreLikeThisDisabledError
from services.errors.message import MessageNotExistsError
@ -37,44 +39,48 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
@overload
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: Literal[True],
*,
session: Session,
) -> Generator[str | Mapping[str, Any], None, None]: ...
@overload
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: Literal[False],
*,
session: Session,
) -> Mapping[str, Any]: ...
@overload
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: bool = False,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]: ...
def generate(
self,
session: Session,
app_model: App,
user: Account | EndUser,
args: Mapping[str, Any],
invoke_from: InvokeFrom,
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -96,7 +102,11 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
conversation = None
# get app model config
app_model_config = self._get_app_model_config(app_model=app_model, conversation=conversation)
app_model_config = self._get_app_model_config(
app_model=app_model,
conversation=conversation,
session=session,
)
# validate override model config
override_model_config_dict = None
@ -106,9 +116,13 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
# validate config
override_model_config_dict = CompletionAppConfigManager.config_validate(
tenant_id=app_model.tenant_id, config=args.get("model_config", {})
tenant_id=app_model.tenant_id, config=args.get("model_config", {}), session=session
)
annotation_reply = (
None if override_model_config_dict else load_annotation_reply_config(session, app_model_config.app_id)
)
# parse files
# TODO(QuantumGhost): Move file parsing logic to the API controller layer
# for better separation of concerns.
@ -118,7 +132,7 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from):
files = args["files"] if args.get("files") else []
file_extra_config = FileUploadConfigManager.convert(
override_model_config_dict or app_model_config.to_dict()
override_model_config_dict or app_model_config.to_dict(annotation_reply=annotation_reply)
)
if file_extra_config:
file_objs = file_factory.build_from_mappings(
@ -132,7 +146,10 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
# convert to app config
app_config = CompletionAppConfigManager.get_app_config(
app_model=app_model, app_model_config=app_model_config, override_config_dict=override_model_config_dict
app_model=app_model,
app_model_config=app_model_config,
override_config_dict=override_model_config_dict,
annotation_reply=annotation_reply,
)
# get tracing instance
@ -161,7 +178,10 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
)
# init generate records
(conversation, message) = self._init_generate_records(application_generate_entity)
(conversation, message) = self._init_generate_records(
application_generate_entity,
session=session,
)
# init queue manager
queue_manager = MessageBasedAppQueueManager(
@ -180,7 +200,6 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
def worker_with_context():
return context.run(
self._generate_worker,
session=session,
flask_app=current_app._get_current_object(), # type: ignore
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
@ -206,7 +225,6 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
def _generate_worker(
self,
flask_app: Flask,
session: Session,
application_generate_entity: CompletionAppGenerateEntity,
queue_manager: AppQueueManager,
message_id: str,
@ -226,12 +244,13 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
# chatbot app
runner = CompletionAppRunner()
runner.run(
session=session,
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
message=message,
)
with session_factory.create_session() as session:
runner.run(
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,
message=message,
session=session,
)
except GenerateTaskStoppedError:
pass
except InvokeAuthorizationError:
@ -253,12 +272,13 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
def generate_more_like_this(
self,
session: Session,
app_model: App,
message_id: str,
user: Account | EndUser,
invoke_from: InvokeFrom,
stream: bool = True,
*,
session: Session,
) -> Mapping | Generator[Mapping | str, None, None]:
"""
Generate App response.
@ -276,12 +296,14 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
Message.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
Message.from_account_id == (user.id if isinstance(user, Account) else None),
)
message = db.session.scalar(stmt)
message = session.scalar(stmt)
if not message:
raise MessageNotExistsError()
current_app_model_config = app_model.app_model_config
current_app_model_config = (
session.get(AppModelConfig, app_model.app_model_config_id) if app_model.app_model_config_id else None
)
if not current_app_model_config:
raise MoreLikeThisDisabledError()
@ -290,10 +312,16 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
if not current_app_model_config.more_like_this or more_like_this.get("enabled", False) is False:
raise MoreLikeThisDisabledError()
app_model_config = message.app_model_config
conversation = session.get(Conversation, message.conversation_id) if message.conversation_id else None
app_model_config = (
session.get(AppModelConfig, conversation.app_model_config_id)
if conversation and conversation.app_model_config_id
else None
)
if not app_model_config:
raise ValueError("Message app_model_config is None")
override_model_config_dict = app_model_config.to_dict()
annotation_reply = load_annotation_reply_config(session, app_model_config.app_id)
override_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply)
model_dict = override_model_config_dict["model"]
completion_params = model_dict.get("completion_params", {})
completion_params["temperature"] = 0.9
@ -305,7 +333,7 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
file_extra_config = FileUploadConfigManager.convert(override_model_config_dict)
if file_extra_config:
file_objs = file_factory.build_from_mappings(
mappings=message.message_files,
mappings=message.message_files_with_session(session=session),
tenant_id=app_model.tenant_id,
config=file_extra_config,
access_controller=self._file_access_controller,
@ -315,7 +343,10 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
# convert to app config
app_config = CompletionAppConfigManager.get_app_config(
app_model=app_model, app_model_config=app_model_config, override_config_dict=override_model_config_dict
app_model=app_model,
app_model_config=app_model_config,
override_config_dict=override_model_config_dict,
annotation_reply=annotation_reply,
)
# init application generate entity
@ -323,7 +354,7 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
task_id=str(uuid.uuid4()),
app_config=app_config,
model_conf=ModelConfigConverter.convert(app_config),
inputs=message.inputs,
inputs=message.inputs_with_session(session=session),
query=message.query,
files=list(file_objs),
user_id=user.id,
@ -333,7 +364,10 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
)
# init generate records
(conversation, message) = self._init_generate_records(application_generate_entity)
(conversation, message) = self._init_generate_records(
application_generate_entity,
session=session,
)
# init queue manager
queue_manager = MessageBasedAppQueueManager(
@ -352,7 +386,6 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
def worker_with_context():
return context.run(
self._generate_worker,
session=session,
flask_app=current_app._get_current_object(), # type: ignore
application_generate_entity=application_generate_entity,
queue_manager=queue_manager,

23
api/core/app/apps/completion/app_runner.py
View File

@ -15,7 +15,6 @@ from core.db.session_factory import create_session
from core.model_manager import ModelInstance
from core.moderation.base import ModerationError
from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
from extensions.ext_database import db
from graphon.file import File
from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent
from models.model import App, Message
@ -30,13 +29,16 @@ class CompletionAppRunner(AppRunner):
def run(
self,
session: Session,
application_generate_entity: CompletionAppGenerateEntity,
queue_manager: AppQueueManager,
message: Message,
session: Session,
):
"""
Run application
"""Run the application without retaining ``session`` during model I/O.
Database preparation is committed and the connection is released before
the provider response is requested or consumed.
:param application_generate_entity: application generate entity
:param queue_manager: application queue manager
:param message: message
@ -45,10 +47,10 @@ class CompletionAppRunner(AppRunner):
app_config = application_generate_entity.app_config
app_config = cast(CompletionAppConfig, app_config)
stmt = select(App).where(App.id == app_config.app_id)
with create_session() as session:
app_record = session.scalar(stmt)
with create_session() as read_session:
app_record = read_session.scalar(stmt)
if app_record:
session.expunge(app_record)
read_session.expunge(app_record)
if not app_record:
raise ValueError("App not found")
@ -150,6 +152,9 @@ class CompletionAppRunner(AppRunner):
)
context_files = retrieved_files or []
session.commit()
session.close()
# reorganize all inputs and template to prompt messages
# Include: prompt template, inputs, query(optional), files(optional)
# memory(optional), external data, dataset context(optional)
@ -184,10 +189,6 @@ class CompletionAppRunner(AppRunner):
model=application_generate_entity.model_conf.model,
)
# Release the Flask scoped session before LLM streaming so a checked-out DB connection
# is not held for the lifetime of the provider response.
db.session.close()
invoke_result = model_instance.invoke_llm(
prompt_messages=prompt_messages,
model_parameters=application_generate_entity.model_conf.parameters,

32
api/core/app/apps/message_based_app_generator.py
View File

@ -89,12 +89,18 @@ class MessageBasedAppGenerator(BaseAppGenerator):
logger.exception("Failed to handle response, conversation_id: %s", conversation.id)
raise e
def _get_app_model_config(self, app_model: App, conversation: Conversation | None = None) -> AppModelConfig:
def _get_app_model_config(
self,
app_model: App,
conversation: Conversation | None = None,
*,
session: Session,
) -> AppModelConfig:
if conversation:
stmt = select(AppModelConfig).where(
AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id
)
app_model_config = db.session.scalar(stmt)
app_model_config = session.scalar(stmt)
if not app_model_config:
raise AppModelConfigBrokenError()
@ -102,7 +108,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
if app_model.app_model_config_id is None:
raise AppModelConfigBrokenError()
app_model_config = app_model.app_model_config
app_model_config = session.get(AppModelConfig, app_model.app_model_config_id)
if not app_model_config:
raise AppModelConfigBrokenError()
@ -118,6 +124,8 @@ class MessageBasedAppGenerator(BaseAppGenerator):
AdvancedChatAppGenerateEntity,
],
conversation: Conversation | None = None,
*,
session: Session,
) -> tuple[Conversation, Message]:
"""
Initialize generate records
@ -183,9 +191,9 @@ class MessageBasedAppGenerator(BaseAppGenerator):
from_account_id=account_id,
)
db.session.add(conversation)
db.session.flush()
db.session.refresh(conversation)
session.add(conversation)
session.flush()
session.refresh(conversation)
else:
conversation.updated_at = naive_utc_now()
@ -216,9 +224,9 @@ class MessageBasedAppGenerator(BaseAppGenerator):
app_mode=app_config.app_mode,
)
db.session.add(message)
db.session.flush()
db.session.refresh(message)
session.add(message)
session.flush()
session.refresh(message)
message_files = []
for file in application_generate_entity.files:
@ -235,16 +243,16 @@ class MessageBasedAppGenerator(BaseAppGenerator):
message_files.append(message_file)
if message_files:
db.session.add_all(message_files)
session.add_all(message_files)
db.session.commit()
session.commit()
if isinstance(application_generate_entity, ConversationAppGenerateEntity):
application_generate_entity.conversation_id = conversation.id
application_generate_entity.is_new_conversation = created_new_conversation
return conversation, message
except Exception:
db.session.rollback()
session.rollback()
raise
def _get_conversation_introduction(self, application_generate_entity: AppGenerateEntity) -> str:

50
api/core/app/apps/pipeline/pipeline_generator.py
View File

@ -64,6 +64,7 @@ class PipelineGenerator(BaseAppGenerator):
def generate(
self,
*,
session: Session,
pipeline: Pipeline,
workflow: Workflow,
user: Account | EndUser,
@ -79,6 +80,7 @@ class PipelineGenerator(BaseAppGenerator):
def generate(
self,
*,
session: Session,
pipeline: Pipeline,
workflow: Workflow,
user: Account | EndUser,
@ -94,6 +96,7 @@ class PipelineGenerator(BaseAppGenerator):
def generate(
self,
*,
session: Session,
pipeline: Pipeline,
workflow: Workflow,
user: Account | EndUser,
@ -108,6 +111,7 @@ class PipelineGenerator(BaseAppGenerator):
def generate(
self,
*,
session: Session,
pipeline: Pipeline,
workflow: Workflow,
user: Account | EndUser,
@ -120,10 +124,9 @@ class PipelineGenerator(BaseAppGenerator):
) -> Mapping[str, Any] | Generator[Mapping | str, None, None] | None:
# Add null check for dataset
with Session(db.engine, expire_on_commit=False) as session:
dataset = pipeline.retrieve_dataset(session)
if not dataset:
raise ValueError("Pipeline dataset is required")
dataset = pipeline.retrieve_dataset(session)
if not dataset:
raise ValueError("Pipeline dataset is required")
inputs: Mapping[str, Any] = args["inputs"]
start_node_id: str = args["start_node_id"]
datasource_type = DatasourceProviderType(args["datasource_type"])
@ -157,9 +160,9 @@ class PipelineGenerator(BaseAppGenerator):
batch=batch,
document_form=dataset.chunk_structure,
)
db.session.add(document)
session.add(document)
documents.append(document)
db.session.commit()
session.flush()
# run in child thread
rag_pipeline_invoke_entities = []
@ -177,8 +180,7 @@ class PipelineGenerator(BaseAppGenerator):
pipeline_id=pipeline.id,
created_by=user.id,
)
db.session.add(document_pipeline_execution_log)
db.session.commit()
session.add(document_pipeline_execution_log)
application_generate_entity = RagPipelineGenerateEntity(
task_id=str(uuid.uuid4()),
app_config=pipeline_config,
@ -227,6 +229,7 @@ class PipelineGenerator(BaseAppGenerator):
)
if invoke_from == InvokeFrom.DEBUGGER or is_retry:
return self._generate(
session=session,
flask_app=current_app._get_current_object(), # type: ignore
context=contextvars.copy_context(),
pipeline=pipeline,
@ -253,6 +256,8 @@ class PipelineGenerator(BaseAppGenerator):
)
)
if invoke_from == InvokeFrom.PUBLISHED_PIPELINE and not is_retry:
session.commit()
if rag_pipeline_invoke_entities:
RagPipelineTaskProxy(dataset.tenant_id, user.id, rag_pipeline_invoke_entities).delay()
# return batch, dataset, documents
@ -282,6 +287,7 @@ class PipelineGenerator(BaseAppGenerator):
def _generate(
self,
*,
session: Session,
flask_app: Flask,
context: contextvars.Context,
pipeline: Pipeline,
@ -310,7 +316,7 @@ class PipelineGenerator(BaseAppGenerator):
"""
with preserve_flask_contexts(flask_app, context_vars=context):
# init queue manager
workflow = db.session.get(Workflow, workflow_id)
workflow = session.get(Workflow, workflow_id)
if not workflow:
raise ValueError(f"Workflow not found: {workflow_id}")
queue_manager = PipelineQueueManager(
@ -362,6 +368,8 @@ class PipelineGenerator(BaseAppGenerator):
user: Account | EndUser,
args: Mapping[str, Any],
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -372,6 +380,7 @@ class PipelineGenerator(BaseAppGenerator):
:param user: account or end user
:param args: request args
:param streaming: is streamed
:param session: database session supplied by the caller
"""
if not node_id:
raise ValueError("node_id is required")
@ -384,10 +393,9 @@ class PipelineGenerator(BaseAppGenerator):
pipeline=pipeline, workflow=workflow, start_node_id=args.get("start_node_id", "shared")
)
with Session(db.engine) as session:
dataset = pipeline.retrieve_dataset(session)
if not dataset:
raise ValueError("Pipeline dataset is required")
dataset = pipeline.retrieve_dataset(session)
if not dataset:
raise ValueError("Pipeline dataset is required")
# init application generate entity - use RagPipelineGenerateEntity instead
application_generate_entity = RagPipelineGenerateEntity(
@ -428,7 +436,7 @@ class PipelineGenerator(BaseAppGenerator):
app_id=application_generate_entity.app_config.app_id,
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
)
draft_var_srv = WorkflowDraftVariableService(db.session())
draft_var_srv = WorkflowDraftVariableService(session)
draft_var_srv.prefill_conversation_variable_default_values(workflow, user_id=user.id)
var_loader = DraftVarLoader(
engine=db.engine,
@ -438,6 +446,7 @@ class PipelineGenerator(BaseAppGenerator):
)
return self._generate(
session=session,
flask_app=current_app._get_current_object(), # type: ignore
pipeline=pipeline,
workflow_id=workflow.id,
@ -459,6 +468,8 @@ class PipelineGenerator(BaseAppGenerator):
user: Account | EndUser,
args: Mapping[str, Any],
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -469,6 +480,7 @@ class PipelineGenerator(BaseAppGenerator):
:param user: account or end user
:param args: request args
:param streaming: is streamed
:param session: database session supplied by the caller
"""
if not node_id:
raise ValueError("node_id is required")
@ -476,10 +488,9 @@ class PipelineGenerator(BaseAppGenerator):
if args.get("inputs") is None:
raise ValueError("inputs is required")
with Session(db.engine) as session:
dataset = pipeline.retrieve_dataset(session)
if not dataset:
raise ValueError("Pipeline dataset is required")
dataset = pipeline.retrieve_dataset(session)
if not dataset:
raise ValueError("Pipeline dataset is required")
# convert to app config
pipeline_config = PipelineConfigManager.get_pipeline_config(
@ -524,7 +535,7 @@ class PipelineGenerator(BaseAppGenerator):
app_id=application_generate_entity.app_config.app_id,
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
)
draft_var_srv = WorkflowDraftVariableService(db.session())
draft_var_srv = WorkflowDraftVariableService(session)
draft_var_srv.prefill_conversation_variable_default_values(workflow, user_id=user.id)
var_loader = DraftVarLoader(
engine=db.engine,
@ -534,6 +545,7 @@ class PipelineGenerator(BaseAppGenerator):
)
return self._generate(
session=session,
flask_app=current_app._get_current_object(), # type: ignore
pipeline=pipeline,
workflow_id=workflow.id,

10
api/core/app/apps/workflow/app_generator.py
View File

@ -419,6 +419,8 @@ class WorkflowAppGenerator(BaseAppGenerator):
user: Account | EndUser,
args: Mapping[str, Any],
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -429,6 +431,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
:param user: account or end user
:param args: request args
:param streaming: is streamed
:param session: database session supplied by the caller
"""
if not node_id:
raise ValueError("node_id is required")
@ -478,7 +481,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
app_id=application_generate_entity.app_config.app_id,
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
)
draft_var_srv = WorkflowDraftVariableService(db.session())
draft_var_srv = WorkflowDraftVariableService(session)
draft_var_srv.prefill_conversation_variable_default_values(workflow, user_id=user.id)
var_loader = DraftVarLoader(
engine=db.engine,
@ -508,6 +511,8 @@ class WorkflowAppGenerator(BaseAppGenerator):
user: Account | EndUser,
args: LoopNodeRunPayload,
streaming: bool = True,
*,
session: Session,
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
"""
Generate App response.
@ -518,6 +523,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
:param user: account or end user
:param args: request args
:param streaming: is streamed
:param session: database session supplied by the caller
"""
if not node_id:
raise ValueError("node_id is required")
@ -565,7 +571,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
app_id=application_generate_entity.app_config.app_id,
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
)
draft_var_srv = WorkflowDraftVariableService(db.session())
draft_var_srv = WorkflowDraftVariableService(session)
draft_var_srv.prefill_conversation_variable_default_values(workflow, user_id=user.id)
var_loader = DraftVarLoader(
engine=db.engine,

40
api/core/app/features/annotation_reply/annotation_reply.py
View File

@ -1,15 +1,14 @@
import logging
from typing import cast
from sqlalchemy import select
from sqlalchemy.orm import Session
from core.app.entities.app_invoke_entities import InvokeFrom
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from extensions.ext_database import db
from models.dataset import Dataset, DatasetCollectionBinding
from models.dataset import Dataset
from models.enums import CollectionBindingType, ConversationFromSource
from models.model import App, AppAnnotationSetting, Message, MessageAnnotation
from models.model import AnnotationReplyEnabledConfig, App, Message, MessageAnnotation, load_annotation_reply_config
from services.annotation_service import AppAnnotationService
from services.dataset_service import DatasetCollectionBindingService
@ -25,34 +24,27 @@ class AnnotationReplyFeature:
user_id: str,
invoke_from: InvokeFrom,
*,
session: Session | None = None,
session: Session,
) -> MessageAnnotation | None:
"""Return the closest annotation reply and record a hit in ``session``.
The caller may provide its transaction so the setting lookup, annotation
lookup, and hit-history write share one session. Runtime callers that do
not provide one continue to use Flask-SQLAlchemy's scoped session.
Vector-search failures are logged and return ``None``; transaction
cleanup remains the caller's responsibility.
The setting lookup, vector access, annotation lookup, and hit-history
write share the caller-owned session. Vector-search failures are logged
and return ``None``; transaction cleanup remains the caller's responsibility.
"""
if session is None:
session = db.session()
stmt = select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_record.id)
annotation_setting = session.scalar(stmt)
if not annotation_setting:
try:
annotation_reply_config = load_annotation_reply_config(session, app_record.id)
except ValueError:
return None
collection_binding_detail = session.get(DatasetCollectionBinding, annotation_setting.collection_binding_id)
if not collection_binding_detail:
if not annotation_reply_config["enabled"]:
return None
enabled_config = cast(AnnotationReplyEnabledConfig, annotation_reply_config)
try:
score_threshold = annotation_setting.score_threshold or 1
embedding_provider_name = collection_binding_detail.provider_name
embedding_model_name = collection_binding_detail.model_name
score_threshold = enabled_config["score_threshold"] or 1
embedding_provider_name = enabled_config["embedding_model"]["embedding_provider_name"]
embedding_model_name = enabled_config["embedding_model"]["embedding_model_name"]
dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
embedding_provider_name, embedding_model_name, session, CollectionBindingType.ANNOTATION
@ -67,7 +59,7 @@ class AnnotationReplyFeature:
collection_binding_id=dataset_collection_binding.id,
)
vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"])
vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"], session=session)
documents = vector.search_by_vector(
query=query, top_k=1, score_threshold=score_threshold, filter={"group_id": [dataset.id]}

Some files were not shown because too many files have changed in this diff Show More