Compare commits

...
Author SHA1 Message Date
yyh 6f582c1b2f fix(web): remove indicator from WorkplaceSelector and restore original item padding
The indicator was unnecessary (current workspace is shown in the trigger)
and shifted layout, breaking avatar alignment. Restore original
gap-2/pl-3/pr-2 spacing. Mark SelectItemIndicator as @public for Knip.

Made-with: Cursor
2026-04-16 11:13:07 +08:00
yyh f734c35443 refactor(web): redesign SelectTrigger as pure re-export and migrate WorkplaceSelector to Select
SelectTrigger was a heavily styled function coupling ARIA trigger
behavior with form-input appearance (input bg, size variants, built-in
chevron, clearable/loading). This made it unusable for non-form contexts
like WorkplaceSelector and inconsistent with DropdownMenuTrigger which
is a pure re-export.

- SelectTrigger is now a pure re-export of BaseSelect.Trigger (aligned
  with DropdownMenuTrigger pattern)
- Added SelectIcon re-export for composed triggers needing chevron state
- Removed SelectPrimitiveTrigger, selectTriggerVariants, contentPadding,
  and SelectTriggerProps
- 3 form callers migrated to inline composition with className
- WorkplaceSelector migrated from @headlessui/react Menu to base-ui
  Select with correct listbox semantics, fixing the focus ring issue

Made-with: Cursor
2026-04-16 11:06:03 +08:00
99andGitHub cf4d7afb9c chore(api): prune redundant direct dependency declarations (#35272) 2026-04-15 16:53:41 +00:00
e6b5923ff1 test(types): replace Account/Tenant status string literals with enum values (#35267)
Co-authored-by: xr843 <xianren843@protonmail.com>
2026-04-15 15:52:38 +00:00
yyhGitHubcopilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
538093855b refactor(web): align Switch API with Base UI naming convention (#35269)
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
2026-04-15 15:47:59 +00:00
yyhandGitHub e6b8cbe657 fix(web): include dify-ui workspace package in docker install filter (#35268) 2026-04-15 14:46:50 +00:00
yyhGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
af7d5e60b4 feat(ui): scaffold @langgenius/dify-ui and migrate design tokens (#35256)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-15 13:11:20 +00:00
dbceb3067e refactor(api): migrate console tag responses from marshal_with to BaseModel (#35208)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-15 09:57:27 +00:00
JamesandGitHub 425457cb16 test: remove legacy workflow draft variable api test (#35226) 2026-04-15 09:53:53 +00:00
JamesandGitHub e5bd18132c test: remove legacy jinja2 code executor integration test (#35222) 2026-04-15 09:53:46 +00:00
JamesandGitHub 2f33867d07 test: remove legacy trigger provider permissions test (#35227) 2026-04-15 09:53:42 +00:00
JamesandGitHub fd71c56f16 test: remove legacy storage key loader integration test (#35225) 2026-04-15 09:52:31 +00:00
zyssyz123GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>Joel
e3c2116501 fix: remove enable for get (#35245)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Joel <iamjoel007@gmail.com>
2026-04-15 09:18:29 +00:00
fb17339d89 feat(web): unify create_app tracking and persist external attribution (#35241)
Co-authored-by: CodingOnStar <hanxujiang@dify.com>
2026-04-15 08:59:31 +00:00
zyssyz123GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
9fd196642d feat: tidb endpoint (#35158)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-15 08:46:25 +00:00
JamesGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
98897a5379 test: migrate webhook service additional mock tests to testcontainers (#35199)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-15 08:14:36 +00:00
Wu TianweiGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>yyh
5542329554 fix(dataset): fix dataset list overlay issue (#35244)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: yyh <yuanyouhuilyz@gmail.com>
2026-04-15 08:03:02 +00:00
zyssyz123andGitHub 79332c0e5e fix: Change 'commit' to 'flush' to prevent subsequent transaction fai… (#35243) 2026-04-15 07:50:52 +00:00
yyhGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
50a55513d4 refactor(ui): decouple CSS dependencies and improve test quality (#35242)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-15 07:03:40 +00:00
3bccdd6c9a test: migrate Service API site controller tests to Testcontainers (#32454) (#35183)
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-04-15 05:51:34 +00:00
非法操作andGitHub 76af80e332 fix: open restore version panel raise 500 (#35240) 2026-04-15 04:43:38 +00:00
非法操作GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
7a880ae60c fix: import DSL and copy app not work (#35239)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-15 04:43:22 +00:00
JamesandGitHub 5bc0f9513b test: remove legacy python3 code executor integration test (#35223) 2026-04-15 02:20:06 +00:00
JamesandGitHub b77801ece9 test: remove legacy code executor integration test (#35224) 2026-04-15 02:19:42 +00:00
JamesGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
7de92c598f test: migrate schedule service mock tests to testcontainers (#35196)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-14 20:34:53 +00:00
JamesGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
693080aa12 test: migrate dataset service dataset mock tests to testcontainers (#35194)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-14 19:52:31 +00:00
25c388d0db refactor(api): migrate console workflow-trigger responses to BaseModel (#35200)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 19:52:17 +00:00
b1722c8af9 refactor(api): migrate console conversation variables response model to BaseModel (#35193)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 19:51:33 +00:00
b65a5fcd97 refactor(api): migrate service api workflow responses from marshal_with to BaseModel (#35195)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 19:50:59 +00:00
1c3cba281a refactor(api): migrate console message responses from marshal_with to BaseModel (#35204)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 19:49:33 +00:00
800954f8ce refactor(api): migrate service conversation-variable responses to BaseModel (#35205)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 19:49:21 +00:00
f66a3c49c4 refactor(api): migrate console recommended-app response to BaseModel (#35206)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
2026-04-14 19:48:43 +00:00
ef396ac84e refactor(api): migrate workspace current response from marshal_with to BaseModel (#35207)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 19:48:09 +00:00
7e7b27fdec refactor: replace bare dict with dict[str, Any] in response converter… (#35212)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-14 19:45:04 +00:00
wdeveloper16GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
9c90c1c455 refactor: replace bare dict with dict[str, Any] in services and hosti… (#35211)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-14 19:44:40 +00:00
b1df52b8ff refactor(api): migrate console workflow app-log responses to BaseModel (#35201)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 18:43:09 +00:00
e527b7c5f1 refactor(api): migrate console installed-app list response to BaseModel (#35202)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 18:40:40 +00:00
wdeveloper16andGitHub 149b9d4c0f refactor: replace bare dict with dict[str, Any] in services unit test helpers (#35182) 2026-04-14 18:37:59 +00:00
tmimmanuelGitHubtmimmanuel <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>
ef28a63ad3 refactor(api): add null safety to extractor_processor and firecrawl (#35209)
Co-authored-by: tmimmanuel <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>
2026-04-14 18:23:20 +00:00
e78558bc06 refactor(api): migrate dataset hit-testing response model to BaseModel (#35192)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 18:12:40 +00:00
JamesandGitHub f63d7c4121 test: remove document service status mock tests superseded by testcontainers (#35197) 2026-04-14 18:07:00 +00:00
ef062fb397 refactor(api): migrate console extension endpoint from api.model to BaseModel (#35189)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 17:58:33 +00:00
a2ea7ca039 refactor(api): migrate workspace account marshal_with responses to BaseModel (#35190)
Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>
2026-04-14 17:57:34 +00:00
bohdansolovieGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
6876cd787b test: migrate web site controller tests to Testcontainers (#32454) (#35180)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-14 17:52:13 +00:00
wdeveloper16andGitHub 50a6892c3a refactor: replace bare dict with dict[str, Any] in controller and core unit tests (#35181) 2026-04-14 17:51:49 +00:00
wdeveloper16andGitHub 1bcc7f78c7 refactor: replace bare dict with dict[str, Any] in models, providers, extensions, libs, and test utilities (#35186) 2026-04-14 17:51:25 +00:00
wdeveloper16GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2fd5b76ac1 refactor: replace bare dict with dict[str, Any] in enterprise telemetry, external data tool, and moderation tests (#35185)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-14 17:51:02 +00:00
wdeveloper16GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
62f42b3f24 refactor: replace bare dict with dict[str, Any] in RAG and service unit tests (#35184)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-14 17:50:43 +00:00
1201 changed files with 9342 additions and 10662 deletions
+14 -16
View File
@@ -6,10 +6,9 @@ from typing import Any, Literal
from flask import request
from flask_restx import Resource
from graphon.enums import WorkflowExecutionStatus
from graphon.file import helpers as file_helpers
from pydantic import AliasChoices, BaseModel, Field, computed_field, field_validator
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session
from werkzeug.exceptions import BadRequest
from controllers.common.helpers import FileInfo
@@ -31,13 +30,14 @@ from core.rag.retrieval.retrieval_methods import RetrievalMethod
from core.trigger.constants import TRIGGER_NODE_TYPES
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.helper import build_icon_url
from libs.login import current_account_with_tenant, login_required
from models import App, DatasetPermissionEnum, Workflow
from models.model import IconType
from services.app_dsl_service import AppDslService
from services.app_service import AppService
from services.enterprise.enterprise_service import EnterpriseService
from services.entities.dsl_entities import ImportMode
from services.entities.dsl_entities import ImportMode, ImportStatus
from services.entities.knowledge_entities.knowledge_entities import (
DataSource,
InfoList,
@@ -161,15 +161,6 @@ def _to_timestamp(value: datetime | int | None) -> int | None:
return value
def _build_icon_url(icon_type: str | IconType | None, icon: str | None) -> str | None:
if icon is None or icon_type is None:
return None
icon_type_value = icon_type.value if isinstance(icon_type, IconType) else str(icon_type)
if icon_type_value.lower() != IconType.IMAGE:
return None
return file_helpers.get_signed_file_url(icon)
class Tag(ResponseModel):
id: str
name: str
@@ -292,7 +283,7 @@ class Site(ResponseModel):
@computed_field(return_type=str | None) # type: ignore
@property
def icon_url(self) -> str | None:
return _build_icon_url(self.icon_type, self.icon)
return build_icon_url(self.icon_type, self.icon)
@field_validator("icon_type", mode="before")
@classmethod
@@ -342,7 +333,7 @@ class AppPartial(ResponseModel):
@computed_field(return_type=str | None) # type: ignore
@property
def icon_url(self) -> str | None:
return _build_icon_url(self.icon_type, self.icon)
return build_icon_url(self.icon_type, self.icon)
@field_validator("created_at", "updated_at", mode="before")
@classmethod
@@ -390,7 +381,7 @@ class AppDetailWithSite(AppDetail):
@computed_field(return_type=str | None) # type: ignore
@property
def icon_url(self) -> str | None:
return _build_icon_url(self.icon_type, self.icon)
return build_icon_url(self.icon_type, self.icon)
class AppPagination(ResponseModel):
@@ -632,7 +623,7 @@ class AppCopyApi(Resource):
args = CopyAppPayload.model_validate(console_ns.payload or {})
with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
with Session(db.engine, expire_on_commit=False) as session:
import_service = AppDslService(session)
yaml_content = import_service.export_dsl(app_model=app_model, include_secret=True)
result = import_service.import_app(
@@ -645,6 +636,13 @@ class AppCopyApi(Resource):
icon=args.icon,
icon_background=args.icon_background,
)
if result.status == ImportStatus.FAILED:
session.rollback()
return result.model_dump(mode="json"), 400
if result.status == ImportStatus.PENDING:
session.rollback()
return result.model_dump(mode="json"), 202
session.commit()
# Inherit web app permission from original app
if result.app_id and FeatureService.get_system_features().webapp_auth.enabled:
+14 -6
View File
@@ -1,6 +1,6 @@
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session
from controllers.common.schema import register_schema_models
from controllers.console.app.wraps import get_app_model
@@ -52,8 +52,9 @@ class AppImportApi(Resource):
current_user, _ = current_account_with_tenant()
args = AppImportPayload.model_validate(console_ns.payload)
# Create service with session
with sessionmaker(db.engine).begin() as session:
# AppDslService performs internal commits for some creation paths, so use a plain
# Session here instead of nesting it inside sessionmaker(...).begin().
with Session(db.engine, expire_on_commit=False) as session:
import_service = AppDslService(session)
# Import app
account = current_user
@@ -69,6 +70,10 @@ class AppImportApi(Resource):
icon_background=args.icon_background,
app_id=args.app_id,
)
if result.status == ImportStatus.FAILED:
session.rollback()
else:
session.commit()
if result.app_id and FeatureService.get_system_features().webapp_auth.enabled:
# update web app setting as private
EnterpriseService.WebAppAuth.update_app_access_mode(result.app_id, "private")
@@ -95,12 +100,15 @@ class AppImportConfirmApi(Resource):
# Check user role first
current_user, _ = current_account_with_tenant()
# Create service with session
with sessionmaker(db.engine).begin() as session:
with Session(db.engine, expire_on_commit=False) as session:
import_service = AppDslService(session)
# Confirm import
account = current_user
result = import_service.confirm_import(import_id=import_id, account=account)
if result.status == ImportStatus.FAILED:
session.rollback()
else:
session.commit()
# Return appropriate status code based on result
if result.status == ImportStatus.FAILED:
@@ -117,7 +125,7 @@ class AppImportCheckDependenciesApi(Resource):
@account_initialization_required
@edit_permission_required
def get(self, app_model: App):
with sessionmaker(db.engine).begin() as session:
with Session(db.engine, expire_on_commit=False) as session:
import_service = AppDslService(session)
result = import_service.check_dependencies(app_model=app_model)
@@ -1,44 +1,86 @@
from __future__ import annotations
from datetime import datetime
from typing import Any
from flask import request
from flask_restx import Resource, fields, marshal_with
from pydantic import BaseModel, Field
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from controllers.common.schema import register_schema_models
from controllers.console import console_ns
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import account_initialization_required, setup_required
from extensions.ext_database import db
from fields.conversation_variable_fields import (
conversation_variable_fields,
paginated_conversation_variable_fields,
)
from fields._value_type_serializer import serialize_value_type
from fields.base import ResponseModel
from libs.login import login_required
from models import ConversationVariable
from models.model import AppMode
DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}"
class ConversationVariablesQuery(BaseModel):
conversation_id: str = Field(..., description="Conversation ID to filter variables")
console_ns.schema_model(
ConversationVariablesQuery.__name__,
ConversationVariablesQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0),
)
def _to_timestamp(value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
# Register models for flask_restx to avoid dict type issues in Swagger
# Register base model first
conversation_variable_model = console_ns.model("ConversationVariable", conversation_variable_fields)
# For nested models, need to replace nested dict with registered model
paginated_conversation_variable_fields_copy = paginated_conversation_variable_fields.copy()
paginated_conversation_variable_fields_copy["data"] = fields.List(
fields.Nested(conversation_variable_model), attribute="data"
)
paginated_conversation_variable_model = console_ns.model(
"PaginatedConversationVariable", paginated_conversation_variable_fields_copy
class ConversationVariableResponse(ResponseModel):
id: str
name: str
value_type: str
value: str | None = None
description: str | None = None
created_at: int | None = None
updated_at: int | None = None
@field_validator("value_type", mode="before")
@classmethod
def _normalize_value_type(cls, value: Any) -> str:
exposed_type = getattr(value, "exposed_type", None)
if callable(exposed_type):
return str(exposed_type().value)
if isinstance(value, str):
return value
try:
return serialize_value_type(value)
except Exception:
return serialize_value_type({"value_type": value})
@field_validator("value", mode="before")
@classmethod
def _normalize_value(cls, value: Any | None) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return str(value)
@field_validator("created_at", "updated_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
class PaginatedConversationVariableResponse(ResponseModel):
page: int
limit: int
total: int
has_more: bool
data: list[ConversationVariableResponse]
register_schema_models(
console_ns,
ConversationVariablesQuery,
ConversationVariableResponse,
PaginatedConversationVariableResponse,
)
@@ -48,12 +90,15 @@ class ConversationVariablesApi(Resource):
@console_ns.doc(description="Get conversation variables for an application")
@console_ns.doc(params={"app_id": "Application ID"})
@console_ns.expect(console_ns.models[ConversationVariablesQuery.__name__])
@console_ns.response(200, "Conversation variables retrieved successfully", paginated_conversation_variable_model)
@console_ns.response(
200,
"Conversation variables retrieved successfully",
console_ns.models[PaginatedConversationVariableResponse.__name__],
)
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=AppMode.ADVANCED_CHAT)
@marshal_with(paginated_conversation_variable_model)
def get(self, app_model):
args = ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)) # type: ignore
@@ -72,17 +117,22 @@ class ConversationVariablesApi(Resource):
with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
rows = session.scalars(stmt).all()
return {
"page": page,
"limit": page_size,
"total": len(rows),
"has_more": False,
"data": [
{
"created_at": row.created_at,
"updated_at": row.updated_at,
**row.to_variable().model_dump(),
}
for row in rows
],
}
response = PaginatedConversationVariableResponse.model_validate(
{
"page": page,
"limit": page_size,
"total": len(rows),
"has_more": False,
"data": [
ConversationVariableResponse.model_validate(
{
"created_at": row.created_at,
"updated_at": row.updated_at,
**row.to_variable().model_dump(),
}
)
for row in rows
],
}
)
return response.model_dump(mode="json")
+69 -127
View File
@@ -1,8 +1,9 @@
import logging
from datetime import datetime
from typing import Literal
from flask import request
from flask_restx import Resource, fields, marshal_with
from flask_restx import Resource
from graphon.model_runtime.errors.invoke import InvokeError
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import exists, func, select
@@ -25,10 +26,21 @@ from controllers.console.wraps import (
setup_required,
)
from core.app.entities.app_invoke_entities import InvokeFrom
from core.entities.execution_extra_content import ExecutionExtraContentDomainModel
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from extensions.ext_database import db
from fields.raws import FilesContainedField
from libs.helper import TimestampField, uuid_value
from fields.base import ResponseModel
from fields.conversation_fields import (
AgentThought,
ConversationAnnotation,
ConversationAnnotationHitHistory,
Feedback,
JSONValue,
MessageFile,
format_files_contained,
to_timestamp,
)
from libs.helper import uuid_value
from libs.infinite_scroll_pagination import InfiniteScrollPagination
from libs.login import current_account_with_tenant, login_required
from models.enums import FeedbackFromSource, FeedbackRating
@@ -98,6 +110,51 @@ class SuggestedQuestionsResponse(BaseModel):
data: list[str] = Field(description="Suggested question")
class MessageDetailResponse(ResponseModel):
id: str
conversation_id: str
inputs: dict[str, JSONValue]
query: str
message: JSONValue | None = None
message_tokens: int | None = None
answer: str = Field(validation_alias="re_sign_file_url_answer")
answer_tokens: int | None = None
provider_response_latency: float | None = None
from_source: str
from_end_user_id: str | None = None
from_account_id: str | None = None
feedbacks: list[Feedback] = Field(default_factory=list)
workflow_run_id: str | None = None
annotation: ConversationAnnotation | None = None
annotation_hit_history: ConversationAnnotationHitHistory | None = None
created_at: int | None = None
agent_thoughts: list[AgentThought] = Field(default_factory=list)
message_files: list[MessageFile] = Field(default_factory=list)
extra_contents: list[ExecutionExtraContentDomainModel] = Field(default_factory=list)
metadata: JSONValue | None = Field(default=None, validation_alias="message_metadata_dict")
status: str
error: str | None = None
parent_message_id: str | None = None
@field_validator("inputs", mode="before")
@classmethod
def _normalize_inputs(cls, value: JSONValue) -> JSONValue:
return format_files_contained(value)
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return to_timestamp(value)
return value
class MessageInfiniteScrollPaginationResponse(ResponseModel):
limit: int
has_more: bool
data: list[MessageDetailResponse]
register_schema_models(
console_ns,
ChatMessagesQuery,
@@ -105,124 +162,8 @@ register_schema_models(
FeedbackExportQuery,
AnnotationCountResponse,
SuggestedQuestionsResponse,
)
# Register models for flask_restx to avoid dict type issues in Swagger
# Register in dependency order: base models first, then dependent models
# Base models
simple_account_model = console_ns.model(
"SimpleAccount",
{
"id": fields.String,
"name": fields.String,
"email": fields.String,
},
)
message_file_model = console_ns.model(
"MessageFile",
{
"id": fields.String,
"filename": fields.String,
"type": fields.String,
"url": fields.String,
"mime_type": fields.String,
"size": fields.Integer,
"transfer_method": fields.String,
"belongs_to": fields.String(default="user"),
"upload_file_id": fields.String(default=None),
},
)
agent_thought_model = console_ns.model(
"AgentThought",
{
"id": fields.String,
"chain_id": fields.String,
"message_id": fields.String,
"position": fields.Integer,
"thought": fields.String,
"tool": fields.String,
"tool_labels": fields.Raw,
"tool_input": fields.String,
"created_at": TimestampField,
"observation": fields.String,
"files": fields.List(fields.String),
},
)
# Models that depend on simple_account_model
feedback_model = console_ns.model(
"Feedback",
{
"rating": fields.String,
"content": fields.String,
"from_source": fields.String,
"from_end_user_id": fields.String,
"from_account": fields.Nested(simple_account_model, allow_null=True),
},
)
annotation_model = console_ns.model(
"Annotation",
{
"id": fields.String,
"question": fields.String,
"content": fields.String,
"account": fields.Nested(simple_account_model, allow_null=True),
"created_at": TimestampField,
},
)
annotation_hit_history_model = console_ns.model(
"AnnotationHitHistory",
{
"annotation_id": fields.String(attribute="id"),
"annotation_create_account": fields.Nested(simple_account_model, allow_null=True),
"created_at": TimestampField,
},
)
# Message detail model that depends on multiple models
message_detail_model = console_ns.model(
"MessageDetail",
{
"id": fields.String,
"conversation_id": fields.String,
"inputs": FilesContainedField,
"query": fields.String,
"message": fields.Raw,
"message_tokens": fields.Integer,
"answer": fields.String(attribute="re_sign_file_url_answer"),
"answer_tokens": fields.Integer,
"provider_response_latency": fields.Float,
"from_source": fields.String,
"from_end_user_id": fields.String,
"from_account_id": fields.String,
"feedbacks": fields.List(fields.Nested(feedback_model)),
"workflow_run_id": fields.String,
"annotation": fields.Nested(annotation_model, allow_null=True),
"annotation_hit_history": fields.Nested(annotation_hit_history_model, allow_null=True),
"created_at": TimestampField,
"agent_thoughts": fields.List(fields.Nested(agent_thought_model)),
"message_files": fields.List(fields.Nested(message_file_model)),
"extra_contents": fields.List(fields.Raw),
"metadata": fields.Raw(attribute="message_metadata_dict"),
"status": fields.String,
"error": fields.String,
"parent_message_id": fields.String,
},
)
# Message infinite scroll pagination model
message_infinite_scroll_pagination_model = console_ns.model(
"MessageInfiniteScrollPagination",
{
"limit": fields.Integer,
"has_more": fields.Boolean,
"data": fields.List(fields.Nested(message_detail_model)),
},
MessageDetailResponse,
MessageInfiniteScrollPaginationResponse,
)
@@ -232,13 +173,12 @@ class ChatMessageListApi(Resource):
@console_ns.doc(description="Get chat messages for a conversation with pagination")
@console_ns.doc(params={"app_id": "Application ID"})
@console_ns.expect(console_ns.models[ChatMessagesQuery.__name__])
@console_ns.response(200, "Success", message_infinite_scroll_pagination_model)
@console_ns.response(200, "Success", console_ns.models[MessageInfiniteScrollPaginationResponse.__name__])
@console_ns.response(404, "Conversation not found")
@login_required
@account_initialization_required
@setup_required
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])
@marshal_with(message_infinite_scroll_pagination_model)
@edit_permission_required
def get(self, app_model):
args = ChatMessagesQuery.model_validate(request.args.to_dict())
@@ -298,7 +238,10 @@ class ChatMessageListApi(Resource):
history_messages = list(reversed(history_messages))
attach_message_extra_contents(history_messages)
return InfiniteScrollPagination(data=history_messages, limit=args.limit, has_more=has_more)
return MessageInfiniteScrollPaginationResponse.model_validate(
InfiniteScrollPagination(data=history_messages, limit=args.limit, has_more=has_more),
from_attributes=True,
).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/feedbacks")
@@ -468,13 +411,12 @@ class MessageApi(Resource):
@console_ns.doc("get_message")
@console_ns.doc(description="Get message details by ID")
@console_ns.doc(params={"app_id": "Application ID", "message_id": "Message ID"})
@console_ns.response(200, "Message retrieved successfully", message_detail_model)
@console_ns.response(200, "Message retrieved successfully", console_ns.models[MessageDetailResponse.__name__])
@console_ns.response(404, "Message not found")
@get_app_model
@setup_required
@login_required
@account_initialization_required
@marshal_with(message_detail_model)
def get(self, app_model, message_id: str):
message_id = str(message_id)
@@ -486,4 +428,4 @@ class MessageApi(Resource):
raise NotFound("Message Not Exists.")
attach_message_extra_contents([message])
return message
return MessageDetailResponse.model_validate(message, from_attributes=True).model_dump(mode="json")
+3 -3
View File
@@ -4,7 +4,7 @@ from collections.abc import Sequence
from typing import Any
from flask import abort, request
from flask_restx import Resource, fields, marshal_with
from flask_restx import Resource, fields, marshal, marshal_with
from graphon.enums import NodeType
from graphon.file import File
from graphon.graph_engine.manager import GraphEngineManager
@@ -942,7 +942,6 @@ class PublishedAllWorkflowApi(Resource):
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
@marshal_with(workflow_pagination_model)
@edit_permission_required
def get(self, app_model: App):
"""
@@ -970,9 +969,10 @@ class PublishedAllWorkflowApi(Resource):
user_id=user_id,
named_only=named_only,
)
serialized_workflows = marshal(workflows, workflow_fields_copy)
return {
"items": workflows,
"items": serialized_workflows,
"page": page,
"limit": limit,
"has_more": has_more,
+128 -19
View File
@@ -1,27 +1,26 @@
from datetime import datetime
from typing import Any
from dateutil.parser import isoparse
from flask import request
from flask_restx import Resource, marshal_with
from flask_restx import Resource
from graphon.enums import WorkflowExecutionStatus
from pydantic import BaseModel, Field, field_validator
from sqlalchemy.orm import sessionmaker
from controllers.common.schema import register_schema_models
from controllers.console import console_ns
from controllers.console.app.wraps import get_app_model
from controllers.console.wraps import account_initialization_required, setup_required
from extensions.ext_database import db
from fields.workflow_app_log_fields import (
build_workflow_app_log_pagination_model,
build_workflow_archived_log_pagination_model,
)
from fields.base import ResponseModel
from fields.end_user_fields import SimpleEndUser
from fields.member_fields import SimpleAccount
from libs.login import login_required
from models import App
from models.model import AppMode
from services.workflow_app_service import WorkflowAppService
DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}"
class WorkflowAppLogQuery(BaseModel):
keyword: str | None = Field(default=None, description="Search keyword for filtering logs")
@@ -58,13 +57,113 @@ class WorkflowAppLogQuery(BaseModel):
raise ValueError("Invalid boolean value for detail")
console_ns.schema_model(
WorkflowAppLogQuery.__name__, WorkflowAppLogQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0)
)
class WorkflowRunForLogResponse(ResponseModel):
id: str
version: str | None = None
status: str | None = None
triggered_from: str | None = None
error: str | None = None
elapsed_time: float | None = None
total_tokens: int | None = None
total_steps: int | None = None
created_at: int | None = None
finished_at: int | None = None
exceptions_count: int | None = None
# Register model for flask_restx to avoid dict type issues in Swagger
workflow_app_log_pagination_model = build_workflow_app_log_pagination_model(console_ns)
workflow_archived_log_pagination_model = build_workflow_archived_log_pagination_model(console_ns)
@field_validator("status", mode="before")
@classmethod
def _normalize_status(cls, value: Any) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return str(getattr(value, "value", value))
@field_validator("created_at", "finished_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
class WorkflowRunForArchivedLogResponse(ResponseModel):
id: str
status: str | None = None
triggered_from: str | None = None
elapsed_time: float | None = None
total_tokens: int | None = None
@field_validator("status", mode="before")
@classmethod
def _normalize_status(cls, value: Any) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return str(getattr(value, "value", value))
class WorkflowAppLogPartialResponse(ResponseModel):
id: str
workflow_run: WorkflowRunForLogResponse | None = None
details: Any = None
created_from: str | None = None
created_by_role: str | None = None
created_by_account: SimpleAccount | None = None
created_by_end_user: SimpleEndUser | None = None
created_at: int | None = None
@field_validator("created_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
class WorkflowArchivedLogPartialResponse(ResponseModel):
id: str
workflow_run: WorkflowRunForArchivedLogResponse | None = None
trigger_metadata: Any = None
created_by_account: SimpleAccount | None = None
created_by_end_user: SimpleEndUser | None = None
created_at: int | None = None
@field_validator("created_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
class WorkflowAppLogPaginationResponse(ResponseModel):
page: int
limit: int
total: int
has_more: bool
data: list[WorkflowAppLogPartialResponse]
class WorkflowArchivedLogPaginationResponse(ResponseModel):
page: int
limit: int
total: int
has_more: bool
data: list[WorkflowArchivedLogPartialResponse]
register_schema_models(
console_ns,
WorkflowAppLogQuery,
WorkflowRunForLogResponse,
WorkflowRunForArchivedLogResponse,
WorkflowAppLogPartialResponse,
WorkflowArchivedLogPartialResponse,
WorkflowAppLogPaginationResponse,
WorkflowArchivedLogPaginationResponse,
)
@console_ns.route("/apps/<uuid:app_id>/workflow-app-logs")
@@ -73,12 +172,15 @@ class WorkflowAppLogApi(Resource):
@console_ns.doc(description="Get workflow application execution logs")
@console_ns.doc(params={"app_id": "Application ID"})
@console_ns.expect(console_ns.models[WorkflowAppLogQuery.__name__])
@console_ns.response(200, "Workflow app logs retrieved successfully", workflow_app_log_pagination_model)
@console_ns.response(
200,
"Workflow app logs retrieved successfully",
console_ns.models[WorkflowAppLogPaginationResponse.__name__],
)
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW])
@marshal_with(workflow_app_log_pagination_model)
def get(self, app_model: App):
"""
Get workflow app logs
@@ -102,7 +204,9 @@ class WorkflowAppLogApi(Resource):
created_by_account=args.created_by_account,
)
return workflow_app_log_pagination
return WorkflowAppLogPaginationResponse.model_validate(
workflow_app_log_pagination, from_attributes=True
).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/workflow-archived-logs")
@@ -111,12 +215,15 @@ class WorkflowArchivedLogApi(Resource):
@console_ns.doc(description="Get workflow archived execution logs")
@console_ns.doc(params={"app_id": "Application ID"})
@console_ns.expect(console_ns.models[WorkflowAppLogQuery.__name__])
@console_ns.response(200, "Workflow archived logs retrieved successfully", workflow_archived_log_pagination_model)
@console_ns.response(
200,
"Workflow archived logs retrieved successfully",
console_ns.models[WorkflowArchivedLogPaginationResponse.__name__],
)
@setup_required
@login_required
@account_initialization_required
@get_app_model(mode=[AppMode.WORKFLOW])
@marshal_with(workflow_archived_log_pagination_model)
def get(self, app_model: App):
"""
Get workflow archived logs
@@ -132,4 +239,6 @@ class WorkflowArchivedLogApi(Resource):
limit=args.limit,
)
return workflow_app_log_pagination
return WorkflowArchivedLogPaginationResponse.model_validate(
workflow_app_log_pagination, from_attributes=True
).model_dump(mode="json")
+58 -22
View File
@@ -1,16 +1,17 @@
import logging
from datetime import datetime
from flask import request
from flask_restx import Resource, fields, marshal_with
from pydantic import BaseModel
from flask_restx import Resource
from pydantic import BaseModel, field_validator
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from werkzeug.exceptions import NotFound
from configs import dify_config
from controllers.common.schema import get_or_create_model
from controllers.common.schema import register_schema_models
from extensions.ext_database import db
from fields.workflow_trigger_fields import trigger_fields, triggers_list_fields, webhook_trigger_fields
from fields.base import ResponseModel
from libs.login import current_user, login_required
from models.enums import AppTriggerStatus
from models.model import Account, App, AppMode
@@ -21,15 +22,6 @@ from ..app.wraps import get_app_model
from ..wraps import account_initialization_required, edit_permission_required, setup_required
logger = logging.getLogger(__name__)
DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}"
trigger_model = get_or_create_model("WorkflowTrigger", trigger_fields)
triggers_list_fields_copy = triggers_list_fields.copy()
triggers_list_fields_copy["data"] = fields.List(fields.Nested(trigger_model))
triggers_list_model = get_or_create_model("WorkflowTriggerList", triggers_list_fields_copy)
webhook_trigger_model = get_or_create_model("WebhookTrigger", webhook_trigger_fields)
class Parser(BaseModel):
@@ -41,10 +33,52 @@ class ParserEnable(BaseModel):
enable_trigger: bool
console_ns.schema_model(Parser.__name__, Parser.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0))
class WorkflowTriggerResponse(ResponseModel):
id: str
trigger_type: str
title: str
node_id: str
provider_name: str
icon: str
status: str
created_at: datetime | None = None
updated_at: datetime | None = None
console_ns.schema_model(
ParserEnable.__name__, ParserEnable.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0)
@field_validator("id", "trigger_type", "title", "node_id", "provider_name", "icon", "status", mode="before")
@classmethod
def _normalize_string_fields(cls, value: object) -> str:
if isinstance(value, str):
return value
return str(value)
class WorkflowTriggerListResponse(ResponseModel):
data: list[WorkflowTriggerResponse]
class WebhookTriggerResponse(ResponseModel):
id: str
webhook_id: str
webhook_url: str
webhook_debug_url: str
node_id: str
created_at: datetime | None = None
@field_validator("id", "webhook_id", "webhook_url", "webhook_debug_url", "node_id", mode="before")
@classmethod
def _normalize_string_fields(cls, value: object) -> str:
if isinstance(value, str):
return value
return str(value)
register_schema_models(
console_ns,
Parser,
ParserEnable,
WorkflowTriggerResponse,
WorkflowTriggerListResponse,
WebhookTriggerResponse,
)
@@ -57,7 +91,7 @@ class WebhookTriggerApi(Resource):
@login_required
@account_initialization_required
@get_app_model(mode=AppMode.WORKFLOW)
@marshal_with(webhook_trigger_model)
@console_ns.response(200, "Success", console_ns.models[WebhookTriggerResponse.__name__])
def get(self, app_model: App):
"""Get webhook trigger for a node"""
args = Parser.model_validate(request.args.to_dict(flat=True)) # type: ignore
@@ -78,7 +112,7 @@ class WebhookTriggerApi(Resource):
if not webhook_trigger:
raise NotFound("Webhook trigger not found for this node")
return webhook_trigger
return WebhookTriggerResponse.model_validate(webhook_trigger, from_attributes=True).model_dump(mode="json")
@console_ns.route("/apps/<uuid:app_id>/triggers")
@@ -89,7 +123,7 @@ class AppTriggersApi(Resource):
@login_required
@account_initialization_required
@get_app_model(mode=AppMode.WORKFLOW)
@marshal_with(triggers_list_model)
@console_ns.response(200, "Success", console_ns.models[WorkflowTriggerListResponse.__name__])
def get(self, app_model: App):
"""Get app triggers list"""
assert isinstance(current_user, Account)
@@ -118,7 +152,9 @@ class AppTriggersApi(Resource):
else:
trigger.icon = "" # type: ignore
return {"data": triggers}
return WorkflowTriggerListResponse.model_validate({"data": triggers}, from_attributes=True).model_dump(
mode="json"
)
@console_ns.route("/apps/<uuid:app_id>/trigger-enable")
@@ -129,7 +165,7 @@ class AppTriggerEnableApi(Resource):
@account_initialization_required
@edit_permission_required
@get_app_model(mode=AppMode.WORKFLOW)
@marshal_with(trigger_model)
@console_ns.response(200, "Success", console_ns.models[WorkflowTriggerResponse.__name__])
def post(self, app_model: App):
"""Update app trigger (enable/disable)"""
args = ParserEnable.model_validate(console_ns.payload)
@@ -160,4 +196,4 @@ class AppTriggerEnableApi(Resource):
else:
trigger.icon = "" # type: ignore
return trigger
return WorkflowTriggerResponse.model_validate(trigger, from_attributes=True).model_dump(mode="json")
+3 -1
View File
@@ -1,3 +1,5 @@
from typing import Any
from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
@@ -40,7 +42,7 @@ class ActivatePayload(BaseModel):
class ActivationCheckResponse(BaseModel):
is_valid: bool = Field(description="Whether token is valid")
data: dict | None = Field(default=None, description="Activation data if valid")
data: dict[str, Any] | None = Field(default=None, description="Activation data if valid")
class ActivationResponse(BaseModel):
@@ -1026,7 +1026,7 @@ class DocumentMetadataApi(DocumentResource):
if not isinstance(doc_metadata, dict):
raise ValueError("doc_metadata must be a dictionary.")
metadata_schema: dict = cast(dict, DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type])
metadata_schema: dict[str, Any] = cast(dict[str, Any], DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type])
document.doc_metadata = {}
if doc_type == "others":
+93 -36
View File
@@ -1,13 +1,13 @@
from flask_restx import Resource, fields
from __future__ import annotations
from controllers.common.schema import register_schema_model
from fields.hit_testing_fields import (
child_chunk_fields,
document_fields,
files_fields,
hit_testing_record_fields,
segment_fields,
)
from datetime import datetime
from typing import Any
from flask_restx import Resource
from pydantic import Field, field_validator
from controllers.common.schema import register_schema_models
from fields.base import ResponseModel
from libs.login import login_required
from .. import console_ns
@@ -18,39 +18,92 @@ from ..wraps import (
setup_required,
)
register_schema_model(console_ns, HitTestingPayload)
def _to_timestamp(value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
def _get_or_create_model(model_name: str, field_def):
"""Get or create a flask_restx model to avoid dict type issues in Swagger."""
existing = console_ns.models.get(model_name)
if existing is None:
existing = console_ns.model(model_name, field_def)
return existing
class HitTestingDocument(ResponseModel):
id: str | None = None
data_source_type: str | None = None
name: str | None = None
doc_type: str | None = None
doc_metadata: Any | None = None
# Register models for flask_restx to avoid dict type issues in Swagger
document_model = _get_or_create_model("HitTestingDocument", document_fields)
class HitTestingSegment(ResponseModel):
id: str | None = None
position: int | None = None
document_id: str | None = None
content: str | None = None
sign_content: str | None = None
answer: str | None = None
word_count: int | None = None
tokens: int | None = None
keywords: list[str] = Field(default_factory=list)
index_node_id: str | None = None
index_node_hash: str | None = None
hit_count: int | None = None
enabled: bool | None = None
disabled_at: int | None = None
disabled_by: str | None = None
status: str | None = None
created_by: str | None = None
created_at: int | None = None
indexing_at: int | None = None
completed_at: int | None = None
error: str | None = None
stopped_at: int | None = None
document: HitTestingDocument | None = None
segment_fields_copy = segment_fields.copy()
segment_fields_copy["document"] = fields.Nested(document_model)
segment_model = _get_or_create_model("HitTestingSegment", segment_fields_copy)
@field_validator("disabled_at", "created_at", "indexing_at", "completed_at", "stopped_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
child_chunk_model = _get_or_create_model("HitTestingChildChunk", child_chunk_fields)
files_model = _get_or_create_model("HitTestingFile", files_fields)
hit_testing_record_fields_copy = hit_testing_record_fields.copy()
hit_testing_record_fields_copy["segment"] = fields.Nested(segment_model)
hit_testing_record_fields_copy["child_chunks"] = fields.List(fields.Nested(child_chunk_model))
hit_testing_record_fields_copy["files"] = fields.List(fields.Nested(files_model))
hit_testing_record_model = _get_or_create_model("HitTestingRecord", hit_testing_record_fields_copy)
class HitTestingChildChunk(ResponseModel):
id: str | None = None
content: str | None = None
position: int | None = None
score: float | None = None
# Response model for hit testing API
hit_testing_response_fields = {
"query": fields.String,
"records": fields.List(fields.Nested(hit_testing_record_model)),
}
hit_testing_response_model = _get_or_create_model("HitTestingResponse", hit_testing_response_fields)
class HitTestingFile(ResponseModel):
id: str | None = None
name: str | None = None
size: int | None = None
extension: str | None = None
mime_type: str | None = None
source_url: str | None = None
class HitTestingRecord(ResponseModel):
segment: HitTestingSegment | None = None
child_chunks: list[HitTestingChildChunk] = Field(default_factory=list)
score: float | None = None
tsne_position: Any | None = None
files: list[HitTestingFile] = Field(default_factory=list)
summary: str | None = None
class HitTestingResponse(ResponseModel):
query: str
records: list[HitTestingRecord] = Field(default_factory=list)
register_schema_models(
console_ns,
HitTestingPayload,
HitTestingDocument,
HitTestingSegment,
HitTestingChildChunk,
HitTestingFile,
HitTestingRecord,
HitTestingResponse,
)
@console_ns.route("/datasets/<uuid:dataset_id>/hit-testing")
@@ -59,7 +112,11 @@ class HitTestingApi(Resource, DatasetsHitTestingBase):
@console_ns.doc(description="Test dataset knowledge retrieval")
@console_ns.doc(params={"dataset_id": "Dataset ID"})
@console_ns.expect(console_ns.models[HitTestingPayload.__name__])
@console_ns.response(200, "Hit testing completed successfully", model=hit_testing_response_model)
@console_ns.response(
200,
"Hit testing completed successfully",
model=console_ns.models[HitTestingResponse.__name__],
)
@console_ns.response(404, "Dataset not found")
@console_ns.response(400, "Invalid parameters")
@setup_required
@@ -74,4 +131,4 @@ class HitTestingApi(Resource, DatasetsHitTestingBase):
args = payload.model_dump(exclude_none=True)
self.hit_testing_args_check(args)
return self.perform_hit_testing(dataset, args)
return HitTestingResponse.model_validate(self.perform_hit_testing(dataset, args)).model_dump(mode="json")
@@ -1,21 +1,24 @@
import logging
from datetime import datetime
from typing import Any
from flask import request
from flask_restx import Resource, fields, marshal_with
from pydantic import BaseModel, Field
from flask_restx import Resource
from graphon.file import helpers as file_helpers
from pydantic import BaseModel, Field, computed_field, field_validator
from sqlalchemy import and_, select
from werkzeug.exceptions import BadRequest, Forbidden, NotFound
from controllers.common.schema import get_or_create_model
from controllers.common.schema import register_schema_models
from controllers.console import console_ns
from controllers.console.explore.wraps import InstalledAppResource
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
from extensions.ext_database import db
from fields.installed_app_fields import app_fields, installed_app_fields, installed_app_list_fields
from fields.base import ResponseModel
from libs.datetime_utils import naive_utc_now
from libs.login import current_account_with_tenant, login_required
from models import App, InstalledApp, RecommendedApp
from models.model import IconType
from services.account_service import TenantService
from services.enterprise.enterprise_service import EnterpriseService
from services.feature_service import FeatureService
@@ -36,22 +39,97 @@ class InstalledAppsListQuery(BaseModel):
logger = logging.getLogger(__name__)
app_model = get_or_create_model("InstalledAppInfo", app_fields)
def _build_icon_url(icon_type: str | IconType | None, icon: str | None) -> str | None:
if icon is None or icon_type is None:
return None
icon_type_value = icon_type.value if isinstance(icon_type, IconType) else str(icon_type)
if icon_type_value.lower() != IconType.IMAGE:
return None
return file_helpers.get_signed_file_url(icon)
installed_app_fields_copy = installed_app_fields.copy()
installed_app_fields_copy["app"] = fields.Nested(app_model)
installed_app_model = get_or_create_model("InstalledApp", installed_app_fields_copy)
installed_app_list_fields_copy = installed_app_list_fields.copy()
installed_app_list_fields_copy["installed_apps"] = fields.List(fields.Nested(installed_app_model))
installed_app_list_model = get_or_create_model("InstalledAppList", installed_app_list_fields_copy)
def _safe_primitive(value: Any) -> Any:
if value is None or isinstance(value, (str, int, float, bool, datetime)):
return value
return None
class InstalledAppInfoResponse(ResponseModel):
id: str
name: str | None = None
mode: str | None = None
icon_type: str | None = None
icon: str | None = None
icon_background: str | None = None
use_icon_as_answer_icon: bool | None = None
@field_validator("mode", "icon_type", mode="before")
@classmethod
def _normalize_enum_like(cls, value: Any) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return str(getattr(value, "value", value))
@computed_field(return_type=str | None) # type: ignore[prop-decorator]
@property
def icon_url(self) -> str | None:
return _build_icon_url(self.icon_type, self.icon)
class InstalledAppResponse(ResponseModel):
id: str
app: InstalledAppInfoResponse
app_owner_tenant_id: str
is_pinned: bool
last_used_at: int | None = None
editable: bool
uninstallable: bool
@field_validator("app", mode="before")
@classmethod
def _normalize_app(cls, value: Any) -> Any:
if isinstance(value, dict):
return value
return {
"id": _safe_primitive(getattr(value, "id", "")) or "",
"name": _safe_primitive(getattr(value, "name", None)),
"mode": _safe_primitive(getattr(value, "mode", None)),
"icon_type": _safe_primitive(getattr(value, "icon_type", None)),
"icon": _safe_primitive(getattr(value, "icon", None)),
"icon_background": _safe_primitive(getattr(value, "icon_background", None)),
"use_icon_as_answer_icon": _safe_primitive(getattr(value, "use_icon_as_answer_icon", None)),
}
@field_validator("last_used_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
class InstalledAppListResponse(ResponseModel):
installed_apps: list[InstalledAppResponse]
register_schema_models(
console_ns,
InstalledAppCreatePayload,
InstalledAppUpdatePayload,
InstalledAppsListQuery,
InstalledAppInfoResponse,
InstalledAppResponse,
InstalledAppListResponse,
)
@console_ns.route("/installed-apps")
class InstalledAppsListApi(Resource):
@login_required
@account_initialization_required
@marshal_with(installed_app_list_model)
@console_ns.response(200, "Success", console_ns.models[InstalledAppListResponse.__name__])
def get(self):
query = InstalledAppsListQuery.model_validate(request.args.to_dict())
current_user, current_tenant_id = current_account_with_tenant()
@@ -125,7 +203,9 @@ class InstalledAppsListApi(Resource):
)
)
return {"installed_apps": installed_app_list}
return InstalledAppListResponse.model_validate(
{"installed_apps": installed_app_list}, from_attributes=True
).model_dump(mode="json")
@login_required
@account_initialization_required
@@ -1,66 +1,83 @@
from typing import Any
from flask import request
from flask_restx import Resource, fields, marshal_with
from pydantic import BaseModel, Field
from flask_restx import Resource
from pydantic import BaseModel, Field, computed_field, field_validator
from constants.languages import languages
from controllers.common.schema import get_or_create_model
from controllers.common.schema import register_schema_models
from controllers.console import console_ns
from controllers.console.wraps import account_initialization_required
from libs.helper import AppIconUrlField
from fields.base import ResponseModel
from libs.helper import build_icon_url
from libs.login import current_user, login_required
from services.recommended_app_service import RecommendedAppService
app_fields = {
"id": fields.String,
"name": fields.String,
"mode": fields.String,
"icon": fields.String,
"icon_type": fields.String,
"icon_url": AppIconUrlField,
"icon_background": fields.String,
}
app_model = get_or_create_model("RecommendedAppInfo", app_fields)
recommended_app_fields = {
"app": fields.Nested(app_model, attribute="app"),
"app_id": fields.String,
"description": fields.String(attribute="description"),
"copyright": fields.String,
"privacy_policy": fields.String,
"custom_disclaimer": fields.String,
"category": fields.String,
"position": fields.Integer,
"is_listed": fields.Boolean,
"can_trial": fields.Boolean,
}
recommended_app_model = get_or_create_model("RecommendedApp", recommended_app_fields)
recommended_app_list_fields = {
"recommended_apps": fields.List(fields.Nested(recommended_app_model)),
"categories": fields.List(fields.String),
}
recommended_app_list_model = get_or_create_model("RecommendedAppList", recommended_app_list_fields)
class RecommendedAppsQuery(BaseModel):
language: str | None = Field(default=None)
console_ns.schema_model(
RecommendedAppsQuery.__name__,
RecommendedAppsQuery.model_json_schema(ref_template="#/definitions/{model}"),
class RecommendedAppInfoResponse(ResponseModel):
id: str
name: str | None = None
mode: str | None = None
icon: str | None = None
icon_type: str | None = None
icon_background: str | None = None
@staticmethod
def _normalize_enum_like(value: Any) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return str(getattr(value, "value", value))
@field_validator("mode", "icon_type", mode="before")
@classmethod
def _normalize_enum_fields(cls, value: Any) -> str | None:
return cls._normalize_enum_like(value)
@computed_field(return_type=str | None) # type: ignore[prop-decorator]
@property
def icon_url(self) -> str | None:
return build_icon_url(self.icon_type, self.icon)
class RecommendedAppResponse(ResponseModel):
app: RecommendedAppInfoResponse | None = None
app_id: str
description: str | None = None
copyright: str | None = None
privacy_policy: str | None = None
custom_disclaimer: str | None = None
category: str | None = None
position: int | None = None
is_listed: bool | None = None
can_trial: bool | None = None
class RecommendedAppListResponse(ResponseModel):
recommended_apps: list[RecommendedAppResponse]
categories: list[str]
register_schema_models(
console_ns,
RecommendedAppsQuery,
RecommendedAppInfoResponse,
RecommendedAppResponse,
RecommendedAppListResponse,
)
@console_ns.route("/explore/apps")
class RecommendedAppListApi(Resource):
@console_ns.expect(console_ns.models[RecommendedAppsQuery.__name__])
@console_ns.response(200, "Success", console_ns.models[RecommendedAppListResponse.__name__])
@login_required
@account_initialization_required
@marshal_with(recommended_app_list_model)
def get(self):
# language args
args = RecommendedAppsQuery.model_validate(request.args.to_dict(flat=True)) # type: ignore
@@ -72,7 +89,10 @@ class RecommendedAppListApi(Resource):
else:
language_prefix = languages[0]
return RecommendedAppService.get_recommended_apps_and_categories(language_prefix)
return RecommendedAppListResponse.model_validate(
RecommendedAppService.get_recommended_apps_and_categories(language_prefix),
from_attributes=True,
).model_dump(mode="json")
@console_ns.route("/explore/apps/<uuid:app_id>")
+2 -6
View File
@@ -169,6 +169,7 @@ console_ns.schema_model(
class TrialAppWorkflowRunApi(TrialAppResource):
@trial_feature_enable
@console_ns.expect(console_ns.models[WorkflowRunRequest.__name__])
def post(self, trial_app):
"""
@@ -210,6 +211,7 @@ class TrialAppWorkflowRunApi(TrialAppResource):
class TrialAppWorkflowTaskStopApi(TrialAppResource):
@trial_feature_enable
def post(self, trial_app, task_id: str):
"""
Stop workflow task
@@ -290,7 +292,6 @@ class TrialChatApi(TrialAppResource):
class TrialMessageSuggestedQuestionApi(TrialAppResource):
@trial_feature_enable
def get(self, trial_app, message_id):
app_model = trial_app
app_mode = AppMode.value_of(app_model.mode)
@@ -470,7 +471,6 @@ class TrialCompletionApi(TrialAppResource):
class TrialSitApi(Resource):
"""Resource for trial app sites."""
@trial_feature_enable
@get_app_model_with_trial(None)
def get(self, app_model):
"""Retrieve app site info.
@@ -492,7 +492,6 @@ class TrialSitApi(Resource):
class TrialAppParameterApi(Resource):
"""Resource for app variables."""
@trial_feature_enable
@get_app_model_with_trial(None)
def get(self, app_model):
"""Retrieve app parameters."""
@@ -521,7 +520,6 @@ class TrialAppParameterApi(Resource):
class AppApi(Resource):
@trial_feature_enable
@get_app_model_with_trial(None)
@marshal_with(app_detail_with_site_model)
def get(self, app_model):
@@ -534,7 +532,6 @@ class AppApi(Resource):
class AppWorkflowApi(Resource):
@trial_feature_enable
@get_app_model_with_trial(None)
@marshal_with(workflow_model)
def get(self, app_model):
@@ -547,7 +544,6 @@ class AppWorkflowApi(Resource):
class DatasetListApi(Resource):
@trial_feature_enable
@get_app_model_with_trial(None)
def get(self, app_model):
page = request.args.get("page", default=1, type=int)
+68 -24
View File
@@ -1,15 +1,18 @@
from datetime import datetime
from typing import Any
from flask import request
from flask_restx import Resource, fields, marshal_with
from pydantic import BaseModel, Field
from flask_restx import Resource
from pydantic import BaseModel, Field, TypeAdapter, field_validator
from constants import HIDDEN_VALUE
from fields.api_based_extension_fields import api_based_extension_fields
from fields.base import ResponseModel
from libs.login import current_account_with_tenant, login_required
from models.api_based_extension import APIBasedExtension
from services.api_based_extension_service import APIBasedExtensionService
from services.code_based_extension_service import CodeBasedExtensionService
from ..common.schema import register_schema_models
from ..common.schema import DEFAULT_REF_TEMPLATE_SWAGGER_2_0, register_schema_models
from . import console_ns
from .wraps import account_initialization_required, setup_required
@@ -24,12 +27,52 @@ class APIBasedExtensionPayload(BaseModel):
api_key: str = Field(description="API key for authentication")
register_schema_models(console_ns, APIBasedExtensionPayload)
class CodeBasedExtensionResponse(ResponseModel):
module: str = Field(description="Module name")
data: Any = Field(description="Extension data")
api_based_extension_model = console_ns.model("ApiBasedExtensionModel", api_based_extension_fields)
def _mask_api_key(api_key: str) -> str:
if not api_key:
return api_key
if len(api_key) <= 8:
return api_key[0] + "******" + api_key[-1]
return api_key[:3] + "******" + api_key[-3:]
api_based_extension_list_model = fields.List(fields.Nested(api_based_extension_model))
def _to_timestamp(value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
class APIBasedExtensionResponse(ResponseModel):
id: str
name: str
api_endpoint: str
api_key: str
created_at: int | None = None
@field_validator("api_key", mode="before")
@classmethod
def _normalize_api_key(cls, value: str) -> str:
return _mask_api_key(value)
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
register_schema_models(console_ns, APIBasedExtensionPayload, CodeBasedExtensionResponse, APIBasedExtensionResponse)
console_ns.schema_model(
"APIBasedExtensionListResponse",
TypeAdapter(list[APIBasedExtensionResponse]).json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0),
)
def _serialize_api_based_extension(extension: APIBasedExtension) -> dict[str, Any]:
return APIBasedExtensionResponse.model_validate(extension, from_attributes=True).model_dump(mode="json")
@console_ns.route("/code-based-extension")
@@ -40,10 +83,7 @@ class CodeBasedExtensionAPI(Resource):
@console_ns.response(
200,
"Success",
console_ns.model(
"CodeBasedExtensionResponse",
{"module": fields.String(description="Module name"), "data": fields.Raw(description="Extension data")},
),
console_ns.models[CodeBasedExtensionResponse.__name__],
)
@setup_required
@login_required
@@ -51,30 +91,34 @@ class CodeBasedExtensionAPI(Resource):
def get(self):
query = CodeBasedExtensionQuery.model_validate(request.args.to_dict(flat=True)) # type: ignore
return {"module": query.module, "data": CodeBasedExtensionService.get_code_based_extension(query.module)}
return CodeBasedExtensionResponse(
module=query.module,
data=CodeBasedExtensionService.get_code_based_extension(query.module),
).model_dump(mode="json")
@console_ns.route("/api-based-extension")
class APIBasedExtensionAPI(Resource):
@console_ns.doc("get_api_based_extensions")
@console_ns.doc(description="Get all API-based extensions for current tenant")
@console_ns.response(200, "Success", api_based_extension_list_model)
@console_ns.response(200, "Success", console_ns.models["APIBasedExtensionListResponse"])
@setup_required
@login_required
@account_initialization_required
@marshal_with(api_based_extension_model)
def get(self):
_, tenant_id = current_account_with_tenant()
return APIBasedExtensionService.get_all_by_tenant_id(tenant_id)
return [
_serialize_api_based_extension(extension)
for extension in APIBasedExtensionService.get_all_by_tenant_id(tenant_id)
]
@console_ns.doc("create_api_based_extension")
@console_ns.doc(description="Create a new API-based extension")
@console_ns.expect(console_ns.models[APIBasedExtensionPayload.__name__])
@console_ns.response(201, "Extension created successfully", api_based_extension_model)
@console_ns.response(201, "Extension created successfully", console_ns.models[APIBasedExtensionResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@marshal_with(api_based_extension_model)
def post(self):
payload = APIBasedExtensionPayload.model_validate(console_ns.payload or {})
_, current_tenant_id = current_account_with_tenant()
@@ -86,7 +130,7 @@ class APIBasedExtensionAPI(Resource):
api_key=payload.api_key,
)
return APIBasedExtensionService.save(extension_data)
return _serialize_api_based_extension(APIBasedExtensionService.save(extension_data))
@console_ns.route("/api-based-extension/<uuid:id>")
@@ -94,26 +138,26 @@ class APIBasedExtensionDetailAPI(Resource):
@console_ns.doc("get_api_based_extension")
@console_ns.doc(description="Get API-based extension by ID")
@console_ns.doc(params={"id": "Extension ID"})
@console_ns.response(200, "Success", api_based_extension_model)
@console_ns.response(200, "Success", console_ns.models[APIBasedExtensionResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@marshal_with(api_based_extension_model)
def get(self, id):
api_based_extension_id = str(id)
_, tenant_id = current_account_with_tenant()
return APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
return _serialize_api_based_extension(
APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
)
@console_ns.doc("update_api_based_extension")
@console_ns.doc(description="Update API-based extension")
@console_ns.doc(params={"id": "Extension ID"})
@console_ns.expect(console_ns.models[APIBasedExtensionPayload.__name__])
@console_ns.response(200, "Extension updated successfully", api_based_extension_model)
@console_ns.response(200, "Extension updated successfully", console_ns.models[APIBasedExtensionResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@marshal_with(api_based_extension_model)
def post(self, id):
api_based_extension_id = str(id)
_, current_tenant_id = current_account_with_tenant()
@@ -128,7 +172,7 @@ class APIBasedExtensionDetailAPI(Resource):
if payload.api_key != HIDDEN_VALUE:
extension_data_from_db.api_key = payload.api_key
return APIBasedExtensionService.save(extension_data_from_db)
return _serialize_api_based_extension(APIBasedExtensionService.save(extension_data_from_db))
@console_ns.doc("delete_api_based_extension")
@console_ns.doc(description="Delete API-based extension")
+39 -17
View File
@@ -1,13 +1,14 @@
from typing import Literal
from flask import request
from flask_restx import Namespace, Resource, fields, marshal_with
from pydantic import BaseModel, Field
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator
from werkzeug.exceptions import Forbidden
from controllers.common.schema import register_schema_models
from controllers.console import console_ns
from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required
from fields.base import ResponseModel
from libs.login import current_account_with_tenant, login_required
from models.enums import TagType
from services.tag_service import (
@@ -18,17 +19,6 @@ from services.tag_service import (
UpdateTagPayload,
)
dataset_tag_fields = {
"id": fields.String,
"name": fields.String,
"type": fields.String,
"binding_count": fields.String,
}
def build_dataset_tag_fields(api_or_ns: Namespace):
return api_or_ns.model("DataSetTag", dataset_tag_fields)
class TagBasePayload(BaseModel):
name: str = Field(description="Tag name", min_length=1, max_length=50)
@@ -52,12 +42,36 @@ class TagListQueryParam(BaseModel):
keyword: str | None = Field(None, description="Search keyword")
class TagResponse(ResponseModel):
id: str
name: str
type: str | None = None
binding_count: str | None = None
@field_validator("type", mode="before")
@classmethod
def normalize_type(cls, value: TagType | str | None) -> str | None:
if value is None:
return None
if isinstance(value, TagType):
return value.value
return value
@field_validator("binding_count", mode="before")
@classmethod
def normalize_binding_count(cls, value: int | str | None) -> str | None:
if value is None:
return None
return str(value)
register_schema_models(
console_ns,
TagBasePayload,
TagBindingPayload,
TagBindingRemovePayload,
TagListQueryParam,
TagResponse,
)
@@ -69,14 +83,18 @@ class TagListApi(Resource):
@console_ns.doc(
params={"type": 'Tag type filter. Can be "knowledge" or "app".', "keyword": "Search keyword for tag name."}
)
@marshal_with(dataset_tag_fields)
@console_ns.doc(responses={200: ("Success", [console_ns.models[TagResponse.__name__]])})
def get(self):
_, current_tenant_id = current_account_with_tenant()
raw_args = request.args.to_dict()
param = TagListQueryParam.model_validate(raw_args)
tags = TagService.get_tags(param.type, current_tenant_id, param.keyword)
return tags, 200
serialized_tags = [
TagResponse.model_validate(tag, from_attributes=True).model_dump(mode="json") for tag in tags
]
return serialized_tags, 200
@console_ns.expect(console_ns.models[TagBasePayload.__name__])
@setup_required
@@ -91,7 +109,9 @@ class TagListApi(Resource):
payload = TagBasePayload.model_validate(console_ns.payload or {})
tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=payload.type))
response = {"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": 0}
response = TagResponse.model_validate(
{"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": 0}
).model_dump(mode="json")
return response, 200
@@ -114,7 +134,9 @@ class TagUpdateDeleteApi(Resource):
binding_count = TagService.get_tag_binding_count(tag_id)
response = {"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": binding_count}
response = TagResponse.model_validate(
{"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": binding_count}
).model_dump(mode="json")
return response, 200
+68 -38
View File
@@ -5,7 +5,7 @@ from typing import Any, Literal
import pytz
from flask import request
from flask_restx import Resource, fields, marshal_with
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator, model_validator
from sqlalchemy import select
@@ -37,9 +37,10 @@ from controllers.console.wraps import (
setup_required,
)
from extensions.ext_database import db
from fields.base import ResponseModel
from fields.member_fields import Account as AccountResponse
from libs.datetime_utils import naive_utc_now
from libs.helper import EmailStr, TimestampField, extract_remote_ip, timezone
from libs.helper import EmailStr, extract_remote_ip, timezone
from libs.login import current_account_with_tenant, login_required
from models import AccountIntegrate, InvitationCode
from models.account import AccountStatus, InvitationCodeStatus
@@ -178,17 +179,57 @@ def _serialize_account(account) -> dict[str, Any]:
return AccountResponse.model_validate(account, from_attributes=True).model_dump(mode="json")
integrate_fields = {
"provider": fields.String,
"created_at": TimestampField,
"is_bound": fields.Boolean,
"link": fields.String,
}
def _to_timestamp(value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
integrate_model = console_ns.model("AccountIntegrate", integrate_fields)
integrate_list_model = console_ns.model(
"AccountIntegrateList",
{"data": fields.List(fields.Nested(integrate_model))},
class AccountIntegrateResponse(ResponseModel):
provider: str
created_at: int | None = None
is_bound: bool
link: str | None = None
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
class AccountIntegrateListResponse(ResponseModel):
data: list[AccountIntegrateResponse]
class EducationVerifyResponse(ResponseModel):
token: str | None = None
class EducationStatusResponse(ResponseModel):
result: bool | None = None
is_student: bool | None = None
expire_at: int | None = None
allow_refresh: bool | None = None
@field_validator("expire_at", mode="before")
@classmethod
def _normalize_expire_at(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
class EducationAutocompleteResponse(ResponseModel):
data: list[str] = Field(default_factory=list)
curr_page: int | None = None
has_next: bool | None = None
register_schema_models(
console_ns,
AccountIntegrateResponse,
AccountIntegrateListResponse,
EducationVerifyResponse,
EducationStatusResponse,
EducationAutocompleteResponse,
)
@@ -359,7 +400,7 @@ class AccountIntegrateApi(Resource):
@setup_required
@login_required
@account_initialization_required
@marshal_with(integrate_list_model)
@console_ns.response(200, "Success", console_ns.models[AccountIntegrateListResponse.__name__])
def get(self):
account, _ = current_account_with_tenant()
@@ -395,7 +436,9 @@ class AccountIntegrateApi(Resource):
}
)
return {"data": integrate_data}
return AccountIntegrateListResponse(
data=[AccountIntegrateResponse.model_validate(item) for item in integrate_data]
).model_dump(mode="json")
@console_ns.route("/account/delete/verify")
@@ -447,31 +490,22 @@ class AccountDeleteUpdateFeedbackApi(Resource):
@console_ns.route("/account/education/verify")
class EducationVerifyApi(Resource):
verify_fields = {
"token": fields.String,
}
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@cloud_edition_billing_enabled
@marshal_with(verify_fields)
@console_ns.response(200, "Success", console_ns.models[EducationVerifyResponse.__name__])
def get(self):
account, _ = current_account_with_tenant()
return BillingService.EducationIdentity.verify(account.id, account.email)
return EducationVerifyResponse.model_validate(
BillingService.EducationIdentity.verify(account.id, account.email) or {}
).model_dump(mode="json")
@console_ns.route("/account/education")
class EducationApi(Resource):
status_fields = {
"result": fields.Boolean,
"is_student": fields.Boolean,
"expire_at": TimestampField,
"allow_refresh": fields.Boolean,
}
@console_ns.expect(console_ns.models[EducationActivatePayload.__name__])
@setup_required
@login_required
@@ -491,37 +525,33 @@ class EducationApi(Resource):
@account_initialization_required
@only_edition_cloud
@cloud_edition_billing_enabled
@marshal_with(status_fields)
@console_ns.response(200, "Success", console_ns.models[EducationStatusResponse.__name__])
def get(self):
account, _ = current_account_with_tenant()
res = BillingService.EducationIdentity.status(account.id)
res = BillingService.EducationIdentity.status(account.id) or {}
# convert expire_at to UTC timestamp from isoformat
if res and "expire_at" in res:
res["expire_at"] = datetime.fromisoformat(res["expire_at"]).astimezone(pytz.utc)
return res
return EducationStatusResponse.model_validate(res).model_dump(mode="json")
@console_ns.route("/account/education/autocomplete")
class EducationAutoCompleteApi(Resource):
data_fields = {
"data": fields.List(fields.String),
"curr_page": fields.Integer,
"has_next": fields.Boolean,
}
@console_ns.expect(console_ns.models[EducationAutocompleteQuery.__name__])
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@cloud_edition_billing_enabled
@marshal_with(data_fields)
@console_ns.response(200, "Success", console_ns.models[EducationAutocompleteResponse.__name__])
def get(self):
payload = request.args.to_dict(flat=True)
args = EducationAutocompleteQuery.model_validate(payload)
return BillingService.EducationIdentity.autocomplete(args.keywords, args.page, args.limit)
return EducationAutocompleteResponse.model_validate(
BillingService.EducationIdentity.autocomplete(args.keywords, args.page, args.limit) or {}
).model_dump(mode="json")
@console_ns.route("/account/change-email")
+1 -1
View File
@@ -465,7 +465,7 @@ class ModelProviderModelDisableApi(Resource):
class ParserValidate(BaseModel):
model: str
model_type: ModelType
credentials: dict
credentials: dict[str, Any]
console_ns.schema_model(
+44 -4
View File
@@ -1,8 +1,9 @@
import logging
from datetime import datetime
from flask import request
from flask_restx import Resource, fields, marshal, marshal_with
from pydantic import BaseModel, Field
from flask_restx import Resource, fields, marshal
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import select
from werkzeug.exceptions import Unauthorized
@@ -26,6 +27,7 @@ from controllers.console.wraps import (
)
from enums.cloud_plan import CloudPlan
from extensions.ext_database import db
from fields.base import ResponseModel
from libs.helper import TimestampField
from libs.login import current_account_with_tenant, login_required
from models.account import Tenant, TenantCustomConfigDict, TenantStatus
@@ -58,6 +60,37 @@ class WorkspaceInfoPayload(BaseModel):
name: str
class TenantInfoResponse(ResponseModel):
id: str
name: str | None = None
plan: str | None = None
status: str | None = None
created_at: int | None = None
role: str | None = None
in_trial: bool | None = None
trial_end_reason: str | None = None
custom_config: dict | None = None
trial_credits: int | None = None
trial_credits_used: int | None = None
next_credit_reset_date: int | None = None
@field_validator("plan", "status", "trial_end_reason", mode="before")
@classmethod
def _normalize_enum_like(cls, value):
if value is None:
return None
if isinstance(value, str):
return value
return str(getattr(value, "value", value))
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | int | None):
if isinstance(value, datetime):
return int(value.timestamp())
return value
def reg(cls: type[BaseModel]):
console_ns.schema_model(cls.__name__, cls.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0))
@@ -66,6 +99,7 @@ reg(WorkspaceListQuery)
reg(SwitchWorkspacePayload)
reg(WorkspaceCustomConfigPayload)
reg(WorkspaceInfoPayload)
reg(TenantInfoResponse)
provider_fields = {
"provider_name": fields.String,
@@ -180,7 +214,7 @@ class TenantApi(Resource):
@setup_required
@login_required
@account_initialization_required
@marshal_with(tenant_fields)
@console_ns.response(200, "Success", console_ns.models[TenantInfoResponse.__name__])
def post(self):
if request.path == "/info":
logger.warning("Deprecated URL /info was used.")
@@ -200,7 +234,13 @@ class TenantApi(Resource):
else:
raise Unauthorized("workspace is archived")
return WorkspaceService.get_tenant_info(tenant), 200
return (
TenantInfoResponse.model_validate(
WorkspaceService.get_tenant_info(tenant),
from_attributes=True,
).model_dump(mode="json"),
200,
)
@console_ns.route("/workspaces/switch")
+6 -2
View File
@@ -9,7 +9,7 @@ from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session
from controllers.common.schema import register_schema_model
from controllers.console.wraps import setup_required
@@ -56,7 +56,7 @@ class EnterpriseAppDSLImport(Resource):
account.set_tenant_id(workspace_id)
with sessionmaker(db.engine).begin() as session:
with Session(db.engine, expire_on_commit=False) as session:
dsl_service = AppDslService(session)
result = dsl_service.import_app(
account=account,
@@ -65,6 +65,10 @@ class EnterpriseAppDSLImport(Resource):
name=args.name,
description=args.description,
)
if result.status == ImportStatus.FAILED:
session.rollback()
else:
session.commit()
if result.status == ImportStatus.FAILED:
return result.model_dump(mode="json"), 400
@@ -1,7 +1,9 @@
from datetime import datetime
from typing import Any, Literal
from flask import request
from flask_restx import Resource
from graphon.variables.types import SegmentType
from pydantic import BaseModel, Field, TypeAdapter, field_validator
from sqlalchemy.orm import sessionmaker
from werkzeug.exceptions import BadRequest, NotFound
@@ -14,14 +16,12 @@ from controllers.service_api.app.error import NotChatAppError
from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from fields._value_type_serializer import serialize_value_type
from fields.base import ResponseModel
from fields.conversation_fields import (
ConversationInfiniteScrollPagination,
SimpleConversation,
)
from fields.conversation_variable_fields import (
build_conversation_variable_infinite_scroll_pagination_model,
build_conversation_variable_model,
)
from libs.helper import UUIDStrOrEmpty
from models.model import App, AppMode, EndUser
from services.conversation_service import ConversationService
@@ -70,12 +70,70 @@ class ConversationVariableUpdatePayload(BaseModel):
value: Any
class ConversationVariableResponse(ResponseModel):
id: str
name: str
value_type: str
value: str | None = None
description: str | None = None
created_at: int | None = None
updated_at: int | None = None
@field_validator("value_type", mode="before")
@classmethod
def normalize_value_type(cls, value: Any) -> str:
exposed_type = getattr(value, "exposed_type", None)
if callable(exposed_type):
return str(exposed_type().value)
if isinstance(value, str):
try:
return str(SegmentType(value).exposed_type().value)
except ValueError:
return value
try:
return serialize_value_type(value)
except (AttributeError, TypeError, ValueError):
pass
try:
return serialize_value_type({"value_type": value})
except (AttributeError, TypeError, ValueError):
value_attr = getattr(value, "value", None)
if value_attr is not None:
return str(value_attr)
return str(value)
@field_validator("value", mode="before")
@classmethod
def normalize_value(cls, value: Any | None) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return str(value)
@field_validator("created_at", "updated_at", mode="before")
@classmethod
def normalize_timestamp(cls, value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
class ConversationVariableInfiniteScrollPaginationResponse(ResponseModel):
limit: int
has_more: bool
data: list[ConversationVariableResponse]
register_schema_models(
service_api_ns,
ConversationListQuery,
ConversationRenamePayload,
ConversationVariablesQuery,
ConversationVariableUpdatePayload,
ConversationVariableResponse,
ConversationVariableInfiniteScrollPaginationResponse,
)
@@ -204,8 +262,12 @@ class ConversationVariablesApi(Resource):
404: "Conversation not found",
}
)
@service_api_ns.response(
200,
"Variables retrieved successfully",
service_api_ns.models[ConversationVariableInfiniteScrollPaginationResponse.__name__],
)
@validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.QUERY))
@service_api_ns.marshal_with(build_conversation_variable_infinite_scroll_pagination_model(service_api_ns))
def get(self, app_model: App, end_user: EndUser, c_id):
"""List all variables for a conversation.
@@ -222,9 +284,12 @@ class ConversationVariablesApi(Resource):
last_id = str(query_args.last_id) if query_args.last_id else None
try:
return ConversationService.get_conversational_variable(
pagination = ConversationService.get_conversational_variable(
app_model, conversation_id, end_user, query_args.limit, last_id, query_args.variable_name
)
return ConversationVariableInfiniteScrollPaginationResponse.model_validate(
pagination, from_attributes=True
).model_dump(mode="json")
except services.errors.conversation.ConversationNotExistsError:
raise NotFound("Conversation Not Exists.")
@@ -243,8 +308,12 @@ class ConversationVariableDetailApi(Resource):
404: "Conversation or variable not found",
}
)
@service_api_ns.response(
200,
"Variable updated successfully",
service_api_ns.models[ConversationVariableResponse.__name__],
)
@validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.JSON))
@service_api_ns.marshal_with(build_conversation_variable_model(service_api_ns))
def put(self, app_model: App, end_user: EndUser, c_id, variable_id):
"""Update a conversation variable's value.
@@ -261,9 +330,10 @@ class ConversationVariableDetailApi(Resource):
payload = ConversationVariableUpdatePayload.model_validate(service_api_ns.payload or {})
try:
return ConversationService.update_conversation_variable(
variable = ConversationService.update_conversation_variable(
app_model, conversation_id, variable_id, end_user, payload.value
)
return ConversationVariableResponse.model_validate(variable, from_attributes=True).model_dump(mode="json")
except services.errors.conversation.ConversationNotExistsError:
raise NotFound("Conversation Not Exists.")
except services.errors.conversation.ConversationVariableNotExistsError:
+141 -26
View File
@@ -1,13 +1,15 @@
import logging
from collections.abc import Mapping
from datetime import datetime
from typing import Literal
from dateutil.parser import isoparse
from flask import request
from flask_restx import Namespace, Resource, fields
from flask_restx import Resource, fields
from graphon.enums import WorkflowExecutionStatus
from graphon.graph_engine.manager import GraphEngineManager
from graphon.model_runtime.errors.invoke import InvokeError
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, field_validator
from sqlalchemy.orm import sessionmaker
from werkzeug.exceptions import BadRequest, InternalServerError, NotFound
@@ -33,9 +35,10 @@ from core.errors.error import (
from core.helper.trace_id_helper import get_external_trace_id
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from fields.workflow_app_log_fields import build_workflow_app_log_pagination_model
from fields.base import ResponseModel
from fields.end_user_fields import SimpleEndUser
from fields.member_fields import SimpleAccount
from libs import helper
from libs.helper import OptionalTimestampField, TimestampField
from models.model import App, AppMode, EndUser
from models.workflow import WorkflowRun
from repositories.factory import DifyAPIRepositoryFactory
@@ -65,38 +68,142 @@ class WorkflowLogQuery(BaseModel):
register_schema_models(service_api_ns, WorkflowRunPayload, WorkflowLogQuery)
def _to_timestamp(value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
def _enum_value(value):
return getattr(value, "value", value)
class WorkflowRunStatusField(fields.Raw):
def output(self, key, obj: WorkflowRun, **kwargs):
return obj.status.value
return _enum_value(obj.status)
class WorkflowRunOutputsField(fields.Raw):
def output(self, key, obj: WorkflowRun, **kwargs):
if obj.status == WorkflowExecutionStatus.PAUSED:
status = _enum_value(obj.status)
if status == WorkflowExecutionStatus.PAUSED.value:
return {}
outputs = obj.outputs_dict
return outputs or {}
workflow_run_fields = {
"id": fields.String,
"workflow_id": fields.String,
"status": WorkflowRunStatusField,
"inputs": fields.Raw,
"outputs": WorkflowRunOutputsField,
"error": fields.String,
"total_steps": fields.Integer,
"total_tokens": fields.Integer,
"created_at": TimestampField,
"finished_at": OptionalTimestampField,
"elapsed_time": fields.Float,
}
class WorkflowRunResponse(ResponseModel):
id: str
workflow_id: str
status: str
inputs: dict | list | str | int | float | bool | None = None
outputs: dict = Field(default_factory=dict)
error: str | None = None
total_steps: int | None = None
total_tokens: int | None = None
created_at: int | None = None
finished_at: int | None = None
elapsed_time: float | int | None = None
@field_validator("created_at", "finished_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
def build_workflow_run_model(api_or_ns: Namespace):
"""Build the workflow run model for the API or Namespace."""
return api_or_ns.model("WorkflowRun", workflow_run_fields)
class WorkflowRunForLogResponse(ResponseModel):
id: str
version: str | None = None
status: str | None = None
triggered_from: str | None = None
error: str | None = None
elapsed_time: float | int | None = None
total_tokens: int | None = None
total_steps: int | None = None
created_at: int | None = None
finished_at: int | None = None
exceptions_count: int | None = None
@field_validator("status", "triggered_from", mode="before")
@classmethod
def _normalize_enum(cls, value):
return _enum_value(value)
@field_validator("created_at", "finished_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
class WorkflowAppLogPartialResponse(ResponseModel):
id: str
workflow_run: WorkflowRunForLogResponse | None = None
details: dict | list | str | int | float | bool | None = None
created_from: str | None = None
created_by_role: str | None = None
created_by_account: SimpleAccount | None = None
created_by_end_user: SimpleEndUser | None = None
created_at: int | None = None
@field_validator("created_from", "created_by_role", mode="before")
@classmethod
def _normalize_enum(cls, value):
return _enum_value(value)
@field_validator("created_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
class WorkflowAppLogPaginationResponse(ResponseModel):
page: int
limit: int
total: int
has_more: bool
data: list[WorkflowAppLogPartialResponse]
register_schema_models(
service_api_ns,
WorkflowRunResponse,
WorkflowRunForLogResponse,
WorkflowAppLogPartialResponse,
WorkflowAppLogPaginationResponse,
)
def _serialize_workflow_run(workflow_run: WorkflowRun) -> dict:
status = _enum_value(workflow_run.status)
raw_outputs = workflow_run.outputs_dict
if status == WorkflowExecutionStatus.PAUSED.value or raw_outputs is None:
outputs: dict = {}
elif isinstance(raw_outputs, dict):
outputs = raw_outputs
elif isinstance(raw_outputs, Mapping):
outputs = dict(raw_outputs)
else:
outputs = {}
return WorkflowRunResponse.model_validate(
{
"id": workflow_run.id,
"workflow_id": workflow_run.workflow_id,
"status": status,
"inputs": workflow_run.inputs,
"outputs": outputs,
"error": workflow_run.error,
"total_steps": workflow_run.total_steps,
"total_tokens": workflow_run.total_tokens,
"created_at": workflow_run.created_at,
"finished_at": workflow_run.finished_at,
"elapsed_time": workflow_run.elapsed_time,
}
).model_dump(mode="json")
def _serialize_workflow_log_pagination(pagination) -> dict:
return WorkflowAppLogPaginationResponse.model_validate(pagination, from_attributes=True).model_dump(mode="json")
@service_api_ns.route("/workflows/run/<string:workflow_run_id>")
@@ -112,7 +219,11 @@ class WorkflowRunDetailApi(Resource):
}
)
@validate_app_token
@service_api_ns.marshal_with(build_workflow_run_model(service_api_ns))
@service_api_ns.response(
200,
"Workflow run details retrieved successfully",
service_api_ns.models[WorkflowRunResponse.__name__],
)
def get(self, app_model: App, workflow_run_id: str):
"""Get a workflow task running detail.
@@ -133,7 +244,7 @@ class WorkflowRunDetailApi(Resource):
)
if not workflow_run:
raise NotFound("Workflow run not found.")
return workflow_run
return _serialize_workflow_run(workflow_run)
@service_api_ns.route("/workflows/run")
@@ -299,7 +410,11 @@ class WorkflowAppLogApi(Resource):
}
)
@validate_app_token
@service_api_ns.marshal_with(build_workflow_app_log_pagination_model(service_api_ns))
@service_api_ns.response(
200,
"Logs retrieved successfully",
service_api_ns.models[WorkflowAppLogPaginationResponse.__name__],
)
def get(self, app_model: App):
"""Get workflow app logs.
@@ -327,4 +442,4 @@ class WorkflowAppLogApi(Resource):
created_by_account=args.created_by_account,
)
return workflow_app_log_pagination
return _serialize_workflow_log_pagination(workflow_app_log_pagination)
@@ -57,7 +57,7 @@ class AdvancedChatAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, Any, None]:
) -> Generator[dict[str, Any] | str, Any, None]:
"""
Convert stream full response.
:param stream_response: stream response
@@ -88,7 +88,7 @@ class AdvancedChatAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, Any, None]:
) -> Generator[dict[str, Any] | str, Any, None]:
"""
Convert stream simple response.
:param stream_response: stream response
@@ -1,5 +1,5 @@
from collections.abc import Generator
from typing import cast
from typing import Any, cast
from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
from core.app.entities.task_entities import (
@@ -56,7 +56,7 @@ class AgentChatAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream full response.
:param stream_response: stream response
@@ -87,7 +87,7 @@ class AgentChatAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream simple response.
:param stream_response: stream response
@@ -24,7 +24,7 @@ class AppGenerateResponseConverter(ABC):
return cls.convert_blocking_full_response(response)
else:
def _generate_full_response() -> Generator[dict | str, Any, None]:
def _generate_full_response() -> Generator[dict[str, Any] | str, Any, None]:
yield from cls.convert_stream_full_response(response)
return _generate_full_response()
@@ -33,7 +33,7 @@ class AppGenerateResponseConverter(ABC):
return cls.convert_blocking_simple_response(response)
else:
def _generate_simple_response() -> Generator[dict | str, Any, None]:
def _generate_simple_response() -> Generator[dict[str, Any] | str, Any, None]:
yield from cls.convert_stream_simple_response(response)
return _generate_simple_response()
@@ -52,14 +52,14 @@ class AppGenerateResponseConverter(ABC):
@abstractmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
raise NotImplementedError
@classmethod
@abstractmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
raise NotImplementedError
@classmethod
@@ -1,5 +1,5 @@
from collections.abc import Generator
from typing import cast
from typing import Any, cast
from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
from core.app.entities.task_entities import (
@@ -56,7 +56,7 @@ class ChatAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream full response.
:param stream_response: stream response
@@ -87,7 +87,7 @@ class ChatAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream simple response.
:param stream_response: stream response
@@ -1,5 +1,5 @@
from collections.abc import Generator
from typing import cast
from typing import Any, cast
from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
from core.app.entities.task_entities import (
@@ -55,7 +55,7 @@ class CompletionAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream full response.
:param stream_response: stream response
@@ -85,7 +85,7 @@ class CompletionAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream simple response.
:param stream_response: stream response
@@ -37,7 +37,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream full response.
:param stream_response: stream response
@@ -66,7 +66,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream simple response.
:param stream_response: stream response
@@ -1,5 +1,5 @@
from collections.abc import Generator
from typing import cast
from typing import Any, cast
from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
from core.app.entities.task_entities import (
@@ -37,7 +37,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_full_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream full response.
:param stream_response: stream response
@@ -66,7 +66,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
@classmethod
def convert_stream_simple_response(
cls, stream_response: Generator[AppStreamResponse, None, None]
) -> Generator[dict | str, None, None]:
) -> Generator[dict[str, Any] | str, None, None]:
"""
Convert stream simple response.
:param stream_response: stream response
+3 -1
View File
@@ -1,3 +1,5 @@
from typing import Any
from flask import Flask
from graphon.model_runtime.entities.model_entities import ModelType
from pydantic import BaseModel
@@ -28,7 +30,7 @@ class FreeHostingQuota(HostingQuota):
class HostingProvider(BaseModel):
enabled: bool = False
credentials: dict | None = None
credentials: dict[str, Any] | None = None
quota_unit: QuotaUnit | None = None
quotas: list[HostingQuota] = []
@@ -1,6 +1,7 @@
import tempfile
from binascii import hexlify, unhexlify
from collections.abc import Generator
from typing import Any
from graphon.model_runtime.entities.llm_entities import (
LLMResult,
@@ -226,7 +227,7 @@ class PluginModelBackwardsInvocation(BaseBackwardsInvocation):
# invoke model
response = model_instance.invoke_tts(content_text=payload.content_text, voice=payload.voice)
def handle() -> Generator[dict, None, None]:
def handle() -> Generator[dict[str, Any], None, None]:
for chunk in response:
yield {"result": hexlify(chunk).decode("utf-8")}
+6 -2
View File
@@ -95,9 +95,9 @@ class ExtractProcessor:
) -> list[Document]:
if extract_setting.datasource_type == DatasourceType.FILE:
with tempfile.TemporaryDirectory() as temp_dir:
upload_file = extract_setting.upload_file
if not file_path:
assert extract_setting.upload_file is not None, "upload_file is required"
upload_file: UploadFile = extract_setting.upload_file
assert upload_file is not None, "upload_file is required"
suffix = Path(upload_file.key).suffix
# FIXME mypy: Cannot determine type of 'tempfile._get_candidate_names' better not use it here
file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}" # type: ignore
@@ -113,6 +113,7 @@ class ExtractProcessor:
if file_extension in {".xlsx", ".xls"}:
extractor = ExcelExtractor(file_path)
elif file_extension == ".pdf":
assert upload_file is not None
extractor = PdfExtractor(file_path, upload_file.tenant_id, upload_file.created_by)
elif file_extension in {".md", ".markdown", ".mdx"}:
extractor = (
@@ -123,6 +124,7 @@ class ExtractProcessor:
elif file_extension in {".htm", ".html"}:
extractor = HtmlExtractor(file_path)
elif file_extension == ".docx":
assert upload_file is not None
extractor = WordExtractor(file_path, upload_file.tenant_id, upload_file.created_by)
elif file_extension == ".doc":
extractor = UnstructuredWordExtractor(file_path, unstructured_api_url, unstructured_api_key)
@@ -149,12 +151,14 @@ class ExtractProcessor:
if file_extension in {".xlsx", ".xls"}:
extractor = ExcelExtractor(file_path)
elif file_extension == ".pdf":
assert upload_file is not None
extractor = PdfExtractor(file_path, upload_file.tenant_id, upload_file.created_by)
elif file_extension in {".md", ".markdown", ".mdx"}:
extractor = MarkdownExtractor(file_path, autodetect_encoding=True)
elif file_extension in {".htm", ".html"}:
extractor = HtmlExtractor(file_path)
elif file_extension == ".docx":
assert upload_file is not None
extractor = WordExtractor(file_path, upload_file.tenant_id, upload_file.created_by)
elif file_extension == ".csv":
extractor = CSVExtractor(file_path, autodetect_encoding=True)
@@ -174,21 +174,25 @@ class FirecrawlApp:
return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
def _post_request(self, url, data, headers, retries=3, backoff_factor=0.5) -> httpx.Response:
response: httpx.Response | None = None
for attempt in range(retries):
response = httpx.post(url, headers=headers, json=data)
if response.status_code == 502:
time.sleep(backoff_factor * (2**attempt))
else:
return response
assert response is not None, "retries must be at least 1"
return response
def _get_request(self, url, headers, retries=3, backoff_factor=0.5) -> httpx.Response:
response: httpx.Response | None = None
for attempt in range(retries):
response = httpx.get(url, headers=headers)
if response.status_code == 502:
time.sleep(backoff_factor * (2**attempt))
else:
return response
assert response is not None, "retries must be at least 1"
return response
def _handle_error(self, response, action):
+1 -1
View File
@@ -329,7 +329,7 @@ class EnterpriseMetricHandler:
return
include_content = exporter.include_content
attrs: dict = {
attrs: dict[str, Any] = {
"dify.message.id": payload.get("message_id"),
"dify.tenant_id": envelope.tenant_id,
"dify.event.id": envelope.event_id,
+2 -1
View File
@@ -2,6 +2,7 @@ import logging
import os
from collections.abc import Generator
from pathlib import Path
from typing import Any
import opendal
from dotenv import dotenv_values
@@ -19,7 +20,7 @@ def _get_opendal_kwargs(*, scheme: str, env_file_path: str = ".env", prefix: str
if key.startswith(config_prefix):
kwargs[key[len(config_prefix) :].lower()] = value
file_env_vars: dict = dotenv_values(env_file_path) or {}
file_env_vars: dict[str, Any] = dotenv_values(env_file_path) or {}
for key, value in file_env_vars.items():
if key.startswith(config_prefix) and key[len(config_prefix) :].lower() not in kwargs and value:
kwargs[key[len(config_prefix) :].lower()] = value
+13 -1
View File
@@ -120,10 +120,22 @@ class AppIconUrlField(fields.Raw):
obj = obj["app"]
if isinstance(obj, App | Site) and obj.icon_type == IconType.IMAGE:
return file_helpers.get_signed_file_url(obj.icon)
return build_icon_url(obj.icon_type, obj.icon)
return None
def build_icon_url(icon_type: Any, icon: str | None) -> str | None:
if icon is None or icon_type is None:
return None
from models.model import IconType
icon_type_value = icon_type.value if isinstance(icon_type, IconType) else str(icon_type)
if icon_type_value.lower() != IconType.IMAGE:
return None
return file_helpers.get_signed_file_url(icon)
class AvatarUrlField(fields.Raw):
def output(self, key, obj, **kwargs):
if obj is None:
@@ -0,0 +1,26 @@
"""add qdrant_endpoint to tidb_auth_bindings
Revision ID: 8574b23a38fd
Revises: 6b5f9f8b1a2c
Create Date: 2026-04-14 15:00:00.000000
"""
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision = "8574b23a38fd"
down_revision = "6b5f9f8b1a2c"
branch_labels = None
depends_on = None
def upgrade():
with op.batch_alter_table("tidb_auth_bindings", schema=None) as batch_op:
batch_op.add_column(sa.Column("qdrant_endpoint", sa.String(length=512), nullable=True))
def downgrade():
with op.batch_alter_table("tidb_auth_bindings", schema=None) as batch_op:
batch_op.drop_column("qdrant_endpoint")
+4 -3
View File
@@ -1305,6 +1305,7 @@ class TidbAuthBinding(TypeBase):
)
account: Mapped[str] = mapped_column(String(255), nullable=False)
password: Mapped[str] = mapped_column(String(255), nullable=False)
qdrant_endpoint: Mapped[str | None] = mapped_column(String(512), nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
@@ -1551,7 +1552,7 @@ class PipelineBuiltInTemplate(TypeBase):
name: Mapped[str] = mapped_column(sa.String(255), nullable=False)
description: Mapped[str] = mapped_column(LongText, nullable=False)
chunk_structure: Mapped[str] = mapped_column(sa.String(255), nullable=False)
icon: Mapped[dict] = mapped_column(sa.JSON, nullable=False)
icon: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
yaml_content: Mapped[str] = mapped_column(LongText, nullable=False)
copyright: Mapped[str] = mapped_column(sa.String(255), nullable=False)
privacy_policy: Mapped[str] = mapped_column(sa.String(255), nullable=False)
@@ -1584,7 +1585,7 @@ class PipelineCustomizedTemplate(TypeBase):
name: Mapped[str] = mapped_column(sa.String(255), nullable=False)
description: Mapped[str] = mapped_column(LongText, nullable=False)
chunk_structure: Mapped[str] = mapped_column(sa.String(255), nullable=False)
icon: Mapped[dict] = mapped_column(sa.JSON, nullable=False)
icon: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
yaml_content: Mapped[str] = mapped_column(LongText, nullable=False)
install_count: Mapped[int] = mapped_column(sa.Integer, nullable=False)
@@ -1657,7 +1658,7 @@ class DocumentPipelineExecutionLog(TypeBase):
datasource_type: Mapped[str] = mapped_column(sa.String(255), nullable=False)
datasource_info: Mapped[str] = mapped_column(LongText, nullable=False)
datasource_node_id: Mapped[str] = mapped_column(sa.String(255), nullable=False)
input_data: Mapped[dict] = mapped_column(sa.JSON, nullable=False)
input_data: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
+2 -2
View File
@@ -1007,7 +1007,7 @@ class OAuthProviderApp(TypeBase):
app_icon: Mapped[str] = mapped_column(String(255), nullable=False)
client_id: Mapped[str] = mapped_column(String(255), nullable=False)
client_secret: Mapped[str] = mapped_column(String(255), nullable=False)
app_label: Mapped[dict] = mapped_column(sa.JSON, nullable=False, default_factory=dict)
app_label: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False, default_factory=dict)
redirect_uris: Mapped[list] = mapped_column(sa.JSON, nullable=False, default_factory=list)
scope: Mapped[str] = mapped_column(
String(255),
@@ -2495,7 +2495,7 @@ class TraceAppConfig(TypeBase):
)
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
tracing_provider: Mapped[str | None] = mapped_column(String(255), nullable=True)
tracing_config: Mapped[dict | None] = mapped_column(sa.JSON, nullable=True)
tracing_config: Mapped[dict[str, Any] | None] = mapped_column(sa.JSON, nullable=True)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
+4 -3
View File
@@ -1,4 +1,5 @@
from datetime import datetime
from typing import Any
import sqlalchemy as sa
from sqlalchemy import func
@@ -22,7 +23,7 @@ class DatasourceOauthParamConfig(TypeBase):
)
plugin_id: Mapped[str] = mapped_column(sa.String(255), nullable=False)
provider: Mapped[str] = mapped_column(sa.String(255), nullable=False)
system_credentials: Mapped[dict] = mapped_column(AdjustedJSON, nullable=False)
system_credentials: Mapped[dict[str, Any]] = mapped_column(AdjustedJSON, nullable=False)
class DatasourceProvider(TypeBase):
@@ -40,7 +41,7 @@ class DatasourceProvider(TypeBase):
provider: Mapped[str] = mapped_column(sa.String(128), nullable=False)
plugin_id: Mapped[str] = mapped_column(sa.String(255), nullable=False)
auth_type: Mapped[str] = mapped_column(sa.String(255), nullable=False)
encrypted_credentials: Mapped[dict] = mapped_column(AdjustedJSON, nullable=False)
encrypted_credentials: Mapped[dict[str, Any]] = mapped_column(AdjustedJSON, nullable=False)
avatar_url: Mapped[str] = mapped_column(LongText, nullable=True, default="default")
is_default: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
expires_at: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default="-1", default=-1)
@@ -70,7 +71,7 @@ class DatasourceOauthTenantParamConfig(TypeBase):
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider: Mapped[str] = mapped_column(sa.String(255), nullable=False)
plugin_id: Mapped[str] = mapped_column(sa.String(255), nullable=False)
client_params: Mapped[dict] = mapped_column(AdjustedJSON, nullable=False, default_factory=dict)
client_params: Mapped[dict[str, Any]] = mapped_column(AdjustedJSON, nullable=False, default_factory=dict)
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False)
created_at: Mapped[datetime] = mapped_column(
+1 -1
View File
@@ -25,7 +25,7 @@ class DataSourceOauthBinding(TypeBase):
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
access_token: Mapped[str] = mapped_column(String(255), nullable=False)
provider: Mapped[str] = mapped_column(String(255), nullable=False)
source_info: Mapped[dict] = mapped_column(AdjustedJSON, nullable=False)
source_info: Mapped[dict[str, Any]] = mapped_column(AdjustedJSON, nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
@@ -23,7 +23,7 @@ class ElasticSearchJaVector(ElasticSearchVector):
self,
embeddings: list[list[float]],
metadatas: list[dict[Any, Any]] | None = None,
index_params: dict | None = None,
index_params: dict[str, Any] | None = None,
):
lock_name = f"vector_indexing_lock_{self._collection_name}"
with redis_client.lock(lock_name, timeout=20):
@@ -1,3 +1,4 @@
from typing import Any
from uuid import UUID
from numpy import ndarray
@@ -8,5 +9,5 @@ class CollectionORM(DeclarativeBase):
__tablename__: str
id: Mapped[UUID]
text: Mapped[str]
meta: Mapped[dict]
meta: Mapped[dict[str, Any]]
vector: Mapped[ndarray]
@@ -67,7 +67,7 @@ class PGVectoRS(BaseVector):
primary_key=True,
)
text: Mapped[str]
meta: Mapped[dict] = mapped_column(postgresql.JSONB)
meta: Mapped[dict[str, Any]] = mapped_column(postgresql.JSONB)
vector: Mapped[ndarray] = mapped_column(VECTOR(dim))
self._table = _Table
@@ -1,4 +1,5 @@
import json
import logging
import os
import uuid
from collections.abc import Generator, Iterable, Sequence
@@ -7,6 +8,8 @@ from typing import TYPE_CHECKING, Any
import httpx
import qdrant_client
logger = logging.getLogger(__name__)
from flask import current_app
from httpx import DigestAuth
from pydantic import BaseModel
@@ -421,13 +424,16 @@ class TidbOnQdrantVector(BaseVector):
class TidbOnQdrantVectorFactory(AbstractVectorFactory):
def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TidbOnQdrantVector:
logger.info("init_vector: tenant_id=%s, dataset_id=%s", dataset.tenant_id, dataset.id)
stmt = select(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id)
tidb_auth_binding = db.session.scalars(stmt).one_or_none()
if not tidb_auth_binding:
logger.info("No existing TidbAuthBinding for tenant %s, acquiring lock", dataset.tenant_id)
with redis_client.lock("create_tidb_serverless_cluster_lock", timeout=900):
stmt = select(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id)
tidb_auth_binding = db.session.scalars(stmt).one_or_none()
if tidb_auth_binding:
logger.info("Found binding after lock: cluster_id=%s", tidb_auth_binding.cluster_id)
TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}"
else:
@@ -437,11 +443,18 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
.limit(1)
)
if idle_tidb_auth_binding:
logger.info(
"Assigning idle cluster %s to tenant %s",
idle_tidb_auth_binding.cluster_id,
dataset.tenant_id,
)
idle_tidb_auth_binding.active = True
idle_tidb_auth_binding.tenant_id = dataset.tenant_id
db.session.commit()
tidb_auth_binding = idle_tidb_auth_binding
TIDB_ON_QDRANT_API_KEY = f"{idle_tidb_auth_binding.account}:{idle_tidb_auth_binding.password}"
else:
logger.info("No idle clusters available, creating new cluster for tenant %s", dataset.tenant_id)
new_cluster = TidbService.create_tidb_serverless_cluster(
dify_config.TIDB_PROJECT_ID or "",
dify_config.TIDB_API_URL or "",
@@ -450,21 +463,39 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
dify_config.TIDB_PRIVATE_KEY or "",
dify_config.TIDB_REGION or "",
)
logger.info(
"New cluster created: cluster_id=%s, qdrant_endpoint=%s",
new_cluster["cluster_id"],
new_cluster.get("qdrant_endpoint"),
)
new_tidb_auth_binding = TidbAuthBinding(
cluster_id=new_cluster["cluster_id"],
cluster_name=new_cluster["cluster_name"],
account=new_cluster["account"],
password=new_cluster["password"],
qdrant_endpoint=new_cluster.get("qdrant_endpoint"),
tenant_id=dataset.tenant_id,
active=True,
status=TidbAuthBindingStatus.ACTIVE,
)
db.session.add(new_tidb_auth_binding)
db.session.commit()
tidb_auth_binding = new_tidb_auth_binding
TIDB_ON_QDRANT_API_KEY = f"{new_tidb_auth_binding.account}:{new_tidb_auth_binding.password}"
else:
logger.info("Existing binding found: cluster_id=%s", tidb_auth_binding.cluster_id)
TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}"
qdrant_url = (
(tidb_auth_binding.qdrant_endpoint if tidb_auth_binding else None) or dify_config.TIDB_ON_QDRANT_URL or ""
)
logger.info(
"Using qdrant endpoint: %s (from_binding=%s, fallback_global=%s)",
qdrant_url,
tidb_auth_binding.qdrant_endpoint if tidb_auth_binding else None,
dify_config.TIDB_ON_QDRANT_URL,
)
if dataset.index_struct_dict:
class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
collection_name = class_prefix
@@ -479,7 +510,7 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
collection_name=collection_name,
group_id=dataset.id,
config=TidbOnQdrantConfig(
endpoint=dify_config.TIDB_ON_QDRANT_URL or "",
endpoint=qdrant_url,
api_key=TIDB_ON_QDRANT_API_KEY,
root_path=str(config.root_path),
timeout=dify_config.TIDB_ON_QDRANT_CLIENT_TIMEOUT,
@@ -1,3 +1,4 @@
import logging
import time
import uuid
from collections.abc import Sequence
@@ -12,6 +13,8 @@ from extensions.ext_redis import redis_client
from models.dataset import TidbAuthBinding
from models.enums import TidbAuthBindingStatus
logger = logging.getLogger(__name__)
# Reuse a pooled HTTP client for all TiDB Cloud requests to minimize connection churn
_tidb_http_client: httpx.Client = get_pooled_http_client(
"tidb:cloud",
@@ -20,6 +23,46 @@ _tidb_http_client: httpx.Client = get_pooled_http_client(
class TidbService:
@staticmethod
def extract_qdrant_endpoint(cluster_response: dict) -> str | None:
"""Extract the qdrant endpoint URL from a Get Cluster API response.
Reads ``endpoints.public.host`` (e.g. ``gateway01.xx.tidbcloud.com``),
prepends ``qdrant-`` and wraps it as an ``https://`` URL.
"""
endpoints = cluster_response.get("endpoints") or {}
public = endpoints.get("public") or {}
host = public.get("host")
if host:
return f"https://qdrant-{host}"
return None
@staticmethod
def fetch_qdrant_endpoint(api_url: str, public_key: str, private_key: str, cluster_id: str) -> str | None:
"""Call Get Cluster API and extract the qdrant endpoint.
Use ``extract_qdrant_endpoint`` instead when you already have
the cluster response to avoid a redundant API call.
"""
try:
logger.info("Fetching qdrant endpoint for cluster %s", cluster_id)
cluster_response = TidbService.get_tidb_serverless_cluster(api_url, public_key, private_key, cluster_id)
if not cluster_response:
logger.warning("Empty response from Get Cluster API for cluster %s", cluster_id)
return None
qdrant_url = TidbService.extract_qdrant_endpoint(cluster_response)
if qdrant_url:
logger.info("Resolved qdrant endpoint for cluster %s: %s", cluster_id, qdrant_url)
return qdrant_url
logger.warning(
"No endpoints.public.host found for cluster %s, response keys: %s",
cluster_id,
list(cluster_response.keys()),
)
except Exception:
logger.exception("Failed to fetch qdrant endpoint for cluster %s", cluster_id)
return None
@staticmethod
def create_tidb_serverless_cluster(
project_id: str, api_url: str, iam_url: str, public_key: str, private_key: str, region: str
@@ -57,6 +100,7 @@ class TidbService:
"rootPassword": password,
}
logger.info("Creating TiDB serverless cluster: display_name=%s, region=%s", display_name, region)
response = _tidb_http_client.post(
f"{api_url}/clusters", json=cluster_data, auth=DigestAuth(public_key, private_key)
)
@@ -64,21 +108,39 @@ class TidbService:
if response.status_code == 200:
response_data = response.json()
cluster_id = response_data["clusterId"]
logger.info("Cluster created, cluster_id=%s, waiting for ACTIVE state", cluster_id)
retry_count = 0
max_retries = 30
while retry_count < max_retries:
cluster_response = TidbService.get_tidb_serverless_cluster(api_url, public_key, private_key, cluster_id)
if cluster_response["state"] == "ACTIVE":
user_prefix = cluster_response["userPrefix"]
qdrant_endpoint = TidbService.extract_qdrant_endpoint(cluster_response)
logger.info(
"Cluster %s is ACTIVE, user_prefix=%s, qdrant_endpoint=%s",
cluster_id,
user_prefix,
qdrant_endpoint,
)
return {
"cluster_id": cluster_id,
"cluster_name": display_name,
"account": f"{user_prefix}.root",
"password": password,
"qdrant_endpoint": qdrant_endpoint,
}
time.sleep(30) # wait 30 seconds before retrying
logger.info(
"Cluster %s state=%s, retry %d/%d",
cluster_id,
cluster_response["state"],
retry_count + 1,
max_retries,
)
time.sleep(30)
retry_count += 1
logger.error("Cluster %s did not become ACTIVE after %d retries", cluster_id, max_retries)
else:
logger.error("Failed to create cluster: status=%d, body=%s", response.status_code, response.text)
response.raise_for_status()
@staticmethod
@@ -243,19 +305,29 @@ class TidbService:
if response.status_code == 200:
response_data = response.json()
cluster_infos = []
logger.info("Batch created %d clusters", len(response_data.get("clusters", [])))
for item in response_data["clusters"]:
cache_key = f"tidb_serverless_cluster_password:{item['displayName']}"
cached_password = redis_client.get(cache_key)
if not cached_password:
logger.warning("No cached password for cluster %s, skipping", item["displayName"])
continue
qdrant_endpoint = TidbService.fetch_qdrant_endpoint(api_url, public_key, private_key, item["clusterId"])
logger.info(
"Batch cluster %s: qdrant_endpoint=%s",
item["clusterId"],
qdrant_endpoint,
)
cluster_info = {
"cluster_id": item["clusterId"],
"cluster_name": item["displayName"],
"account": "root",
"password": cached_password.decode("utf-8"),
"qdrant_endpoint": qdrant_endpoint,
}
cluster_infos.append(cluster_info)
return cluster_infos
else:
logger.error("Batch create failed: status=%d, body=%s", response.status_code, response.text)
response.raise_for_status()
return []
@@ -114,14 +114,12 @@ class TestTidbOnQdrantVectorDeleteByIds:
assert exc_info.value.status_code == 500
def test_delete_by_ids_with_large_batch(self, vector_instance):
"""Test deletion with a large batch of IDs."""
# Create 1000 IDs
def test_delete_by_ids_with_exactly_1000(self, vector_instance):
"""Test deletion with exactly 1000 IDs triggers a single batch."""
ids = [f"doc_{i}" for i in range(1000)]
vector_instance.delete_by_ids(ids)
# Verify single delete call with all IDs
vector_instance._client.delete.assert_called_once()
call_args = vector_instance._client.delete.call_args
@@ -129,11 +127,28 @@ class TestTidbOnQdrantVectorDeleteByIds:
filter_obj = filter_selector.filter
field_condition = filter_obj.must[0]
# Verify all 1000 IDs are in the batch
assert len(field_condition.match.any) == 1000
assert "doc_0" in field_condition.match.any
assert "doc_999" in field_condition.match.any
def test_delete_by_ids_splits_into_batches(self, vector_instance):
"""Test deletion with >1000 IDs triggers multiple batched calls."""
ids = [f"doc_{i}" for i in range(2500)]
vector_instance.delete_by_ids(ids)
assert vector_instance._client.delete.call_count == 3
batches = []
for call in vector_instance._client.delete.call_args_list:
filter_selector = call[1]["points_selector"]
field_condition = filter_selector.filter.must[0]
batches.append(field_condition.match.any)
assert len(batches[0]) == 1000
assert len(batches[1]) == 1000
assert len(batches[2]) == 500
def test_delete_by_ids_filter_structure(self, vector_instance):
"""Test that the filter structure is correctly constructed."""
ids = ["doc1", "doc2"]
@@ -157,3 +172,57 @@ class TestTidbOnQdrantVectorDeleteByIds:
# Verify MatchAny structure
assert isinstance(field_condition.match, rest.MatchAny)
assert field_condition.match.any == ids
class TestInitVectorEndpointSelection:
"""Test that init_vector selects the correct qdrant endpoint.
We avoid importing the full module (which triggers Flask app context)
by testing the endpoint selection logic directly on TidbOnQdrantConfig.
"""
def test_uses_binding_endpoint_when_present(self):
binding_endpoint = "https://qdrant-custom.tidb.com"
global_url = "https://qdrant-global.tidb.com"
qdrant_url = binding_endpoint or global_url or ""
assert qdrant_url == "https://qdrant-custom.tidb.com"
config = TidbOnQdrantConfig(endpoint=qdrant_url)
assert config.endpoint == "https://qdrant-custom.tidb.com"
def test_falls_back_to_global_when_binding_endpoint_is_none(self):
binding_endpoint = None
global_url = "https://qdrant-global.tidb.com"
qdrant_url = binding_endpoint or global_url or ""
assert qdrant_url == "https://qdrant-global.tidb.com"
config = TidbOnQdrantConfig(endpoint=qdrant_url)
assert config.endpoint == "https://qdrant-global.tidb.com"
def test_falls_back_to_empty_when_both_none(self):
binding_endpoint = None
global_url = None
qdrant_url = binding_endpoint or global_url or ""
assert qdrant_url == ""
config = TidbOnQdrantConfig(endpoint=qdrant_url)
assert config.endpoint == ""
def test_binding_endpoint_takes_precedence_over_global(self):
binding_endpoint = "https://qdrant-ap-southeast.tidb.com"
global_url = "https://qdrant-us-east.tidb.com"
qdrant_url = binding_endpoint or global_url or ""
assert qdrant_url == "https://qdrant-ap-southeast.tidb.com"
def test_empty_string_binding_endpoint_falls_back_to_global(self):
binding_endpoint = ""
global_url = "https://qdrant-global.tidb.com"
qdrant_url = binding_endpoint or global_url or ""
assert qdrant_url == "https://qdrant-global.tidb.com"
@@ -0,0 +1,218 @@
from unittest.mock import MagicMock, patch
import pytest
from dify_vdb_tidb_on_qdrant.tidb_service import TidbService
class TestExtractQdrantEndpoint:
"""Unit tests for TidbService.extract_qdrant_endpoint."""
def test_returns_endpoint_when_host_present(self):
response = {"endpoints": {"public": {"host": "gateway01.us-east-1.tidbcloud.com", "port": 4000}}}
result = TidbService.extract_qdrant_endpoint(response)
assert result == "https://qdrant-gateway01.us-east-1.tidbcloud.com"
def test_returns_none_when_host_missing(self):
response = {"endpoints": {"public": {}}}
assert TidbService.extract_qdrant_endpoint(response) is None
def test_returns_none_when_public_missing(self):
response = {"endpoints": {}}
assert TidbService.extract_qdrant_endpoint(response) is None
def test_returns_none_when_endpoints_missing(self):
assert TidbService.extract_qdrant_endpoint({}) is None
class TestFetchQdrantEndpoint:
"""Unit tests for TidbService.fetch_qdrant_endpoint."""
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_endpoint_when_host_present(self, mock_get_cluster):
mock_get_cluster.return_value = {
"endpoints": {"public": {"host": "gateway01.us-east-1.tidbcloud.com", "port": 4000}}
}
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result == "https://qdrant-gateway01.us-east-1.tidbcloud.com"
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_cluster_response_is_none(self, mock_get_cluster):
mock_get_cluster.return_value = None
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_host_missing(self, mock_get_cluster):
mock_get_cluster.return_value = {"endpoints": {"public": {}}}
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_endpoints_missing(self, mock_get_cluster):
mock_get_cluster.return_value = {}
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_on_exception(self, mock_get_cluster):
mock_get_cluster.side_effect = RuntimeError("network error")
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
class TestCreateTidbServerlessClusterQdrantEndpoint:
"""Verify that create_tidb_serverless_cluster includes qdrant_endpoint in its result."""
@patch.object(TidbService, "get_tidb_serverless_cluster")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_result_contains_qdrant_endpoint(self, mock_config, mock_http, mock_get_cluster):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
mock_get_cluster.return_value = {
"state": "ACTIVE",
"userPrefix": "pfx",
"endpoints": {"public": {"host": "gw.tidbcloud.com", "port": 4000}},
}
result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
assert result is not None
assert result["qdrant_endpoint"] == "https://qdrant-gw.tidbcloud.com"
@patch.object(TidbService, "get_tidb_serverless_cluster")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_result_qdrant_endpoint_none_when_no_endpoints(self, mock_config, mock_http, mock_get_cluster):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
mock_get_cluster.return_value = {"state": "ACTIVE", "userPrefix": "pfx"}
result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
assert result is not None
assert result["qdrant_endpoint"] is None
class TestBatchCreateTidbServerlessClusterQdrantEndpoint:
"""Verify that batch_create includes qdrant_endpoint per cluster."""
@patch.object(TidbService, "fetch_qdrant_endpoint", return_value="https://qdrant-gw.tidbcloud.com")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.redis_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_batch_result_contains_qdrant_endpoint(self, mock_config, mock_http, mock_redis, mock_fetch_ep):
mock_config.TIDB_SPEND_LIMIT = 10
cluster_name = "abc123"
mock_http.post.return_value = MagicMock(
status_code=200,
json=lambda: {"clusters": [{"clusterId": "c-1", "displayName": cluster_name}]},
)
mock_redis.setex = MagicMock()
mock_redis.get.return_value = b"password123"
result = TidbService.batch_create_tidb_serverless_cluster(
batch_size=1,
project_id="proj",
api_url="url",
iam_url="iam",
public_key="pub",
private_key="priv",
region="us-east-1",
)
assert len(result) == 1
assert result[0]["qdrant_endpoint"] == "https://qdrant-gw.tidbcloud.com"
class TestCreateTidbServerlessClusterRetry:
"""Cover retry/logging paths in create_tidb_serverless_cluster."""
@patch.object(TidbService, "get_tidb_serverless_cluster")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_polls_until_active(self, mock_config, mock_http, mock_get_cluster):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
mock_get_cluster.side_effect = [
{"state": "CREATING", "userPrefix": ""},
{"state": "ACTIVE", "userPrefix": "pfx", "endpoints": {"public": {"host": "gw.tidb.com"}}},
]
with patch("dify_vdb_tidb_on_qdrant.tidb_service.time.sleep"):
result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
assert result is not None
assert result["qdrant_endpoint"] == "https://qdrant-gw.tidb.com"
assert mock_get_cluster.call_count == 2
@patch.object(TidbService, "get_tidb_serverless_cluster")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_returns_none_after_max_retries(self, mock_config, mock_http, mock_get_cluster):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
mock_get_cluster.return_value = {"state": "CREATING", "userPrefix": ""}
with patch("dify_vdb_tidb_on_qdrant.tidb_service.time.sleep"):
result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
assert result is None
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_raises_on_post_failure(self, mock_config, mock_http):
mock_config.TIDB_SPEND_LIMIT = 10
mock_response = MagicMock(status_code=400, text="Bad Request")
mock_response.raise_for_status.side_effect = Exception("HTTP 400")
mock_http.post.return_value = mock_response
with pytest.raises(Exception, match="HTTP 400"):
TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
class TestBatchCreateEdgeCases:
"""Cover logging/edge-case branches in batch_create."""
@patch.object(TidbService, "fetch_qdrant_endpoint", return_value=None)
@patch("dify_vdb_tidb_on_qdrant.tidb_service.redis_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_skips_cluster_when_no_cached_password(self, mock_config, mock_http, mock_redis, mock_fetch_ep):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(
status_code=200,
json=lambda: {"clusters": [{"clusterId": "c-1", "displayName": "name1"}]},
)
mock_redis.setex = MagicMock()
mock_redis.get.return_value = None
result = TidbService.batch_create_tidb_serverless_cluster(
batch_size=1,
project_id="proj",
api_url="url",
iam_url="iam",
public_key="pub",
private_key="priv",
region="us-east-1",
)
assert len(result) == 0
mock_fetch_ep.assert_not_called()
@patch("dify_vdb_tidb_on_qdrant.tidb_service.redis_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_raises_on_post_failure(self, mock_config, mock_http, mock_redis):
mock_config.TIDB_SPEND_LIMIT = 10
mock_response = MagicMock(status_code=500, text="Server Error")
mock_response.raise_for_status.side_effect = Exception("HTTP 500")
mock_http.post.return_value = mock_response
mock_redis.setex = MagicMock()
with pytest.raises(Exception, match="HTTP 500"):
TidbService.batch_create_tidb_serverless_cluster(
batch_size=1,
project_id="proj",
api_url="url",
iam_url="iam",
public_key="pub",
private_key="priv",
region="us-east-1",
)
+98 -138
View File
@@ -4,92 +4,52 @@ version = "1.13.3"
requires-python = "~=3.12.0"
dependencies = [
"aliyun-log-python-sdk~=0.9.44",
"arize-phoenix-otel~=0.15.0",
"azure-identity==1.25.3",
"beautifulsoup4==4.14.3",
"boto3==1.42.88",
"bs4~=0.0.1",
"cachetools~=7.0.5",
"celery~=5.6.3",
"charset-normalizer>=3.4.7",
"flask~=3.1.3",
"flask-compress>=1.24,<1.25",
"flask-cors~=6.0.2",
"flask-login~=0.6.3",
"flask-migrate~=4.1.0",
"flask-orjson~=2.0.0",
"flask-sqlalchemy~=3.1.1",
"gevent~=26.4.0",
"gmpy2~=2.3.0",
"google-api-core>=2.30.3",
"google-api-python-client==2.194.0",
"google-auth>=2.49.2",
"google-auth-httplib2==0.3.1",
"google-cloud-aiplatform>=1.147.0",
"googleapis-common-protos>=1.74.0",
"graphon>=0.1.2",
"gunicorn~=25.3.0",
"httpx[socks]~=0.28.1",
"jieba==0.42.1",
"json-repair>=0.59.2",
"langfuse>=4.2.0,<5.0.0",
"langsmith~=0.7.30",
"markdown~=3.10.2",
"mlflow-skinny>=3.11.1",
"numpy~=2.4.4",
"openpyxl~=3.1.5",
"opik~=1.11.2",
"litellm==1.83.0", # Pinned to avoid madoka dependency issue
"opentelemetry-api==1.41.0",
"opentelemetry-distro==0.62b0",
"opentelemetry-exporter-otlp==1.41.0",
"opentelemetry-exporter-otlp-proto-common==1.41.0",
"opentelemetry-exporter-otlp-proto-grpc==1.41.0",
"opentelemetry-exporter-otlp-proto-http==1.41.0",
"opentelemetry-instrumentation==0.62b0",
"opentelemetry-instrumentation-celery==0.62b0",
"opentelemetry-instrumentation-flask==0.62b0",
"opentelemetry-instrumentation-httpx==0.62b0",
"opentelemetry-instrumentation-redis==0.62b0",
"opentelemetry-instrumentation-sqlalchemy==0.62b0",
"opentelemetry-propagator-b3==1.41.0",
"opentelemetry-proto==1.41.0",
"opentelemetry-sdk==1.41.0",
"opentelemetry-semantic-conventions==0.62b0",
"opentelemetry-util-http==0.62b0",
"pandas[excel,output-formatting,performance]~=3.0.2",
"psycogreen~=1.0.2",
"psycopg2-binary~=2.9.11",
"pycryptodome==3.23.0",
"pydantic~=2.12.5",
"pydantic-settings~=2.13.1",
"pyjwt~=2.12.1",
"pypdfium2==5.6.0",
"python-docx~=1.2.0",
"python-dotenv==1.2.2",
"pyyaml~=6.0.1",
"readabilipy~=0.3.0",
"redis[hiredis]~=7.4.0",
"resend~=2.27.0",
"sentry-sdk[flask]~=2.57.0",
"sqlalchemy~=2.0.49",
"starlette==1.0.0",
"tiktoken~=0.12.0",
"transformers~=5.3.0",
"unstructured[docx,epub,md,ppt,pptx]~=0.21.5",
"pypandoc~=1.13",
"yarl~=1.23.0",
"sseclient-py~=1.9.0",
"httpx-sse~=0.4.0",
"sendgrid~=6.12.5",
"flask-restx~=1.3.2",
"packaging~=26.0",
# Legacy: mature and widely deployed
"bleach>=6.3.0",
"boto3>=1.42.88",
"celery>=5.6.3",
"croniter>=6.2.2",
"apscheduler>=3.11.2",
"weave>=0.52.36",
"fastopenapi[flask]>=0.7.0",
"bleach~=6.3.0",
"flask-cors>=6.0.2",
"gevent>=26.4.0",
"gmpy2>=2.3.0",
"google-api-python-client>=2.194.0",
"gunicorn>=25.3.0",
"psycogreen>=1.0.2",
"psycopg2-binary>=2.9.11",
"redis[hiredis]>=7.4.0",
"sendgrid>=6.12.5",
# Stable: production-proven, cap below the next major
"aliyun-log-python-sdk>=0.9.44,<1.0.0",
"azure-identity>=1.25.3,<2.0.0",
"flask-compress>=1.24,<2.0.0",
"flask-login>=0.6.3,<1.0.0",
"flask-migrate>=4.1.0,<5.0.0",
"flask-orjson>=2.0.0,<3.0.0",
"flask-restx>=1.3.2,<2.0.0",
"google-cloud-aiplatform>=1.147.0,<2.0.0",
"httpx[socks]>=0.28.1,<1.0.0",
"langfuse>=4.2.0,<5.0.0",
"langsmith>=0.7.30,<1.0.0",
"mlflow-skinny>=3.11.1,<4.0.0",
"opentelemetry-distro>=0.62b0,<1.0.0",
"opentelemetry-instrumentation-celery>=0.62b0,<1.0.0",
"opentelemetry-instrumentation-flask>=0.62b0,<1.0.0",
"opentelemetry-instrumentation-httpx>=0.62b0,<1.0.0",
"opentelemetry-instrumentation-redis>=0.62b0,<1.0.0",
"opentelemetry-instrumentation-sqlalchemy>=0.62b0,<1.0.0",
"opentelemetry-propagator-b3>=1.41.0,<2.0.0",
"readabilipy>=0.3.0,<1.0.0",
"resend>=2.27.0,<3.0.0",
"weave>=0.52.36,<1.0.0",
# Emerging: newer and fast-moving, use compatible pins
"arize-phoenix-otel~=0.15.0",
"fastopenapi[flask]~=0.7.0",
"graphon~=0.1.2",
"httpx-sse~=0.4.0",
"json-repair~=0.59.2",
"opik~=1.11.2",
]
# Before adding new dependency, consider place it in
# alphabet order (a-z) and suitable group.
@@ -147,46 +107,46 @@ override-dependencies = [
# Required for development and running tests
############################################################
dev = [
"coverage~=7.13.4",
"dotenv-linter~=0.7.0",
"faker~=40.13.0",
"lxml-stubs~=0.5.1",
"basedpyright~=1.39.0",
"ruff~=0.15.10",
"pytest~=9.0.3",
"pytest-benchmark~=5.2.3",
"pytest-cov~=7.1.0",
"pytest-env~=1.6.0",
"pytest-mock~=3.15.1",
"testcontainers~=4.14.2",
"types-aiofiles~=25.1.0",
"types-beautifulsoup4~=4.12.0",
"types-cachetools~=6.2.0",
"types-colorama~=0.4.15",
"types-defusedxml~=0.7.0",
"types-deprecated~=1.3.1",
"types-docutils~=0.22.3",
"types-flask-cors~=6.0.0",
"types-flask-migrate~=4.1.0",
"types-gevent~=26.4.0",
"types-greenlet~=3.4.0",
"types-html5lib~=1.1.11",
"types-markdown~=3.10.2",
"types-oauthlib~=3.3.0",
"types-objgraph~=3.6.0",
"types-olefile~=0.47.0",
"types-openpyxl~=3.1.5",
"types-pexpect~=4.9.0",
"types-protobuf~=7.34.1",
"types-psutil~=7.2.2",
"types-psycopg2~=2.9.21",
"types-pygments~=2.20.0",
"types-pymysql~=1.1.0",
"types-python-dateutil~=2.9.0",
"types-pywin32~=311.0.0",
"types-pyyaml~=6.0.12",
"types-regex~=2026.4.4",
"types-shapely~=2.1.0",
"coverage>=7.13.4",
"dotenv-linter>=0.7.0",
"faker>=20.1.0",
"lxml-stubs>=0.5.1",
"basedpyright>=1.39.0",
"ruff>=0.15.10",
"pytest>=9.0.3",
"pytest-benchmark>=5.2.3",
"pytest-cov>=7.1.0",
"pytest-env>=1.6.0",
"pytest-mock>=3.15.1",
"testcontainers>=4.14.2",
"types-aiofiles>=25.1.0",
"types-beautifulsoup4>=4.12.0",
"types-cachetools>=6.2.0",
"types-colorama>=0.4.15",
"types-defusedxml>=0.7.0",
"types-deprecated>=1.3.1",
"types-docutils>=0.22.3",
"types-flask-cors>=6.0.0",
"types-flask-migrate>=4.1.0",
"types-gevent>=26.4.0",
"types-greenlet>=3.4.0",
"types-html5lib>=1.1.11",
"types-markdown>=3.10.2",
"types-oauthlib>=3.3.0",
"types-objgraph>=3.6.0",
"types-olefile>=0.47.0",
"types-openpyxl>=3.1.5",
"types-pexpect>=4.9.0",
"types-protobuf>=7.34.1",
"types-psutil>=7.2.2",
"types-psycopg2>=2.9.21",
"types-pygments>=2.20.0",
"types-pymysql>=1.1.0",
"types-python-dateutil>=2.9.0",
"types-pywin32>=311.0.0",
"types-pyyaml>=6.0.12",
"types-regex>=2026.4.4",
"types-shapely>=2.1.0",
"types-simplejson>=3.20.0.20260408",
"types-six>=1.17.0.20260408",
"types-tensorflow>=2.18.0.20260408",
@@ -198,19 +158,19 @@ dev = [
"types_pyOpenSSL>=24.1.0",
"types_cffi>=2.0.0.20260408",
"types_setuptools>=82.0.0.20260408",
"pandas-stubs~=3.0.0",
"pandas-stubs>=3.0.0",
"scipy-stubs>=1.15.3.0",
"types-python-http-client>=3.3.7.20260408",
"import-linter>=2.3",
"types-redis>=4.6.0.20241004",
"celery-types>=0.23.0",
"mypy~=1.20.1",
"mypy>=1.20.1",
# "locust>=2.40.4", # Temporarily removed due to compatibility issues. Uncomment when resolved.
"sseclient-py>=1.8.0",
"pytest-timeout>=2.4.0",
"pytest-xdist>=3.8.0",
"pyrefly>=0.60.0",
"xinference-client~=2.4.0",
"xinference-client>=2.4.0",
]
############################################################
@@ -218,21 +178,21 @@ dev = [
# Required for storage clients
############################################################
storage = [
"azure-storage-blob==12.28.0",
"bce-python-sdk~=0.9.69",
"cos-python-sdk-v5==1.9.41",
"esdk-obs-python==3.26.2",
"azure-storage-blob>=12.28.0",
"bce-python-sdk>=0.9.69",
"cos-python-sdk-v5>=1.9.41",
"esdk-obs-python>=3.22.2",
"google-cloud-storage>=3.10.1",
"opendal~=0.46.0",
"oss2==2.19.1",
"supabase~=2.18.1",
"tos~=2.9.0",
"opendal>=0.46.0",
"oss2>=2.19.1",
"supabase>=2.18.1",
"tos>=2.9.0",
]
############################################################
# [ Tools ] dependency group
############################################################
tools = ["cloudscraper~=1.2.71", "nltk~=3.9.1"]
tools = ["cloudscraper>=1.2.71", "nltk>=3.9.1"]
############################################################
# [ VDB ] workspace plugins — hollow packages under providers/vdb/*
@@ -302,7 +262,7 @@ vdb-vastbase = ["dify-vdb-vastbase"]
vdb-vikingdb = ["dify-vdb-vikingdb"]
vdb-weaviate = ["dify-vdb-weaviate"]
# Optional client used by some tests / integrations (not a vector backend plugin)
vdb-xinference = ["xinference-client~=2.4.0"]
vdb-xinference = ["xinference-client>=2.4.0"]
[tool.pyrefly]
project-includes = ["."]
@@ -57,6 +57,7 @@ def create_clusters(batch_size):
cluster_name=new_cluster["cluster_name"],
account=new_cluster["account"],
password=new_cluster["password"],
qdrant_endpoint=new_cluster.get("qdrant_endpoint"),
active=False,
status=TidbAuthBindingStatus.CREATING,
)
+1 -1
View File
@@ -455,7 +455,7 @@ class AppDslService:
app.updated_by = account.id
self._session.add(app)
self._session.commit()
self._session.flush()
app_was_created.send(app, account=account)
# save dependencies
@@ -87,7 +87,7 @@ class RetrievalModel(BaseModel):
class MetaDataConfig(BaseModel):
doc_type: str
doc_metadata: dict
doc_metadata: dict[str, Any]
class KnowledgeConfig(BaseModel):
+18 -14
View File
@@ -104,7 +104,7 @@ class RagPipelineService:
self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
@classmethod
def get_pipeline_templates(cls, type: str = "built-in", language: str = "en-US") -> dict:
def get_pipeline_templates(cls, type: str = "built-in", language: str = "en-US") -> dict[str, Any]:
if type == "built-in":
mode = dify_config.HOSTED_FETCH_PIPELINE_TEMPLATES_MODE
retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
@@ -120,7 +120,7 @@ class RagPipelineService:
return result
@classmethod
def get_pipeline_template_detail(cls, template_id: str, type: str = "built-in") -> dict | None:
def get_pipeline_template_detail(cls, template_id: str, type: str = "built-in") -> dict[str, Any] | None:
"""
Get pipeline template detail.
@@ -131,7 +131,7 @@ class RagPipelineService:
if type == "built-in":
mode = dify_config.HOSTED_FETCH_PIPELINE_TEMPLATES_MODE
retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
built_in_result: dict | None = retrieval_instance.get_pipeline_template_detail(template_id)
built_in_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(template_id)
if built_in_result is None:
logger.warning(
"pipeline template retrieval returned empty result, template_id: %s, mode: %s",
@@ -142,7 +142,7 @@ class RagPipelineService:
else:
mode = "customized"
retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)()
customized_result: dict | None = retrieval_instance.get_pipeline_template_detail(template_id)
customized_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail(template_id)
return customized_result
@classmethod
@@ -297,7 +297,7 @@ class RagPipelineService:
self,
*,
pipeline: Pipeline,
graph: dict,
graph: dict[str, Any],
unique_hash: str | None,
account: Account,
environment_variables: Sequence[VariableBase],
@@ -467,7 +467,9 @@ class RagPipelineService:
return default_block_configs
def get_default_block_config(self, node_type: str, filters: dict | None = None) -> Mapping[str, object] | None:
def get_default_block_config(
self, node_type: str, filters: dict[str, Any] | None = None
) -> Mapping[str, object] | None:
"""
Get default config of node.
:param node_type: node type
@@ -500,7 +502,7 @@ class RagPipelineService:
return default_config
def run_draft_workflow_node(
self, pipeline: Pipeline, node_id: str, user_inputs: dict, account: Account
self, pipeline: Pipeline, node_id: str, user_inputs: dict[str, Any], account: Account
) -> WorkflowNodeExecutionModel | None:
"""
Run draft workflow node
@@ -582,7 +584,7 @@ class RagPipelineService:
self,
pipeline: Pipeline,
node_id: str,
user_inputs: dict,
user_inputs: dict[str, Any],
account: Account,
datasource_type: str,
is_published: bool,
@@ -749,7 +751,7 @@ class RagPipelineService:
self,
pipeline: Pipeline,
node_id: str,
user_inputs: dict,
user_inputs: dict[str, Any],
account: Account,
datasource_type: str,
is_published: bool,
@@ -979,7 +981,7 @@ class RagPipelineService:
return workflow_node_execution
def update_workflow(
self, *, session: Session, workflow_id: str, tenant_id: str, account_id: str, data: dict
self, *, session: Session, workflow_id: str, tenant_id: str, account_id: str, data: dict[str, Any]
) -> Workflow | None:
"""
Update workflow attributes
@@ -1099,7 +1101,9 @@ class RagPipelineService:
]
return datasource_provider_variables
def get_rag_pipeline_paginate_workflow_runs(self, pipeline: Pipeline, args: dict) -> InfiniteScrollPagination:
def get_rag_pipeline_paginate_workflow_runs(
self, pipeline: Pipeline, args: dict[str, Any]
) -> InfiniteScrollPagination:
"""
Get debug workflow run list
Only return triggered_from == debugging
@@ -1169,7 +1173,7 @@ class RagPipelineService:
return list(node_executions)
@classmethod
def publish_customized_pipeline_template(cls, pipeline_id: str, args: dict):
def publish_customized_pipeline_template(cls, pipeline_id: str, args: dict[str, Any]):
"""
Publish customized pipeline template
"""
@@ -1259,7 +1263,7 @@ class RagPipelineService:
)
return node_exec
def set_datasource_variables(self, pipeline: Pipeline, args: dict, current_user: Account):
def set_datasource_variables(self, pipeline: Pipeline, args: dict[str, Any], current_user: Account):
"""
Set datasource variables
"""
@@ -1346,7 +1350,7 @@ class RagPipelineService:
)
return workflow_node_execution_db_model
def get_recommended_plugins(self, type: str) -> dict:
def get_recommended_plugins(self, type: str) -> dict[str, Any]:
# Query active recommended plugins
stmt = select(PipelineRecommendedPlugin).where(PipelineRecommendedPlugin.active == True)
if type and type != "all":
+8 -8
View File
@@ -241,8 +241,8 @@ class WorkflowService:
self,
*,
app_model: App,
graph: dict,
features: dict,
graph: dict[str, Any],
features: dict[str, Any],
unique_hash: str | None,
account: Account,
environment_variables: Sequence[VariableBase],
@@ -576,7 +576,7 @@ class WorkflowService:
except Exception as e:
raise ValueError(f"Failed to validate default credential for tool provider {provider}: {str(e)}")
def _validate_load_balancing_credentials(self, workflow: Workflow, node_data: dict, node_id: str) -> None:
def _validate_load_balancing_credentials(self, workflow: Workflow, node_data: dict[str, Any], node_id: str) -> None:
"""
Validate load balancing credentials for a workflow node.
@@ -1214,7 +1214,7 @@ class WorkflowService:
return variable_pool
def run_free_workflow_node(
self, node_data: dict, tenant_id: str, user_id: str, node_id: str, user_inputs: dict[str, Any]
self, node_data: dict[str, Any], tenant_id: str, user_id: str, node_id: str, user_inputs: dict[str, Any]
) -> WorkflowNodeExecution:
"""
Run free workflow node
@@ -1361,7 +1361,7 @@ class WorkflowService:
node_execution.status = WorkflowNodeExecutionStatus.FAILED
node_execution.error = error
def convert_to_workflow(self, app_model: App, account: Account, args: dict) -> App:
def convert_to_workflow(self, app_model: App, account: Account, args: dict[str, Any]) -> App:
"""
Basic mode of chatbot app(expert mode) to workflow
Completion App to Workflow App
@@ -1421,7 +1421,7 @@ class WorkflowService:
if node_type == BuiltinNodeTypes.HUMAN_INPUT:
self._validate_human_input_node_data(node_data)
def validate_features_structure(self, app_model: App, features: dict):
def validate_features_structure(self, app_model: App, features: dict[str, Any]):
match app_model.mode:
case AppMode.ADVANCED_CHAT:
return AdvancedChatAppConfigManager.config_validate(
@@ -1434,7 +1434,7 @@ class WorkflowService:
case _:
raise ValueError(f"Invalid app mode: {app_model.mode}")
def _validate_human_input_node_data(self, node_data: dict) -> None:
def _validate_human_input_node_data(self, node_data: dict[str, Any]) -> None:
"""
Validate HumanInput node data format.
@@ -1452,7 +1452,7 @@ class WorkflowService:
raise ValueError(f"Invalid HumanInput node data: {str(e)}")
def update_workflow(
self, *, session: Session, workflow_id: str, tenant_id: str, account_id: str, data: dict
self, *, session: Session, workflow_id: str, tenant_id: str, account_id: str, data: dict[str, Any]
) -> Workflow | None:
"""
Update workflow attributes
@@ -1,47 +0,0 @@
import uuid
from unittest import mock
from controllers.console.app import workflow_draft_variable as draft_variable_api
from controllers.console.app import wraps
from factories.variable_factory import build_segment
from models import App, AppMode
from models.workflow import WorkflowDraftVariable
from services.workflow_draft_variable_service import WorkflowDraftVariableList, WorkflowDraftVariableService
def _get_mock_srv_class() -> type[WorkflowDraftVariableService]:
return mock.create_autospec(WorkflowDraftVariableService)
class TestWorkflowDraftNodeVariableListApi:
def test_get(self, test_client, auth_header, monkeypatch):
srv_class = _get_mock_srv_class()
mock_app_model: App = App()
mock_app_model.id = str(uuid.uuid4())
test_node_id = "test_node_id"
mock_app_model.mode = AppMode.ADVANCED_CHAT
mock_load_app_model = mock.Mock(return_value=mock_app_model)
monkeypatch.setattr(draft_variable_api, "WorkflowDraftVariableService", srv_class)
monkeypatch.setattr(wraps, "_load_app_model", mock_load_app_model)
var1 = WorkflowDraftVariable.new_node_variable(
app_id="test_app_1",
node_id="test_node_1",
name="str_var",
value=build_segment("str_value"),
node_execution_id=str(uuid.uuid4()),
)
srv_instance = mock.create_autospec(WorkflowDraftVariableService, instance=True)
srv_class.return_value = srv_instance
srv_instance.list_node_variables.return_value = WorkflowDraftVariableList(variables=[var1])
response = test_client.get(
f"/console/api/apps/{mock_app_model.id}/workflows/draft/nodes/{test_node_id}/variables",
headers=auth_header,
)
assert response.status_code == 200
response_dict = response.json
assert isinstance(response_dict, dict)
assert "items" in response_dict
assert len(response_dict["items"]) == 1
@@ -1,244 +0,0 @@
"""Integration tests for Trigger Provider subscription permission verification."""
import uuid
from unittest import mock
import pytest
from flask.testing import FlaskClient
from controllers.console.workspace import trigger_providers as trigger_providers_api
from libs.datetime_utils import naive_utc_now
from models import Tenant
from models.account import Account, TenantAccountJoin, TenantAccountRole
class TestTriggerProviderSubscriptionPermissions:
"""Test permission verification for Trigger Provider subscription endpoints."""
@pytest.fixture
def mock_account(self, monkeypatch: pytest.MonkeyPatch):
"""Create a mock Account for testing."""
account = Account(name="Test User", email="test@example.com")
account.id = str(uuid.uuid4())
account.last_active_at = naive_utc_now()
account.created_at = naive_utc_now()
account.updated_at = naive_utc_now()
# Create mock tenant
tenant = Tenant(name="Test Tenant")
tenant.id = str(uuid.uuid4())
mock_session_instance = mock.Mock()
mock_tenant_join = TenantAccountJoin(role=TenantAccountRole.OWNER)
monkeypatch.setattr(mock_session_instance, "scalar", mock.Mock(return_value=mock_tenant_join))
mock_scalars_result = mock.Mock()
mock_scalars_result.one.return_value = tenant
monkeypatch.setattr(mock_session_instance, "scalars", mock.Mock(return_value=mock_scalars_result))
mock_session_context = mock.Mock()
mock_session_context.__enter__.return_value = mock_session_instance
monkeypatch.setattr("models.account.Session", lambda _, expire_on_commit: mock_session_context)
account.current_tenant = tenant
account.current_tenant_id = tenant.id
return account
@pytest.mark.parametrize(
("role", "list_status", "get_status", "update_status", "create_status", "build_status", "delete_status"),
[
# Admin/Owner can do everything
(TenantAccountRole.OWNER, 200, 200, 200, 200, 200, 200),
(TenantAccountRole.ADMIN, 200, 200, 200, 200, 200, 200),
# Editor can list, get, update (parameters), but not create, build, or delete
(TenantAccountRole.EDITOR, 200, 200, 200, 403, 403, 403),
# Normal user cannot do anything
(TenantAccountRole.NORMAL, 403, 403, 403, 403, 403, 403),
# Dataset operator cannot do anything
(TenantAccountRole.DATASET_OPERATOR, 403, 403, 403, 403, 403, 403),
],
)
def test_trigger_subscription_permissions(
self,
test_client: FlaskClient,
auth_header,
monkeypatch,
mock_account,
role: TenantAccountRole,
list_status: int,
get_status: int,
update_status: int,
create_status: int,
build_status: int,
delete_status: int,
):
"""Test that different roles have appropriate permissions for trigger subscription operations."""
# Set user role
mock_account.role = role
# Mock current user
monkeypatch.setattr(trigger_providers_api, "current_user", mock_account)
# Mock AccountService.load_user to prevent authentication issues
from services.account_service import AccountService
mock_load_user = mock.Mock(return_value=mock_account)
monkeypatch.setattr(AccountService, "load_user", mock_load_user)
# Test data
provider = "some_provider/some_trigger"
subscription_builder_id = str(uuid.uuid4())
subscription_id = str(uuid.uuid4())
# Mock service methods
mock_list_subscriptions = mock.Mock(return_value=[])
monkeypatch.setattr(
"services.trigger.trigger_provider_service.TriggerProviderService.list_trigger_provider_subscriptions",
mock_list_subscriptions,
)
mock_get_subscription_builder = mock.Mock(return_value={"id": subscription_builder_id})
monkeypatch.setattr(
"services.trigger.trigger_subscription_builder_service.TriggerSubscriptionBuilderService.get_subscription_builder_by_id",
mock_get_subscription_builder,
)
mock_update_subscription_builder = mock.Mock(return_value={"id": subscription_builder_id})
monkeypatch.setattr(
"services.trigger.trigger_subscription_builder_service.TriggerSubscriptionBuilderService.update_trigger_subscription_builder",
mock_update_subscription_builder,
)
mock_create_subscription_builder = mock.Mock(return_value={"id": subscription_builder_id})
monkeypatch.setattr(
"services.trigger.trigger_subscription_builder_service.TriggerSubscriptionBuilderService.create_trigger_subscription_builder",
mock_create_subscription_builder,
)
mock_update_and_build_builder = mock.Mock()
monkeypatch.setattr(
"services.trigger.trigger_subscription_builder_service.TriggerSubscriptionBuilderService.update_and_build_builder",
mock_update_and_build_builder,
)
mock_delete_provider = mock.Mock()
mock_delete_plugin_trigger = mock.Mock()
mock_db_session = mock.Mock()
mock_db_session.commit = mock.Mock()
def mock_session_func(engine=None):
return mock_session_context
mock_session_context = mock.Mock()
mock_session_context.__enter__.return_value = mock_db_session
mock_session_context.__exit__.return_value = None
monkeypatch.setattr("services.trigger.trigger_provider_service.Session", mock_session_func)
monkeypatch.setattr("services.trigger.trigger_subscription_operator_service.Session", mock_session_func)
monkeypatch.setattr(
"services.trigger.trigger_provider_service.TriggerProviderService.delete_trigger_provider",
mock_delete_provider,
)
monkeypatch.setattr(
"services.trigger.trigger_subscription_operator_service.TriggerSubscriptionOperatorService.delete_plugin_trigger_by_subscription",
mock_delete_plugin_trigger,
)
# Test 1: List subscriptions (should work for Editor, Admin, Owner)
response = test_client.get(
f"/console/api/workspaces/current/trigger-provider/{provider}/subscriptions/list",
headers=auth_header,
)
assert response.status_code == list_status
# Test 2: Get subscription builder (should work for Editor, Admin, Owner)
response = test_client.get(
f"/console/api/workspaces/current/trigger-provider/{provider}/subscriptions/builder/{subscription_builder_id}",
headers=auth_header,
)
assert response.status_code == get_status
# Test 3: Update subscription builder parameters (should work for Editor, Admin, Owner)
response = test_client.post(
f"/console/api/workspaces/current/trigger-provider/{provider}/subscriptions/builder/update/{subscription_builder_id}",
headers=auth_header,
json={"parameters": {"webhook_url": "https://example.com/webhook"}},
)
assert response.status_code == update_status
# Test 4: Create subscription builder (should only work for Admin, Owner)
response = test_client.post(
f"/console/api/workspaces/current/trigger-provider/{provider}/subscriptions/builder/create",
headers=auth_header,
json={"credential_type": "api_key"},
)
assert response.status_code == create_status
# Test 5: Build/activate subscription (should only work for Admin, Owner)
response = test_client.post(
f"/console/api/workspaces/current/trigger-provider/{provider}/subscriptions/builder/build/{subscription_builder_id}",
headers=auth_header,
json={"name": "Test Subscription"},
)
assert response.status_code == build_status
# Test 6: Delete subscription (should only work for Admin, Owner)
response = test_client.post(
f"/console/api/workspaces/current/trigger-provider/{subscription_id}/subscriptions/delete",
headers=auth_header,
)
assert response.status_code == delete_status
@pytest.mark.parametrize(
("role", "status"),
[
(TenantAccountRole.OWNER, 200),
(TenantAccountRole.ADMIN, 200),
# Editor should be able to access logs for debugging
(TenantAccountRole.EDITOR, 200),
(TenantAccountRole.NORMAL, 403),
(TenantAccountRole.DATASET_OPERATOR, 403),
],
)
def test_trigger_subscription_logs_permissions(
self,
test_client: FlaskClient,
auth_header,
monkeypatch,
mock_account,
role: TenantAccountRole,
status: int,
):
"""Test that different roles have appropriate permissions for accessing subscription logs."""
# Set user role
mock_account.role = role
# Mock current user
monkeypatch.setattr(trigger_providers_api, "current_user", mock_account)
# Mock AccountService.load_user to prevent authentication issues
from services.account_service import AccountService
mock_load_user = mock.Mock(return_value=mock_account)
monkeypatch.setattr(AccountService, "load_user", mock_load_user)
# Test data
provider = "some_provider/some_trigger"
subscription_builder_id = str(uuid.uuid4())
# Mock service method
mock_list_logs = mock.Mock(return_value=[])
monkeypatch.setattr(
"services.trigger.trigger_subscription_builder_service.TriggerSubscriptionBuilderService.list_logs",
mock_list_logs,
)
# Test access to logs
response = test_client.get(
f"/console/api/workspaces/current/trigger-provider/{provider}/subscriptions/builder/logs/{subscription_builder_id}",
headers=auth_header,
)
assert response.status_code == status
@@ -1,375 +0,0 @@
import unittest
from datetime import UTC, datetime
from unittest.mock import patch
from uuid import uuid4
import pytest
from graphon.file import File, FileTransferMethod, FileType
from sqlalchemy.orm import Session
from core.app.file_access import DatabaseFileAccessController
from extensions.ext_database import db
from extensions.storage.storage_type import StorageType
from factories.file_factory import StorageKeyLoader
from models import ToolFile, UploadFile
from models.enums import CreatorUserRole
@pytest.mark.usefixtures("flask_req_ctx")
class TestStorageKeyLoader(unittest.TestCase):
"""
Integration tests for StorageKeyLoader class.
Tests the batched loading of storage keys from the database for files
with different transfer methods: LOCAL_FILE, REMOTE_URL, and TOOL_FILE.
"""
def setUp(self):
"""Set up test data before each test method."""
self.session = db.session()
self.tenant_id = str(uuid4())
self.user_id = str(uuid4())
self.conversation_id = str(uuid4())
# Create test data that will be cleaned up after each test
self.test_upload_files = []
self.test_tool_files = []
# Create StorageKeyLoader instance
self.loader = StorageKeyLoader(
self.session,
self.tenant_id,
access_controller=DatabaseFileAccessController(),
)
def tearDown(self):
"""Clean up test data after each test method."""
self.session.rollback()
def _create_upload_file(
self, file_id: str | None = None, storage_key: str | None = None, tenant_id: str | None = None
) -> UploadFile:
"""Helper method to create an UploadFile record for testing."""
if file_id is None:
file_id = str(uuid4())
if storage_key is None:
storage_key = f"test_storage_key_{uuid4()}"
if tenant_id is None:
tenant_id = self.tenant_id
upload_file = UploadFile(
tenant_id=tenant_id,
storage_type=StorageType.LOCAL,
key=storage_key,
name="test_file.txt",
size=1024,
extension=".txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=self.user_id,
created_at=datetime.now(UTC),
used=False,
)
upload_file.id = file_id
self.session.add(upload_file)
self.session.flush()
self.test_upload_files.append(upload_file)
return upload_file
def _create_tool_file(
self, file_id: str | None = None, file_key: str | None = None, tenant_id: str | None = None
) -> ToolFile:
"""Helper method to create a ToolFile record for testing."""
if file_id is None:
file_id = str(uuid4())
if file_key is None:
file_key = f"test_file_key_{uuid4()}"
if tenant_id is None:
tenant_id = self.tenant_id
tool_file = ToolFile(
user_id=self.user_id,
tenant_id=tenant_id,
conversation_id=self.conversation_id,
file_key=file_key,
mimetype="text/plain",
original_url="http://example.com/file.txt",
name="test_tool_file.txt",
size=2048,
)
tool_file.id = file_id
self.session.add(tool_file)
self.session.flush()
self.test_tool_files.append(tool_file)
return tool_file
def _create_file(self, related_id: str, transfer_method: FileTransferMethod, tenant_id: str | None = None) -> File:
"""Helper method to create a File object for testing."""
if tenant_id is None:
tenant_id = self.tenant_id
# Set related_id for LOCAL_FILE and TOOL_FILE transfer methods
file_related_id = None
remote_url = None
if transfer_method in (FileTransferMethod.LOCAL_FILE, FileTransferMethod.TOOL_FILE):
file_related_id = related_id
elif transfer_method == FileTransferMethod.REMOTE_URL:
remote_url = "https://example.com/test_file.txt"
file_related_id = related_id
return File(
id=str(uuid4()), # Generate new UUID for File.id
tenant_id=tenant_id,
type=FileType.DOCUMENT,
transfer_method=transfer_method,
related_id=file_related_id,
remote_url=remote_url,
filename="test_file.txt",
extension=".txt",
mime_type="text/plain",
size=1024,
storage_key="initial_key",
)
def test_load_storage_keys_local_file(self):
"""Test loading storage keys for LOCAL_FILE transfer method."""
# Create test data
upload_file = self._create_upload_file()
file = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
# Load storage keys
self.loader.load_storage_keys([file])
# Verify storage key was loaded correctly
assert file._storage_key == upload_file.key
def test_load_storage_keys_remote_url(self):
"""Test loading storage keys for REMOTE_URL transfer method."""
# Create test data
upload_file = self._create_upload_file()
file = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.REMOTE_URL)
# Load storage keys
self.loader.load_storage_keys([file])
# Verify storage key was loaded correctly
assert file._storage_key == upload_file.key
def test_load_storage_keys_tool_file(self):
"""Test loading storage keys for TOOL_FILE transfer method."""
# Create test data
tool_file = self._create_tool_file()
file = self._create_file(related_id=tool_file.id, transfer_method=FileTransferMethod.TOOL_FILE)
# Load storage keys
self.loader.load_storage_keys([file])
# Verify storage key was loaded correctly
assert file._storage_key == tool_file.file_key
def test_load_storage_keys_mixed_methods(self):
"""Test batch loading with mixed transfer methods."""
# Create test data for different transfer methods
upload_file1 = self._create_upload_file()
upload_file2 = self._create_upload_file()
tool_file = self._create_tool_file()
file1 = self._create_file(related_id=upload_file1.id, transfer_method=FileTransferMethod.LOCAL_FILE)
file2 = self._create_file(related_id=upload_file2.id, transfer_method=FileTransferMethod.REMOTE_URL)
file3 = self._create_file(related_id=tool_file.id, transfer_method=FileTransferMethod.TOOL_FILE)
files = [file1, file2, file3]
# Load storage keys
self.loader.load_storage_keys(files)
# Verify all storage keys were loaded correctly
assert file1._storage_key == upload_file1.key
assert file2._storage_key == upload_file2.key
assert file3._storage_key == tool_file.file_key
def test_load_storage_keys_empty_list(self):
"""Test with empty file list."""
# Should not raise any exceptions
self.loader.load_storage_keys([])
def test_load_storage_keys_ignores_legacy_file_tenant_id(self):
"""Legacy file tenant_id should not override the loader tenant scope."""
upload_file = self._create_upload_file()
file = self._create_file(
related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE, tenant_id=str(uuid4())
)
self.loader.load_storage_keys([file])
assert file._storage_key == upload_file.key
def test_load_storage_keys_missing_file_id(self):
"""Test with None file.related_id."""
# Create a file with valid parameters first, then manually set related_id to None
file = self._create_file(related_id=str(uuid4()), transfer_method=FileTransferMethod.LOCAL_FILE)
file.related_id = None
# Should raise ValueError for None file related_id
with pytest.raises(ValueError) as context:
self.loader.load_storage_keys([file])
assert str(context.value) == "file id should not be None."
def test_load_storage_keys_nonexistent_upload_file_records(self):
"""Test with missing UploadFile database records."""
# Create file with non-existent upload file id
non_existent_id = str(uuid4())
file = self._create_file(related_id=non_existent_id, transfer_method=FileTransferMethod.LOCAL_FILE)
# Should raise ValueError for missing record
with pytest.raises(ValueError):
self.loader.load_storage_keys([file])
def test_load_storage_keys_nonexistent_tool_file_records(self):
"""Test with missing ToolFile database records."""
# Create file with non-existent tool file id
non_existent_id = str(uuid4())
file = self._create_file(related_id=non_existent_id, transfer_method=FileTransferMethod.TOOL_FILE)
# Should raise ValueError for missing record
with pytest.raises(ValueError):
self.loader.load_storage_keys([file])
def test_load_storage_keys_invalid_uuid(self):
"""Test with invalid UUID format."""
# Create a file with valid parameters first, then manually set invalid related_id
file = self._create_file(related_id=str(uuid4()), transfer_method=FileTransferMethod.LOCAL_FILE)
file.related_id = "invalid-uuid-format"
# Should raise ValueError for invalid UUID
with pytest.raises(ValueError):
self.loader.load_storage_keys([file])
def test_load_storage_keys_batch_efficiency(self):
"""Test batched operations use efficient queries."""
# Create multiple files of different types
upload_files = [self._create_upload_file() for _ in range(3)]
tool_files = [self._create_tool_file() for _ in range(2)]
files = []
files.extend(
[self._create_file(related_id=uf.id, transfer_method=FileTransferMethod.LOCAL_FILE) for uf in upload_files]
)
files.extend(
[self._create_file(related_id=tf.id, transfer_method=FileTransferMethod.TOOL_FILE) for tf in tool_files]
)
# Mock the session to count queries
with patch.object(self.session, "scalars", wraps=self.session.scalars) as mock_scalars:
self.loader.load_storage_keys(files)
# Should make exactly 2 queries (one for upload_files, one for tool_files)
assert mock_scalars.call_count == 2
# Verify all storage keys were loaded correctly
for i, file in enumerate(files[:3]):
assert file._storage_key == upload_files[i].key
for i, file in enumerate(files[3:]):
assert file._storage_key == tool_files[i].file_key
def test_load_storage_keys_tenant_isolation(self):
"""Test that tenant isolation works correctly."""
# Create files for different tenants
other_tenant_id = str(uuid4())
# Create upload file for current tenant
upload_file_current = self._create_upload_file()
file_current = self._create_file(
related_id=upload_file_current.id, transfer_method=FileTransferMethod.LOCAL_FILE
)
# Create upload file for other tenant (but don't add to cleanup list)
upload_file_other = UploadFile(
tenant_id=other_tenant_id,
storage_type=StorageType.LOCAL,
key="other_tenant_key",
name="other_file.txt",
size=1024,
extension=".txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=self.user_id,
created_at=datetime.now(UTC),
used=False,
)
upload_file_other.id = str(uuid4())
self.session.add(upload_file_other)
self.session.flush()
# Create file for other tenant but try to load with current tenant's loader
file_other = self._create_file(
related_id=upload_file_other.id, transfer_method=FileTransferMethod.LOCAL_FILE, tenant_id=other_tenant_id
)
# Should raise ValueError due to tenant mismatch
with pytest.raises(ValueError) as context:
self.loader.load_storage_keys([file_other])
assert "Upload file not found for id:" in str(context.value)
# Current tenant's file should still work
self.loader.load_storage_keys([file_current])
assert file_current._storage_key == upload_file_current.key
def test_load_storage_keys_mixed_tenant_batch(self):
"""Test batch with mixed tenant files (should fail on first mismatch)."""
# Create files for current tenant
upload_file_current = self._create_upload_file()
file_current = self._create_file(
related_id=upload_file_current.id, transfer_method=FileTransferMethod.LOCAL_FILE
)
# Create file for different tenant
other_tenant_id = str(uuid4())
file_other = self._create_file(
related_id=str(uuid4()), transfer_method=FileTransferMethod.LOCAL_FILE, tenant_id=other_tenant_id
)
# Should raise ValueError on tenant mismatch
with pytest.raises(ValueError) as context:
self.loader.load_storage_keys([file_current, file_other])
assert "Upload file not found for id:" in str(context.value)
def test_load_storage_keys_duplicate_file_ids(self):
"""Test handling of duplicate file IDs in the batch."""
# Create upload file
upload_file = self._create_upload_file()
# Create two File objects with same related_id
file1 = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
file2 = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
# Should handle duplicates gracefully
self.loader.load_storage_keys([file1, file2])
# Both files should have the same storage key
assert file1._storage_key == upload_file.key
assert file2._storage_key == upload_file.key
def test_load_storage_keys_session_isolation(self):
"""Test that the loader uses the provided session correctly."""
# Create test data
upload_file = self._create_upload_file()
file = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
# Create loader with different session (same underlying connection)
with Session(bind=db.engine) as other_session:
other_loader = StorageKeyLoader(
other_session,
self.tenant_id,
access_controller=DatabaseFileAccessController(),
)
with pytest.raises(ValueError):
other_loader.load_storage_keys([file])
@@ -1,11 +0,0 @@
import pytest
from core.helper.code_executor.code_executor import CodeExecutionError, CodeExecutor
CODE_LANGUAGE = "unsupported_language"
def test_unsupported_with_code_template():
with pytest.raises(CodeExecutionError) as e:
CodeExecutor.execute_workflow_code_template(language=CODE_LANGUAGE, code="", inputs={})
assert str(e.value) == f"Unsupported language {CODE_LANGUAGE}"
@@ -1,95 +0,0 @@
import base64
from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
from core.helper.code_executor.jinja2.jinja2_transformer import Jinja2TemplateTransformer
CODE_LANGUAGE = CodeLanguage.JINJA2
def test_jinja2():
"""Test basic Jinja2 template rendering."""
template = "Hello {{template}}"
# Template must be base64 encoded to match the new safe embedding approach
template_b64 = base64.b64encode(template.encode("utf-8")).decode("utf-8")
inputs = base64.b64encode(b'{"template": "World"}').decode("utf-8")
code = (
Jinja2TemplateTransformer.get_runner_script()
.replace(Jinja2TemplateTransformer._template_b64_placeholder, template_b64)
.replace(Jinja2TemplateTransformer._inputs_placeholder, inputs)
)
result = CodeExecutor.execute_code(
language=CODE_LANGUAGE, preload=Jinja2TemplateTransformer.get_preload_script(), code=code
)
assert result == "<<RESULT>>Hello World<<RESULT>>\n"
def test_jinja2_with_code_template():
"""Test template rendering via the high-level workflow API."""
result = CodeExecutor.execute_workflow_code_template(
language=CODE_LANGUAGE, code="Hello {{template}}", inputs={"template": "World"}
)
assert result == {"result": "Hello World"}
def test_jinja2_get_runner_script():
"""Test that runner script contains required placeholders."""
runner_script = Jinja2TemplateTransformer.get_runner_script()
assert runner_script.count(Jinja2TemplateTransformer._template_b64_placeholder) == 1
assert runner_script.count(Jinja2TemplateTransformer._inputs_placeholder) == 1
assert runner_script.count(Jinja2TemplateTransformer._result_tag) == 2
def test_jinja2_template_with_special_characters():
"""
Test that templates with special characters (quotes, newlines) render correctly.
This is a regression test for issue #26818 where textarea pre-fill values
containing special characters would break template rendering.
"""
# Template with triple quotes, single quotes, double quotes, and newlines
template = """<html>
<body>
<input value="{{ task.get('Task ID', '') }}"/>
<textarea>{{ task.get('Issues', 'No issues reported') }}</textarea>
<p>Status: "{{ status }}"</p>
<pre>'''code block'''</pre>
</body>
</html>"""
inputs = {"task": {"Task ID": "TASK-123", "Issues": "Line 1\nLine 2\nLine 3"}, "status": "completed"}
result = CodeExecutor.execute_workflow_code_template(language=CODE_LANGUAGE, code=template, inputs=inputs)
# Verify the template rendered correctly with all special characters
output = result["result"]
assert 'value="TASK-123"' in output
assert "<textarea>Line 1\nLine 2\nLine 3</textarea>" in output
assert 'Status: "completed"' in output
assert "'''code block'''" in output
def test_jinja2_template_with_html_textarea_prefill():
"""
Specific test for HTML textarea with Jinja2 variable pre-fill.
Verifies fix for issue #26818.
"""
template = "<textarea name='notes'>{{ notes }}</textarea>"
notes_content = "This is a multi-line note.\nWith special chars: 'single' and \"double\" quotes."
inputs = {"notes": notes_content}
result = CodeExecutor.execute_workflow_code_template(language=CODE_LANGUAGE, code=template, inputs=inputs)
expected_output = f"<textarea name='notes'>{notes_content}</textarea>"
assert result["result"] == expected_output
def test_jinja2_assemble_runner_script_encodes_template():
"""Test that assemble_runner_script properly base64 encodes the template."""
template = "Hello {{ name }}!"
inputs = {"name": "World"}
script = Jinja2TemplateTransformer.assemble_runner_script(template, inputs)
# The template should be base64 encoded in the script
template_b64 = base64.b64encode(template.encode("utf-8")).decode("utf-8")
assert template_b64 in script
# The raw template should NOT appear in the script (it's encoded)
assert "Hello {{ name }}!" not in script
@@ -1,36 +0,0 @@
from textwrap import dedent
from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
from core.helper.code_executor.python3.python3_code_provider import Python3CodeProvider
from core.helper.code_executor.python3.python3_transformer import Python3TemplateTransformer
CODE_LANGUAGE = CodeLanguage.PYTHON3
def test_python3_plain():
code = 'print("Hello World")'
result = CodeExecutor.execute_code(language=CODE_LANGUAGE, preload="", code=code)
assert result == "Hello World\n"
def test_python3_json():
code = dedent("""
import json
print(json.dumps({'Hello': 'World'}))
""")
result = CodeExecutor.execute_code(language=CODE_LANGUAGE, preload="", code=code)
assert result == '{"Hello": "World"}\n'
def test_python3_with_code_template():
result = CodeExecutor.execute_workflow_code_template(
language=CODE_LANGUAGE, code=Python3CodeProvider.get_default_code(), inputs={"arg1": "Hello", "arg2": "World"}
)
assert result == {"result": "HelloWorld"}
def test_python3_get_runner_script():
runner_script = Python3TemplateTransformer.get_runner_script()
assert runner_script.count(Python3TemplateTransformer._code_placeholder) == 1
assert runner_script.count(Python3TemplateTransformer._inputs_placeholder) == 1
assert runner_script.count(Python3TemplateTransformer._result_tag) == 2
@@ -432,7 +432,7 @@ class TestWorkflowAppLogEndpoints:
monkeypatch.setattr(workflow_app_log_module, "sessionmaker", DummySessionMaker)
def fake_get_paginate(self, **_kwargs):
return {"items": [], "total": 0}
return {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []}
monkeypatch.setattr(
workflow_app_log_module.WorkflowAppService,
@@ -443,7 +443,7 @@ class TestWorkflowAppLogEndpoints:
with app.test_request_context("/?page=1&limit=20"):
result = method(app_model=SimpleNamespace(id="app-1"))
assert result == {"items": [], "total": 0}
assert result == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []}
class TestWorkflowDraftVariableEndpoints:
@@ -608,7 +608,8 @@ class TestWorkflowTriggerEndpoints:
with app.test_request_context("/?node_id=node-1"):
result = method(app_model=SimpleNamespace(id="app-1"))
assert result is trigger
assert isinstance(result, dict)
assert {"id", "webhook_id", "webhook_url", "webhook_debug_url", "node_id", "created_at"} <= set(result.keys())
class TestWrapsEndpoints:
@@ -96,6 +96,56 @@ class TestAppImportApi:
assert status == 200
assert response["status"] == ImportStatus.COMPLETED
def test_import_post_commits_session_on_success(self, app, monkeypatch: pytest.MonkeyPatch) -> None:
api = app_import_module.AppImportApi()
method = _unwrap(api.post)
_install_features(monkeypatch, enabled=False)
monkeypatch.setattr(
app_import_module.AppDslService,
"import_app",
lambda *_args, **_kwargs: _Result(ImportStatus.COMPLETED, app_id="app-123"),
)
monkeypatch.setattr(app_import_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
fake_session = MagicMock()
fake_session.__enter__.return_value = fake_session
fake_session.__exit__.return_value = None
monkeypatch.setattr(app_import_module, "Session", lambda *_args, **_kwargs: fake_session)
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
response, status = method()
fake_session.commit.assert_called_once_with()
fake_session.rollback.assert_not_called()
assert status == 200
assert response["status"] == ImportStatus.COMPLETED
def test_import_post_rolls_back_session_on_failure(self, app, monkeypatch: pytest.MonkeyPatch) -> None:
api = app_import_module.AppImportApi()
method = _unwrap(api.post)
_install_features(monkeypatch, enabled=False)
monkeypatch.setattr(
app_import_module.AppDslService,
"import_app",
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED, app_id=None),
)
monkeypatch.setattr(app_import_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
fake_session = MagicMock()
fake_session.__enter__.return_value = fake_session
fake_session.__exit__.return_value = None
monkeypatch.setattr(app_import_module, "Session", lambda *_args, **_kwargs: fake_session)
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
response, status = method()
fake_session.rollback.assert_called_once_with()
fake_session.commit.assert_not_called()
assert status == 400
assert response["status"] == ImportStatus.FAILED
class TestAppImportConfirmApi:
@pytest.fixture
@@ -12,7 +12,7 @@ from constants import HEADER_NAME_CSRF_TOKEN
from libs.datetime_utils import naive_utc_now
from libs.token import _real_cookie_name, generate_csrf_token
from models import Account, DifySetup, Tenant, TenantAccountJoin
from models.account import AccountStatus, TenantAccountRole
from models.account import AccountStatus, TenantAccountRole, TenantStatus
from models.enums import ConversationFromSource, CreatorUserRole
from models.model import App, AppMode, Conversation, Message
from models.workflow import WorkflowRun
@@ -30,7 +30,7 @@ def _create_account_and_tenant(db_session: Session) -> tuple[Account, Tenant]:
db_session.add(account)
db_session.commit()
tenant = Tenant(name="Test Tenant", status="normal")
tenant = Tenant(name="Test Tenant", status=TenantStatus.NORMAL)
db_session.add(tenant)
db_session.commit()
@@ -11,7 +11,7 @@ from constants import HEADER_NAME_CSRF_TOKEN
from libs.datetime_utils import naive_utc_now
from libs.token import _real_cookie_name, generate_csrf_token
from models import Account, DifySetup, Tenant, TenantAccountJoin
from models.account import AccountStatus, TenantAccountRole
from models.account import AccountStatus, TenantAccountRole, TenantStatus
from models.model import App, AppMode
from services.account_service import AccountService
@@ -37,7 +37,7 @@ def create_console_account_and_tenant(db_session: Session) -> tuple[Account, Ten
db_session.add(account)
db_session.commit()
tenant = Tenant(name="Test Tenant", status="normal")
tenant = Tenant(name="Test Tenant", status=TenantStatus.NORMAL)
db_session.add(tenant)
db_session.commit()
@@ -0,0 +1,110 @@
"""
Testcontainers integration tests for Service API Site controller.
"""
from __future__ import annotations
import pytest
from flask import Flask
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
from controllers.service_api.app.site import AppSiteApi
from models.account import Tenant, TenantStatus
from models.model import App, AppMode, Site
@pytest.fixture
def app(flask_app_with_containers) -> Flask:
return flask_app_with_containers
def _unwrap(method):
fn = method
while hasattr(fn, "__wrapped__"):
fn = fn.__wrapped__
return fn
def _create_tenant(db_session: Session, *, status: TenantStatus = TenantStatus.NORMAL) -> Tenant:
tenant = Tenant(name="service-api-site-tenant", status=status)
db_session.add(tenant)
db_session.commit()
return tenant
def _create_app(db_session: Session, tenant_id: str) -> App:
app_model = App(
tenant_id=tenant_id,
mode=AppMode.CHAT,
name="service-api-site-app",
enable_site=True,
enable_api=True,
status="normal",
)
db_session.add(app_model)
db_session.commit()
return app_model
def _create_site(db_session: Session, app_id: str) -> Site:
site = Site(
app_id=app_id,
title="Service API Site",
icon_type="emoji",
icon="robot",
icon_background="#ffffff",
description="Service API test site",
default_language="en-US",
prompt_public=True,
show_workflow_steps=True,
customize_token_strategy="not_allow",
use_icon_as_answer_icon=False,
chat_color_theme="light",
chat_color_theme_inverted=False,
)
db_session.add(site)
db_session.commit()
return site
class TestAppSiteApi:
def test_get_site_success(self, app: Flask, db_session_with_containers: Session) -> None:
tenant = _create_tenant(db_session_with_containers)
app_model = _create_app(db_session_with_containers, tenant.id)
_create_site(db_session_with_containers, app_model.id)
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test-token"}):
api = AppSiteApi()
response = _unwrap(api.get)(api, app_model=app_model)
assert response["title"] == "Service API Site"
assert response["icon"] == "robot"
assert response["description"] == "Service API test site"
def test_get_site_not_found(self, app: Flask, db_session_with_containers: Session) -> None:
tenant = _create_tenant(db_session_with_containers)
app_model = _create_app(db_session_with_containers, tenant.id)
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test-token"}):
api = AppSiteApi()
with pytest.raises(Forbidden):
_unwrap(api.get)(api, app_model=app_model)
def test_get_site_tenant_archived(self, app: Flask, db_session_with_containers: Session) -> None:
tenant = _create_tenant(db_session_with_containers)
app_model = _create_app(db_session_with_containers, tenant.id)
_create_site(db_session_with_containers, app_model.id)
archived_tenant = db_session_with_containers.get(Tenant, tenant.id)
assert archived_tenant is not None
archived_tenant.status = TenantStatus.ARCHIVE
db_session_with_containers.commit()
app_model = db_session_with_containers.get(App, app_model.id)
assert app_model is not None
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test-token"}):
api = AppSiteApi()
with pytest.raises(Forbidden):
_unwrap(api.get)(api, app_model=app_model)
@@ -1,28 +1,48 @@
"""Unit tests for controllers.web.site endpoints."""
"""Testcontainers integration tests for controllers.web.site endpoints."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from flask import Flask
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
from controllers.web.site import AppSiteApi, AppSiteInfo
from models import Tenant, TenantStatus
from models.model import App, AppMode, CustomizeTokenStrategy, Site
def _tenant(*, status: str = "normal") -> SimpleNamespace:
return SimpleNamespace(
id="tenant-1",
status=status,
plan="basic",
custom_config_dict={"remove_webapp_brand": False, "replace_webapp_logo": False},
@pytest.fixture
def app(flask_app_with_containers) -> Flask:
return flask_app_with_containers
def _create_tenant(db_session: Session, *, status: TenantStatus = TenantStatus.NORMAL) -> Tenant:
tenant = Tenant(name="test-tenant", status=status)
db_session.add(tenant)
db_session.commit()
return tenant
def _create_app(db_session: Session, tenant_id: str, *, enable_site: bool = True) -> App:
app_model = App(
tenant_id=tenant_id,
mode=AppMode.CHAT,
name="test-app",
enable_site=enable_site,
enable_api=True,
)
db_session.add(app_model)
db_session.commit()
return app_model
def _site() -> SimpleNamespace:
return SimpleNamespace(
def _create_site(db_session: Session, app_id: str) -> Site:
site = Site(
app_id=app_id,
title="Site",
icon_type="emoji",
icon="robot",
@@ -31,77 +51,64 @@ def _site() -> SimpleNamespace:
default_language="en",
chat_color_theme="light",
chat_color_theme_inverted=False,
copyright=None,
privacy_policy=None,
custom_disclaimer=None,
customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW,
code=f"code-{app_id[-6:]}",
prompt_public=False,
show_workflow_steps=True,
use_icon_as_answer_icon=False,
)
db_session.add(site)
db_session.commit()
return site
# ---------------------------------------------------------------------------
# AppSiteApi
# ---------------------------------------------------------------------------
class TestAppSiteApi:
@patch("controllers.web.site.FeatureService.get_features")
@patch("controllers.web.site.db")
def test_happy_path(self, mock_db: MagicMock, mock_features: MagicMock, app: Flask) -> None:
def test_happy_path(self, mock_features, app: Flask, db_session_with_containers: Session) -> None:
app.config["RESTX_MASK_HEADER"] = "X-Fields"
mock_features.return_value = SimpleNamespace(can_replace_logo=False)
site_obj = _site()
mock_db.session.scalar.return_value = site_obj
tenant = _tenant()
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant, enable_site=True)
tenant = _create_tenant(db_session_with_containers)
app_model = _create_app(db_session_with_containers, tenant.id)
_create_site(db_session_with_containers, app_model.id)
end_user = SimpleNamespace(id="eu-1")
mock_features.return_value = SimpleNamespace(can_replace_logo=False)
with app.test_request_context("/site"):
result = AppSiteApi().get(app_model, end_user)
# marshal_with serializes AppSiteInfo to a dict
assert result["app_id"] == "app-1"
assert result["app_id"] == app_model.id
assert result["plan"] == "basic"
assert result["enable_site"] is True
@patch("controllers.web.site.db")
def test_missing_site_raises_forbidden(self, mock_db: MagicMock, app: Flask) -> None:
def test_missing_site_raises_forbidden(self, app: Flask, db_session_with_containers: Session) -> None:
app.config["RESTX_MASK_HEADER"] = "X-Fields"
mock_db.session.scalar.return_value = None
tenant = _tenant()
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant, enable_site=True)
tenant = _create_tenant(db_session_with_containers)
app_model = _create_app(db_session_with_containers, tenant.id)
end_user = SimpleNamespace(id="eu-1")
with app.test_request_context("/site"):
with pytest.raises(Forbidden):
AppSiteApi().get(app_model, end_user)
@patch("controllers.web.site.db")
def test_archived_tenant_raises_forbidden(self, mock_db: MagicMock, app: Flask) -> None:
@patch("controllers.web.site.FeatureService.get_features")
def test_archived_tenant_raises_forbidden(
self, mock_features, app: Flask, db_session_with_containers: Session
) -> None:
app.config["RESTX_MASK_HEADER"] = "X-Fields"
from models.account import TenantStatus
mock_db.session.scalar.return_value = _site()
tenant = SimpleNamespace(
id="tenant-1",
status=TenantStatus.ARCHIVE,
plan="basic",
custom_config_dict={},
)
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant)
tenant = _create_tenant(db_session_with_containers, status=TenantStatus.ARCHIVE)
app_model = _create_app(db_session_with_containers, tenant.id)
_create_site(db_session_with_containers, app_model.id)
end_user = SimpleNamespace(id="eu-1")
mock_features.return_value = SimpleNamespace(can_replace_logo=False)
with app.test_request_context("/site"):
with pytest.raises(Forbidden):
AppSiteApi().get(app_model, end_user)
# ---------------------------------------------------------------------------
# AppSiteInfo
# ---------------------------------------------------------------------------
class TestAppSiteInfo:
def test_basic_fields(self) -> None:
tenant = _tenant()
site_obj = _site()
tenant = SimpleNamespace(id="tenant-1", plan="basic", custom_config_dict={})
site_obj = SimpleNamespace()
info = AppSiteInfo(tenant, SimpleNamespace(id="app-1", enable_site=True), site_obj, "eu-1", False)
assert info.app_id == "app-1"
@@ -118,7 +125,7 @@ class TestAppSiteInfo:
plan="pro",
custom_config_dict={"remove_webapp_brand": True, "replace_webapp_logo": True},
)
site_obj = _site()
site_obj = SimpleNamespace()
info = AppSiteInfo(tenant, SimpleNamespace(id="app-1", enable_site=True), site_obj, "eu-1", True)
assert info.can_replace_logo is True
@@ -88,11 +88,11 @@ class TestPauseStatePersistenceLayerTestContainers:
def setup_test_data(self, db_session_with_containers, file_service, workflow_run_service):
"""Set up test data for each test method using TestContainers."""
# Create test tenant and account
from models.account import Tenant, TenantAccountJoin, TenantAccountRole
from models.account import AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus
tenant = Tenant(
name="Test Tenant",
status="normal",
status=TenantStatus.NORMAL,
)
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
@@ -101,7 +101,7 @@ class TestPauseStatePersistenceLayerTestContainers:
email="test@example.com",
name="Test User",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
db_session_with_containers.add(account)
db_session_with_containers.commit()
@@ -18,7 +18,14 @@ from core.workflow.human_input_compat import (
MemberRecipient,
WebAppDeliveryMethod,
)
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.account import (
Account,
AccountStatus,
Tenant,
TenantAccountJoin,
TenantAccountRole,
TenantStatus,
)
from models.human_input import (
EmailExternalRecipientPayload,
EmailMemberRecipientPayload,
@@ -29,7 +36,7 @@ from models.human_input import (
def _create_tenant_with_members(session: Session, member_emails: list[str]) -> tuple[Tenant, list[Account]]:
tenant = Tenant(name="Test Tenant", status="normal")
tenant = Tenant(name="Test Tenant", status=TenantStatus.NORMAL)
session.add(tenant)
session.flush()
@@ -39,7 +46,7 @@ def _create_tenant_with_members(session: Session, member_emails: list[str]) -> t
email=email,
name=f"Member {index}",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
session.add(account)
session.flush()
@@ -29,7 +29,7 @@ from core.workflow.node_runtime import DifyHumanInputNodeRuntime
from core.workflow.system_variables import build_system_variables
from libs.datetime_utils import naive_utc_now
from models import Account
from models.account import Tenant, TenantAccountJoin, TenantAccountRole
from models.account import AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus
from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
from models.model import App, AppMode, IconType
from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun
@@ -175,7 +175,7 @@ class TestHumanInputResumeNodeExecutionIntegration:
def setup_test_data(self, db_session_with_containers: Session):
tenant = Tenant(
name="Test Tenant",
status="normal",
status=TenantStatus.NORMAL,
)
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
@@ -184,7 +184,7 @@ class TestHumanInputResumeNodeExecutionIntegration:
email="test@example.com",
name="Test User",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
db_session_with_containers.add(account)
db_session_with_containers.commit()
@@ -9,7 +9,7 @@ from werkzeug.exceptions import Unauthorized
from configs import dify_config
from controllers.console.error import AccountNotFound, NotAllowedCreateWorkspace
from models import AccountStatus, TenantAccountJoin
from models import AccountStatus, TenantAccountJoin, TenantStatus
from services.account_service import AccountService, RegisterService, TenantService, TokenPair
from services.errors.account import (
AccountAlreadyInTenantError,
@@ -2851,7 +2851,7 @@ class TestRegisterService:
interface_language="en-US",
password=existing_pending_member_password,
)
existing_account.status = "pending"
existing_account.status = AccountStatus.PENDING
db_session_with_containers.commit()
@@ -2941,7 +2941,7 @@ class TestRegisterService:
interface_language="en-US",
password=already_in_tenant_password,
)
existing_account.status = "active"
existing_account.status = AccountStatus.ACTIVE
db_session_with_containers.commit()
@@ -3331,7 +3331,7 @@ class TestRegisterService:
TenantService.create_tenant_member(tenant, account, role="normal")
# Change tenant status to non-normal
tenant.status = "archive"
tenant.status = TenantStatus.ARCHIVE
db_session_with_containers.commit()
@@ -0,0 +1,613 @@
"""Testcontainers integration tests for DatasetService permission and lifecycle SQL paths."""
from datetime import datetime
from types import SimpleNamespace
from unittest.mock import patch
from uuid import uuid4
import pytest
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import (
AppDatasetJoin,
Dataset,
DatasetAutoDisableLog,
DatasetCollectionBinding,
DatasetPermission,
DatasetPermissionEnum,
)
from models.enums import DataSourceType
from services.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService
from services.errors.account import NoPermissionError
class DatasetPermissionIntegrationFactory:
@staticmethod
def create_account_with_tenant(
db_session_with_containers: Session,
role: TenantAccountRole = TenantAccountRole.OWNER,
) -> tuple[Account, Tenant]:
account = Account(
email=f"{uuid4()}@example.com",
name=f"user-{uuid4()}",
interface_language="en-US",
status="active",
)
tenant = Tenant(name=f"tenant-{uuid4()}", status="normal")
db_session_with_containers.add_all([account, tenant])
db_session_with_containers.flush()
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=role,
current=True,
)
db_session_with_containers.add(join)
db_session_with_containers.commit()
account.role = role
account._current_tenant = tenant
return account, tenant
@staticmethod
def create_account_in_tenant(
db_session_with_containers: Session,
tenant: Tenant,
role: TenantAccountRole = TenantAccountRole.EDITOR,
) -> Account:
account = Account(
email=f"{uuid4()}@example.com",
name=f"user-{uuid4()}",
interface_language="en-US",
status="active",
)
db_session_with_containers.add(account)
db_session_with_containers.flush()
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=role,
current=True,
)
db_session_with_containers.add(join)
db_session_with_containers.commit()
account.role = role
account._current_tenant = tenant
return account
@staticmethod
def create_dataset(
db_session_with_containers: Session,
*,
tenant_id: str,
created_by: str,
name: str | None = None,
permission: DatasetPermissionEnum = DatasetPermissionEnum.ONLY_ME,
indexing_technique: str | None = IndexTechniqueType.HIGH_QUALITY,
enable_api: bool = True,
) -> Dataset:
dataset = Dataset(
tenant_id=tenant_id,
name=name or f"dataset-{uuid4()}",
description="desc",
data_source_type=DataSourceType.UPLOAD_FILE,
indexing_technique=indexing_technique,
created_by=created_by,
provider="vendor",
permission=permission,
retrieval_model={"top_k": 2},
)
dataset.enable_api = enable_api
db_session_with_containers.add(dataset)
db_session_with_containers.commit()
return dataset
@staticmethod
def create_dataset_permission(
db_session_with_containers: Session,
*,
dataset_id: str,
tenant_id: str,
account_id: str,
) -> DatasetPermission:
permission = DatasetPermission(
dataset_id=dataset_id,
tenant_id=tenant_id,
account_id=account_id,
has_permission=True,
)
db_session_with_containers.add(permission)
db_session_with_containers.commit()
return permission
@staticmethod
def create_app_dataset_join(
db_session_with_containers: Session,
*,
dataset_id: str,
) -> AppDatasetJoin:
join = AppDatasetJoin(
app_id=str(uuid4()),
dataset_id=dataset_id,
)
db_session_with_containers.add(join)
db_session_with_containers.commit()
return join
@staticmethod
def create_collection_binding(
db_session_with_containers: Session,
*,
provider_name: str,
model_name: str,
collection_type: str = "dataset",
) -> DatasetCollectionBinding:
binding = DatasetCollectionBinding(
provider_name=provider_name,
model_name=model_name,
collection_name=f"collection_{uuid4().hex}",
type=collection_type,
)
db_session_with_containers.add(binding)
db_session_with_containers.commit()
return binding
@staticmethod
def create_auto_disable_log(
db_session_with_containers: Session,
*,
tenant_id: str,
dataset_id: str,
document_id: str,
) -> DatasetAutoDisableLog:
log = DatasetAutoDisableLog(
tenant_id=tenant_id,
dataset_id=dataset_id,
document_id=document_id,
)
db_session_with_containers.add(log)
db_session_with_containers.commit()
return log
class TestDatasetServicePermissionsAndLifecycle:
def test_delete_dataset_returns_false_when_dataset_is_missing(self, db_session_with_containers: Session):
owner, _tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
result = DatasetService.delete_dataset(str(uuid4()), user=owner)
assert result is False
def test_delete_dataset_checks_permission_and_deletes_dataset(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
)
with patch("services.dataset_service.dataset_was_deleted.send") as send_deleted_signal:
result = DatasetService.delete_dataset(dataset.id, user=owner)
assert result is True
assert db_session_with_containers.get(Dataset, dataset.id) is None
send_deleted_signal.assert_called_once_with(dataset)
def test_dataset_use_check_returns_true_when_join_exists(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
)
DatasetPermissionIntegrationFactory.create_app_dataset_join(
db_session_with_containers,
dataset_id=dataset.id,
)
assert DatasetService.dataset_use_check(dataset.id) is True
def test_dataset_use_check_returns_false_when_join_missing(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
)
assert DatasetService.dataset_use_check(dataset.id) is False
def test_check_dataset_permission_rejects_cross_tenant_access(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
outsider, _other_tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(
db_session_with_containers
)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
)
with pytest.raises(NoPermissionError, match="do not have permission"):
DatasetService.check_dataset_permission(dataset, outsider)
def test_check_dataset_permission_rejects_only_me_dataset_for_non_creator(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
member = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.ONLY_ME,
)
with pytest.raises(NoPermissionError, match="do not have permission"):
DatasetService.check_dataset_permission(dataset, member)
def test_check_dataset_permission_rejects_partial_team_user_without_binding(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
member = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
with pytest.raises(NoPermissionError, match="do not have permission"):
DatasetService.check_dataset_permission(dataset, member)
def test_check_dataset_permission_allows_partial_team_creator(self, db_session_with_containers: Session):
creator, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(
db_session_with_containers,
role=TenantAccountRole.EDITOR,
)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=creator.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
DatasetService.check_dataset_permission(dataset, creator)
def test_check_dataset_permission_allows_partial_team_member_with_binding(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
member = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
DatasetPermissionIntegrationFactory.create_dataset_permission(
db_session_with_containers,
dataset_id=dataset.id,
tenant_id=tenant.id,
account_id=member.id,
)
DatasetService.check_dataset_permission(dataset, member)
def test_check_dataset_operator_permission_rejects_only_me_for_non_creator(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
operator = DatasetPermissionIntegrationFactory.create_account_in_tenant(
db_session_with_containers,
tenant,
role=TenantAccountRole.EDITOR,
)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.ONLY_ME,
)
with pytest.raises(NoPermissionError, match="do not have permission"):
DatasetService.check_dataset_operator_permission(user=operator, dataset=dataset)
def test_check_dataset_operator_permission_rejects_partial_team_without_binding(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
operator = DatasetPermissionIntegrationFactory.create_account_in_tenant(
db_session_with_containers,
tenant,
role=TenantAccountRole.EDITOR,
)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
with pytest.raises(NoPermissionError, match="do not have permission"):
DatasetService.check_dataset_operator_permission(user=operator, dataset=dataset)
def test_check_dataset_operator_permission_allows_partial_team_with_binding(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
operator = DatasetPermissionIntegrationFactory.create_account_in_tenant(
db_session_with_containers,
tenant,
role=TenantAccountRole.EDITOR,
)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
DatasetPermissionIntegrationFactory.create_dataset_permission(
db_session_with_containers,
dataset_id=dataset.id,
tenant_id=tenant.id,
account_id=operator.id,
)
DatasetService.check_dataset_operator_permission(user=operator, dataset=dataset)
def test_update_dataset_api_status_raises_not_found_for_missing_dataset(self, flask_app_with_containers):
with flask_app_with_containers.app_context():
with pytest.raises(NotFound, match="Dataset not found"):
DatasetService.update_dataset_api_status(str(uuid4()), True)
def test_update_dataset_api_status_requires_current_user_id(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
enable_api=False,
)
with patch("services.dataset_service.current_user", SimpleNamespace(id=None)):
with pytest.raises(ValueError, match="Current user or current user id not found"):
DatasetService.update_dataset_api_status(dataset.id, True)
def test_update_dataset_api_status_updates_fields_and_commits(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
enable_api=False,
)
now = datetime(2026, 4, 14, 18, 0, 0)
with (
patch("services.dataset_service.current_user", owner),
patch("services.dataset_service.naive_utc_now", return_value=now),
):
DatasetService.update_dataset_api_status(dataset.id, True)
db_session_with_containers.refresh(dataset)
assert dataset.enable_api is True
assert dataset.updated_by == owner.id
assert dataset.updated_at == now
def test_get_dataset_auto_disable_logs_returns_empty_when_billing_is_disabled(
self, db_session_with_containers: Session
):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
features = SimpleNamespace(
billing=SimpleNamespace(enabled=False, subscription=SimpleNamespace(plan="professional"))
)
with (
patch("services.dataset_service.current_user", owner),
patch("services.dataset_service.FeatureService.get_features", return_value=features),
):
result = DatasetService.get_dataset_auto_disable_logs(str(uuid4()))
assert result == {"document_ids": [], "count": 0}
def test_get_dataset_auto_disable_logs_returns_recent_document_ids(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
)
DatasetPermissionIntegrationFactory.create_auto_disable_log(
db_session_with_containers,
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=str(uuid4()),
)
DatasetPermissionIntegrationFactory.create_auto_disable_log(
db_session_with_containers,
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=str(uuid4()),
)
features = SimpleNamespace(
billing=SimpleNamespace(enabled=True, subscription=SimpleNamespace(plan="professional"))
)
with (
patch("services.dataset_service.current_user", owner),
patch("services.dataset_service.FeatureService.get_features", return_value=features),
):
result = DatasetService.get_dataset_auto_disable_logs(dataset.id)
assert result["count"] == 2
assert len(result["document_ids"]) == 2
class TestDatasetCollectionBindingServiceIntegration:
def test_get_dataset_collection_binding_returns_existing_binding(self, db_session_with_containers: Session):
binding = DatasetPermissionIntegrationFactory.create_collection_binding(
db_session_with_containers,
provider_name="provider",
model_name="model",
)
result = DatasetCollectionBindingService.get_dataset_collection_binding("provider", "model")
assert result.id == binding.id
def test_get_dataset_collection_binding_creates_binding_when_missing(self, db_session_with_containers: Session):
result = DatasetCollectionBindingService.get_dataset_collection_binding("provider", "missing-model")
persisted = db_session_with_containers.get(DatasetCollectionBinding, result.id)
assert persisted is not None
assert persisted.provider_name == "provider"
assert persisted.model_name == "missing-model"
assert persisted.type == "dataset"
assert persisted.collection_name
def test_get_dataset_collection_binding_by_id_and_type_raises_when_missing(self, flask_app_with_containers):
with flask_app_with_containers.app_context():
with pytest.raises(ValueError, match="Dataset collection binding not found"):
DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type(str(uuid4()))
def test_get_dataset_collection_binding_by_id_and_type_returns_binding(self, db_session_with_containers: Session):
binding = DatasetPermissionIntegrationFactory.create_collection_binding(
db_session_with_containers,
provider_name="provider",
model_name="model",
)
result = DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type(binding.id)
assert result.id == binding.id
class TestDatasetPermissionServiceIntegration:
def test_get_dataset_partial_member_list_returns_scalar_results(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
member_a = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
member_b = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
DatasetPermissionIntegrationFactory.create_dataset_permission(
db_session_with_containers,
dataset_id=dataset.id,
tenant_id=tenant.id,
account_id=member_a.id,
)
DatasetPermissionIntegrationFactory.create_dataset_permission(
db_session_with_containers,
dataset_id=dataset.id,
tenant_id=tenant.id,
account_id=member_b.id,
)
result = DatasetPermissionService.get_dataset_partial_member_list(dataset.id)
assert set(result) == {member_a.id, member_b.id}
def test_update_partial_member_list_replaces_permissions_and_commits(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
member_a = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
member_b = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
stale_member = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
DatasetPermissionIntegrationFactory.create_dataset_permission(
db_session_with_containers,
dataset_id=dataset.id,
tenant_id=tenant.id,
account_id=stale_member.id,
)
DatasetPermissionService.update_partial_member_list(
tenant.id,
dataset.id,
[{"user_id": member_a.id}, {"user_id": member_b.id}],
)
permissions = db_session_with_containers.query(DatasetPermission).filter_by(dataset_id=dataset.id).all()
assert {permission.account_id for permission in permissions} == {member_a.id, member_b.id}
def test_check_permission_requires_dataset_editor(self):
user = SimpleNamespace(is_dataset_editor=False, is_dataset_operator=False)
dataset = SimpleNamespace(id="dataset-1", permission=DatasetPermissionEnum.ALL_TEAM)
with pytest.raises(NoPermissionError, match="does not have permission"):
DatasetPermissionService.check_permission(user, dataset, DatasetPermissionEnum.ALL_TEAM, [])
def test_check_permission_prevents_dataset_operator_from_changing_permission_mode(self):
user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True)
dataset = SimpleNamespace(id="dataset-1", permission=DatasetPermissionEnum.ALL_TEAM)
with pytest.raises(NoPermissionError, match="cannot change the dataset permissions"):
DatasetPermissionService.check_permission(user, dataset, DatasetPermissionEnum.ONLY_ME, [])
def test_check_permission_requires_partial_member_list_for_partial_members_mode(self):
user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True)
dataset = SimpleNamespace(id="dataset-1", permission=DatasetPermissionEnum.PARTIAL_TEAM)
with pytest.raises(ValueError, match="Partial member list is required"):
DatasetPermissionService.check_permission(user, dataset, DatasetPermissionEnum.PARTIAL_TEAM, [])
def test_check_permission_rejects_dataset_operator_member_list_changes(self):
user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True)
dataset = SimpleNamespace(id="dataset-1", permission=DatasetPermissionEnum.PARTIAL_TEAM)
with patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["user-1"]):
with pytest.raises(ValueError, match="cannot change the dataset permissions"):
DatasetPermissionService.check_permission(
user,
dataset,
DatasetPermissionEnum.PARTIAL_TEAM,
[{"user_id": "user-2"}],
)
def test_check_permission_allows_dataset_operator_when_member_list_is_unchanged(self):
user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True)
dataset = SimpleNamespace(id="dataset-1", permission=DatasetPermissionEnum.PARTIAL_TEAM)
with patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["user-1"]):
DatasetPermissionService.check_permission(
user,
dataset,
DatasetPermissionEnum.PARTIAL_TEAM,
[{"user_id": "user-1"}],
)
def test_clear_partial_member_list_deletes_permissions_and_commits(self, db_session_with_containers: Session):
owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers)
member = DatasetPermissionIntegrationFactory.create_account_in_tenant(db_session_with_containers, tenant)
dataset = DatasetPermissionIntegrationFactory.create_dataset(
db_session_with_containers,
tenant_id=tenant.id,
created_by=owner.id,
permission=DatasetPermissionEnum.PARTIAL_TEAM,
)
DatasetPermissionIntegrationFactory.create_dataset_permission(
db_session_with_containers,
dataset_id=dataset.id,
tenant_id=tenant.id,
account_id=member.id,
)
DatasetPermissionService.clear_partial_member_list(dataset.id)
remaining = db_session_with_containers.query(DatasetPermission).filter_by(dataset_id=dataset.id).all()
assert remaining == []
@@ -7,7 +7,14 @@ from graphon.model_runtime.entities.model_entities import ModelType
from sqlalchemy.orm import Session
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.account import (
Account,
AccountStatus,
Tenant,
TenantAccountJoin,
TenantAccountRole,
TenantStatus,
)
from models.dataset import Dataset, ExternalKnowledgeApis, ExternalKnowledgeBindings
from models.enums import DataSourceType
from services.dataset_service import DatasetService
@@ -26,12 +33,12 @@ class DatasetUpdateTestDataFactory:
email=f"{uuid4()}@example.com",
name=f"user-{uuid4()}",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
db_session_with_containers.add(account)
db_session_with_containers.commit()
tenant = Tenant(name=f"tenant-{account.id}", status="normal")
tenant = Tenant(name=f"tenant-{account.id}", status=TenantStatus.NORMAL)
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
@@ -0,0 +1,387 @@
"""Testcontainers integration tests for schedule service SQL-backed behavior."""
from datetime import datetime
from types import SimpleNamespace
from uuid import uuid4
import pytest
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from core.workflow.nodes.trigger_schedule.entities import ScheduleConfig, SchedulePlanUpdate
from core.workflow.nodes.trigger_schedule.exc import ScheduleNotFoundError
from events.event_handlers.sync_workflow_schedule_when_app_published import sync_schedule_from_workflow
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.trigger import WorkflowSchedulePlan
from services.errors.account import AccountNotFoundError
from services.trigger.schedule_service import ScheduleService
class ScheduleServiceIntegrationFactory:
@staticmethod
def create_account_with_tenant(
db_session_with_containers: Session,
role: TenantAccountRole = TenantAccountRole.OWNER,
) -> tuple[Account, Tenant]:
account = Account(
email=f"{uuid4()}@example.com",
name=f"user-{uuid4()}",
interface_language="en-US",
status="active",
)
tenant = Tenant(name=f"tenant-{uuid4()}", status="normal")
db_session_with_containers.add_all([account, tenant])
db_session_with_containers.flush()
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=role,
current=True,
)
db_session_with_containers.add(join)
db_session_with_containers.commit()
account.current_tenant = tenant
return account, tenant
@staticmethod
def create_schedule_plan(
db_session_with_containers: Session,
*,
tenant_id: str,
app_id: str | None = None,
node_id: str = "start",
cron_expression: str = "30 10 * * *",
timezone: str = "UTC",
next_run_at: datetime | None = None,
) -> WorkflowSchedulePlan:
schedule = WorkflowSchedulePlan(
tenant_id=tenant_id,
app_id=app_id or str(uuid4()),
node_id=node_id,
cron_expression=cron_expression,
timezone=timezone,
next_run_at=next_run_at,
)
db_session_with_containers.add(schedule)
db_session_with_containers.commit()
return schedule
def _cron_workflow(
*,
node_id: str = "start",
cron_expression: str = "30 10 * * *",
timezone: str = "UTC",
):
return SimpleNamespace(
graph_dict={
"nodes": [
{
"id": node_id,
"data": {
"type": "trigger-schedule",
"mode": "cron",
"cron_expression": cron_expression,
"timezone": timezone,
},
}
]
}
)
def _no_schedule_workflow():
return SimpleNamespace(
graph_dict={
"nodes": [
{
"id": "node-1",
"data": {"type": "llm"},
}
]
}
)
class TestScheduleServiceIntegration:
def test_create_schedule_persists_schedule(self, db_session_with_containers: Session):
account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
expected_next_run = datetime(2026, 1, 1, 10, 30, 0)
config = ScheduleConfig(
node_id="start",
cron_expression="30 10 * * *",
timezone="UTC",
)
with pytest.MonkeyPatch.context() as monkeypatch:
monkeypatch.setattr(
"services.trigger.schedule_service.calculate_next_run_at",
lambda *_args, **_kwargs: expected_next_run,
)
schedule = ScheduleService.create_schedule(
session=db_session_with_containers,
tenant_id=tenant.id,
app_id=str(uuid4()),
config=config,
)
persisted = db_session_with_containers.get(WorkflowSchedulePlan, schedule.id)
assert persisted is not None
assert persisted.tenant_id == tenant.id
assert persisted.node_id == "start"
assert persisted.cron_expression == "30 10 * * *"
assert persisted.timezone == "UTC"
assert persisted.next_run_at == expected_next_run
def test_update_schedule_updates_fields_and_recomputes_next_run(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
schedule = ScheduleServiceIntegrationFactory.create_schedule_plan(
db_session_with_containers,
tenant_id=tenant.id,
cron_expression="30 10 * * *",
timezone="UTC",
)
expected_next_run = datetime(2026, 1, 2, 12, 0, 0)
with pytest.MonkeyPatch.context() as monkeypatch:
monkeypatch.setattr(
"services.trigger.schedule_service.calculate_next_run_at",
lambda *_args, **_kwargs: expected_next_run,
)
updated = ScheduleService.update_schedule(
session=db_session_with_containers,
schedule_id=schedule.id,
updates=SchedulePlanUpdate(
cron_expression="0 12 * * *",
timezone="America/New_York",
),
)
db_session_with_containers.refresh(updated)
assert updated.cron_expression == "0 12 * * *"
assert updated.timezone == "America/New_York"
assert updated.next_run_at == expected_next_run
def test_update_schedule_updates_only_node_id_without_recomputing_time(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
initial_next_run = datetime(2026, 1, 1, 10, 0, 0)
schedule = ScheduleServiceIntegrationFactory.create_schedule_plan(
db_session_with_containers,
tenant_id=tenant.id,
next_run_at=initial_next_run,
)
with pytest.MonkeyPatch.context() as monkeypatch:
calls: list[tuple] = []
def _track(*args, **kwargs):
calls.append((args, kwargs))
return datetime(2026, 1, 9, 10, 0, 0)
monkeypatch.setattr("services.trigger.schedule_service.calculate_next_run_at", _track)
updated = ScheduleService.update_schedule(
session=db_session_with_containers,
schedule_id=schedule.id,
updates=SchedulePlanUpdate(node_id="node-new"),
)
db_session_with_containers.refresh(updated)
assert updated.node_id == "node-new"
assert updated.next_run_at == initial_next_run
assert calls == []
def test_update_schedule_not_found_raises(self, db_session_with_containers: Session):
with pytest.raises(ScheduleNotFoundError, match="Schedule not found"):
ScheduleService.update_schedule(
session=db_session_with_containers,
schedule_id=str(uuid4()),
updates=SchedulePlanUpdate(node_id="node-new"),
)
def test_delete_schedule_removes_row(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
schedule = ScheduleServiceIntegrationFactory.create_schedule_plan(
db_session_with_containers,
tenant_id=tenant.id,
)
ScheduleService.delete_schedule(
session=db_session_with_containers,
schedule_id=schedule.id,
)
db_session_with_containers.commit()
assert db_session_with_containers.get(WorkflowSchedulePlan, schedule.id) is None
def test_delete_schedule_not_found_raises(self, db_session_with_containers: Session):
with pytest.raises(ScheduleNotFoundError, match="Schedule not found"):
ScheduleService.delete_schedule(
session=db_session_with_containers,
schedule_id=str(uuid4()),
)
def test_get_tenant_owner_returns_owner_account(self, db_session_with_containers: Session):
owner, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(
db_session_with_containers,
role=TenantAccountRole.OWNER,
)
result = ScheduleService.get_tenant_owner(
session=db_session_with_containers,
tenant_id=tenant.id,
)
assert result.id == owner.id
def test_get_tenant_owner_falls_back_to_admin(self, db_session_with_containers: Session):
admin, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(
db_session_with_containers,
role=TenantAccountRole.ADMIN,
)
result = ScheduleService.get_tenant_owner(
session=db_session_with_containers,
tenant_id=tenant.id,
)
assert result.id == admin.id
def test_get_tenant_owner_raises_when_account_record_missing(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
db_session_with_containers.execute(delete(TenantAccountJoin))
missing_account_id = str(uuid4())
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=missing_account_id,
role=TenantAccountRole.OWNER,
current=True,
)
db_session_with_containers.add(join)
db_session_with_containers.commit()
with pytest.raises(AccountNotFoundError, match=missing_account_id):
ScheduleService.get_tenant_owner(session=db_session_with_containers, tenant_id=tenant.id)
def test_get_tenant_owner_raises_when_no_owner_or_admin_found(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
db_session_with_containers.execute(delete(TenantAccountJoin))
db_session_with_containers.commit()
with pytest.raises(AccountNotFoundError, match=tenant.id):
ScheduleService.get_tenant_owner(session=db_session_with_containers, tenant_id=tenant.id)
def test_update_next_run_at_updates_persisted_value(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
schedule = ScheduleServiceIntegrationFactory.create_schedule_plan(
db_session_with_containers,
tenant_id=tenant.id,
)
expected_next_run = datetime(2026, 1, 3, 10, 30, 0)
with pytest.MonkeyPatch.context() as monkeypatch:
monkeypatch.setattr(
"services.trigger.schedule_service.calculate_next_run_at",
lambda *_args, **_kwargs: expected_next_run,
)
result = ScheduleService.update_next_run_at(
session=db_session_with_containers,
schedule_id=schedule.id,
)
db_session_with_containers.refresh(schedule)
assert result == expected_next_run
assert schedule.next_run_at == expected_next_run
def test_update_next_run_at_raises_when_schedule_not_found(self, db_session_with_containers: Session):
with pytest.raises(ScheduleNotFoundError, match="Schedule not found"):
ScheduleService.update_next_run_at(
session=db_session_with_containers,
schedule_id=str(uuid4()),
)
class TestSyncScheduleFromWorkflowIntegration:
def test_sync_schedule_create_new(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
app_id = str(uuid4())
expected_next_run = datetime(2026, 1, 4, 10, 30, 0)
with pytest.MonkeyPatch.context() as monkeypatch:
monkeypatch.setattr(
"services.trigger.schedule_service.calculate_next_run_at",
lambda *_args, **_kwargs: expected_next_run,
)
result = sync_schedule_from_workflow(
tenant_id=tenant.id,
app_id=app_id,
workflow=_cron_workflow(),
)
assert result is not None
persisted = db_session_with_containers.execute(
select(WorkflowSchedulePlan).where(WorkflowSchedulePlan.app_id == app_id)
).scalar_one()
assert persisted.node_id == "start"
assert persisted.cron_expression == "30 10 * * *"
assert persisted.timezone == "UTC"
assert persisted.next_run_at == expected_next_run
def test_sync_schedule_update_existing(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
app_id = str(uuid4())
existing = ScheduleServiceIntegrationFactory.create_schedule_plan(
db_session_with_containers,
tenant_id=tenant.id,
app_id=app_id,
node_id="old-start",
cron_expression="30 10 * * *",
timezone="UTC",
)
existing_id = existing.id
expected_next_run = datetime(2026, 1, 5, 12, 0, 0)
with pytest.MonkeyPatch.context() as monkeypatch:
monkeypatch.setattr(
"services.trigger.schedule_service.calculate_next_run_at",
lambda *_args, **_kwargs: expected_next_run,
)
result = sync_schedule_from_workflow(
tenant_id=tenant.id,
app_id=app_id,
workflow=_cron_workflow(
node_id="start",
cron_expression="0 12 * * *",
timezone="America/New_York",
),
)
assert result is not None
db_session_with_containers.expire_all()
persisted = db_session_with_containers.get(WorkflowSchedulePlan, existing_id)
assert persisted is not None
assert persisted.node_id == "start"
assert persisted.cron_expression == "0 12 * * *"
assert persisted.timezone == "America/New_York"
assert persisted.next_run_at == expected_next_run
def test_sync_schedule_remove_when_no_config(self, db_session_with_containers: Session):
_account, tenant = ScheduleServiceIntegrationFactory.create_account_with_tenant(db_session_with_containers)
app_id = str(uuid4())
existing = ScheduleServiceIntegrationFactory.create_schedule_plan(
db_session_with_containers,
tenant_id=tenant.id,
app_id=app_id,
)
existing_id = existing.id
result = sync_schedule_from_workflow(
tenant_id=tenant.id,
app_id=app_id,
workflow=_no_schedule_workflow(),
)
assert result is None
db_session_with_containers.expire_all()
assert db_session_with_containers.get(WorkflowSchedulePlan, existing_id) is None
@@ -0,0 +1,507 @@
from __future__ import annotations
import json
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from uuid import uuid4
import pytest
from sqlalchemy import select
from sqlalchemy.orm import Session
from core.trigger.constants import TRIGGER_WEBHOOK_NODE_TYPE
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.enums import AppTriggerStatus, AppTriggerType
from models.model import App
from models.trigger import AppTrigger, WorkflowWebhookTrigger
from models.workflow import Workflow
from services.errors.app import QuotaExceededError
from services.trigger.webhook_service import WebhookService
class WebhookServiceRelationshipFactory:
@staticmethod
def create_account_and_tenant(db_session_with_containers: Session) -> tuple[Account, Tenant]:
account = Account(
name=f"Account {uuid4()}",
email=f"webhook-{uuid4()}@example.com",
password="hashed-password",
password_salt="salt",
interface_language="en-US",
timezone="UTC",
)
db_session_with_containers.add(account)
db_session_with_containers.commit()
tenant = Tenant(name=f"Tenant {uuid4()}", plan="basic", status="normal")
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER,
current=True,
)
db_session_with_containers.add(join)
db_session_with_containers.commit()
account.current_tenant = tenant
return account, tenant
@staticmethod
def create_app(db_session_with_containers: Session, tenant: Tenant, account: Account) -> App:
app = App(
tenant_id=tenant.id,
name=f"Webhook App {uuid4()}",
description="",
mode="workflow",
icon_type="emoji",
icon="bot",
icon_background="#FFFFFF",
enable_site=False,
enable_api=True,
api_rpm=100,
api_rph=100,
is_demo=False,
is_public=False,
is_universal=False,
created_by=account.id,
updated_by=account.id,
)
db_session_with_containers.add(app)
db_session_with_containers.commit()
return app
@staticmethod
def create_workflow(
db_session_with_containers: Session,
*,
app: App,
account: Account,
node_ids: list[str],
version: str,
) -> Workflow:
graph = {
"nodes": [
{
"id": node_id,
"data": {
"type": TRIGGER_WEBHOOK_NODE_TYPE,
"title": f"Webhook {node_id}",
"method": "post",
"content_type": "application/json",
"headers": [],
"params": [],
"body": [],
"status_code": 200,
"response_body": '{"status": "ok"}',
"timeout": 30,
},
}
for node_id in node_ids
],
"edges": [],
}
workflow = Workflow(
tenant_id=app.tenant_id,
app_id=app.id,
type="workflow",
graph=json.dumps(graph),
features=json.dumps({}),
created_by=account.id,
updated_by=account.id,
environment_variables=[],
conversation_variables=[],
version=version,
)
db_session_with_containers.add(workflow)
db_session_with_containers.commit()
return workflow
@staticmethod
def create_webhook_trigger(
db_session_with_containers: Session,
*,
app: App,
account: Account,
node_id: str,
webhook_id: str | None = None,
) -> WorkflowWebhookTrigger:
webhook_trigger = WorkflowWebhookTrigger(
app_id=app.id,
node_id=node_id,
tenant_id=app.tenant_id,
webhook_id=webhook_id or uuid4().hex[:24],
created_by=account.id,
)
db_session_with_containers.add(webhook_trigger)
db_session_with_containers.commit()
return webhook_trigger
@staticmethod
def create_app_trigger(
db_session_with_containers: Session,
*,
app: App,
node_id: str,
status: AppTriggerStatus,
) -> AppTrigger:
app_trigger = AppTrigger(
tenant_id=app.tenant_id,
app_id=app.id,
node_id=node_id,
trigger_type=AppTriggerType.TRIGGER_WEBHOOK,
provider_name="webhook",
title=f"Webhook {node_id}",
status=status,
)
db_session_with_containers.add(app_trigger)
db_session_with_containers.commit()
return app_trigger
class TestWebhookServiceLookupWithContainers:
def test_get_webhook_trigger_and_workflow_raises_when_app_trigger_missing(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version="2026-04-14.001"
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
with pytest.raises(ValueError, match="App trigger not found"):
WebhookService.get_webhook_trigger_and_workflow(webhook_trigger.webhook_id)
def test_get_webhook_trigger_and_workflow_raises_when_app_trigger_rate_limited(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version="2026-04-14.001"
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
factory.create_app_trigger(
db_session_with_containers, app=app, node_id="node-1", status=AppTriggerStatus.RATE_LIMITED
)
with pytest.raises(ValueError, match="rate limited"):
WebhookService.get_webhook_trigger_and_workflow(webhook_trigger.webhook_id)
def test_get_webhook_trigger_and_workflow_raises_when_app_trigger_disabled(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version="2026-04-14.001"
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
factory.create_app_trigger(
db_session_with_containers, app=app, node_id="node-1", status=AppTriggerStatus.DISABLED
)
with pytest.raises(ValueError, match="disabled"):
WebhookService.get_webhook_trigger_and_workflow(webhook_trigger.webhook_id)
def test_get_webhook_trigger_and_workflow_raises_when_workflow_missing(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
factory.create_app_trigger(
db_session_with_containers, app=app, node_id="node-1", status=AppTriggerStatus.ENABLED
)
with pytest.raises(ValueError, match="Workflow not found"):
WebhookService.get_webhook_trigger_and_workflow(webhook_trigger.webhook_id)
def test_get_webhook_trigger_and_workflow_returns_debug_draft_workflow(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
factory.create_workflow(
db_session_with_containers,
app=app,
account=account,
node_ids=["published-node"],
version="2026-04-14.001",
)
draft_workflow = factory.create_workflow(
db_session_with_containers,
app=app,
account=account,
node_ids=["debug-node"],
version=Workflow.VERSION_DRAFT,
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="debug-node"
)
got_trigger, got_workflow, got_node_config = WebhookService.get_webhook_trigger_and_workflow(
webhook_trigger.webhook_id,
is_debug=True,
)
assert got_trigger.id == webhook_trigger.id
assert got_workflow.id == draft_workflow.id
assert got_node_config["id"] == "debug-node"
class TestWebhookServiceTriggerExecutionWithContainers:
def test_trigger_workflow_execution_triggers_async_workflow_successfully(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
workflow = factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version="2026-04-14.001"
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
end_user = SimpleNamespace(id=str(uuid4()))
webhook_data = {"body": {"value": 1}, "headers": {}, "query_params": {}, "files": {}, "method": "POST"}
with (
patch(
"services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type",
return_value=end_user,
),
patch("services.trigger.webhook_service.QuotaType.TRIGGER.consume") as mock_consume,
patch("services.trigger.webhook_service.AsyncWorkflowService.trigger_workflow_async") as mock_trigger,
):
WebhookService.trigger_workflow_execution(webhook_trigger, webhook_data, workflow)
mock_consume.assert_called_once_with(webhook_trigger.tenant_id)
mock_trigger.assert_called_once()
trigger_args = mock_trigger.call_args.args
assert trigger_args[1] is end_user
assert trigger_args[2].workflow_id == workflow.id
assert trigger_args[2].root_node_id == webhook_trigger.node_id
def test_trigger_workflow_execution_marks_tenant_rate_limited_when_quota_exceeded(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
workflow = factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version="2026-04-14.001"
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
with (
patch(
"services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type",
return_value=SimpleNamespace(id=str(uuid4())),
),
patch(
"services.trigger.webhook_service.QuotaType.TRIGGER.consume",
side_effect=QuotaExceededError(feature="trigger", tenant_id=tenant.id, required=1),
),
patch(
"services.trigger.webhook_service.AppTriggerService.mark_tenant_triggers_rate_limited"
) as mock_mark_rate_limited,
):
with pytest.raises(QuotaExceededError):
WebhookService.trigger_workflow_execution(
webhook_trigger,
{"body": {}, "headers": {}, "query_params": {}, "files": {}, "method": "POST"},
workflow,
)
mock_mark_rate_limited.assert_called_once_with(tenant.id)
def test_trigger_workflow_execution_logs_and_reraises_unexpected_errors(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
workflow = factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version="2026-04-14.001"
)
webhook_trigger = factory.create_webhook_trigger(
db_session_with_containers, app=app, account=account, node_id="node-1"
)
with (
patch(
"services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type",
side_effect=RuntimeError("boom"),
),
patch("services.trigger.webhook_service.logger.exception") as mock_logger_exception,
):
with pytest.raises(RuntimeError, match="boom"):
WebhookService.trigger_workflow_execution(
webhook_trigger,
{"body": {}, "headers": {}, "query_params": {}, "files": {}, "method": "POST"},
workflow,
)
mock_logger_exception.assert_called_once()
class TestWebhookServiceRelationshipSyncWithContainers:
def test_sync_webhook_relationships_raises_when_workflow_exceeds_node_limit(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
node_ids = [f"node-{index}" for index in range(WebhookService.MAX_WEBHOOK_NODES_PER_WORKFLOW + 1)]
workflow = factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=node_ids, version=Workflow.VERSION_DRAFT
)
with pytest.raises(ValueError, match="maximum webhook node limit"):
WebhookService.sync_webhook_relationships(app, workflow)
def test_sync_webhook_relationships_raises_when_lock_not_acquired(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
workflow = factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=["node-1"], version=Workflow.VERSION_DRAFT
)
lock = MagicMock()
lock.acquire.return_value = False
with patch("services.trigger.webhook_service.redis_client.lock", return_value=lock):
with pytest.raises(RuntimeError, match="Failed to acquire lock"):
WebhookService.sync_webhook_relationships(app, workflow)
def test_sync_webhook_relationships_creates_missing_records_and_deletes_stale_records(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
stale_trigger = factory.create_webhook_trigger(
db_session_with_containers,
app=app,
account=account,
node_id="node-stale",
webhook_id="stale-webhook-id-000001",
)
stale_trigger_id = stale_trigger.id
workflow = factory.create_workflow(
db_session_with_containers,
app=app,
account=account,
node_ids=["node-new"],
version=Workflow.VERSION_DRAFT,
)
with patch(
"services.trigger.webhook_service.WebhookService.generate_webhook_id", return_value="new-webhook-id-000001"
):
WebhookService.sync_webhook_relationships(app, workflow)
db_session_with_containers.expire_all()
records = db_session_with_containers.scalars(
select(WorkflowWebhookTrigger).where(WorkflowWebhookTrigger.app_id == app.id)
).all()
assert [record.node_id for record in records] == ["node-new"]
assert records[0].webhook_id == "new-webhook-id-000001"
assert db_session_with_containers.get(WorkflowWebhookTrigger, stale_trigger_id) is None
def test_sync_webhook_relationships_sets_redis_cache_for_new_record(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
workflow = factory.create_workflow(
db_session_with_containers,
app=app,
account=account,
node_ids=["node-cache"],
version=Workflow.VERSION_DRAFT,
)
cache_key = f"{WebhookService.__WEBHOOK_NODE_CACHE_KEY__}:{app.id}:node-cache"
with patch(
"services.trigger.webhook_service.WebhookService.generate_webhook_id", return_value="cache-webhook-id-00001"
):
WebhookService.sync_webhook_relationships(app, workflow)
cached_payload = WebhookServiceRelationshipFactory._read_cache(cache_key)
assert cached_payload is not None
assert cached_payload["node_id"] == "node-cache"
assert cached_payload["webhook_id"] == "cache-webhook-id-00001"
def test_sync_webhook_relationships_logs_when_lock_release_fails(
self, db_session_with_containers: Session, flask_app_with_containers
):
del flask_app_with_containers
factory = WebhookServiceRelationshipFactory
account, tenant = factory.create_account_and_tenant(db_session_with_containers)
app = factory.create_app(db_session_with_containers, tenant, account)
workflow = factory.create_workflow(
db_session_with_containers, app=app, account=account, node_ids=[], version=Workflow.VERSION_DRAFT
)
lock = MagicMock()
lock.acquire.return_value = True
lock.release.side_effect = RuntimeError("release failed")
with (
patch("services.trigger.webhook_service.redis_client.lock", return_value=lock),
patch("services.trigger.webhook_service.logger.exception") as mock_logger_exception,
):
WebhookService.sync_webhook_relationships(app, workflow)
mock_logger_exception.assert_called_once()
def _read_cache(cache_key: str) -> dict[str, str] | None:
from extensions.ext_redis import redis_client
cached = redis_client.get(cache_key)
if not cached:
return None
if isinstance(cached, bytes):
cached = cached.decode("utf-8")
return json.loads(cached)
WebhookServiceRelationshipFactory._read_cache = staticmethod(_read_cache)
@@ -16,7 +16,7 @@ from sqlalchemy import delete, func, select, update
from core.indexing_runner import DocumentIsPausedError, IndexingRunner
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus
from models.dataset import Dataset, Document, DocumentSegment
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus
from tasks.document_indexing_sync_task import document_indexing_sync_task
@@ -31,12 +31,12 @@ class DocumentIndexingSyncTaskTestDataFactory:
email=f"{uuid4()}@example.com",
name=f"user-{uuid4()}",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
db_session_with_containers.add(account)
db_session_with_containers.flush()
tenant = Tenant(name=f"tenant-{account.id}", status="normal")
tenant = Tenant(name=f"tenant-{account.id}", status=TenantStatus.NORMAL)
db_session_with_containers.add(tenant)
db_session_with_containers.flush()
@@ -33,7 +33,7 @@ from extensions.ext_storage import storage
from libs.datetime_utils import naive_utc_now
from models import Account
from models import WorkflowPause as WorkflowPauseModel
from models.account import Tenant, TenantAccountJoin, TenantAccountRole
from models.account import AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus
from models.model import UploadFile
from models.workflow import Workflow, WorkflowRun
from repositories.sqlalchemy_api_workflow_run_repository import (
@@ -181,7 +181,7 @@ class TestWorkflowPauseIntegration:
tenant = Tenant(
name="Test Tenant",
status="normal",
status=TenantStatus.NORMAL,
)
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
@@ -190,7 +190,7 @@ class TestWorkflowPauseIntegration:
email="test@example.com",
name="Test User",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
db_session_with_containers.add(account)
db_session_with_containers.commit()
@@ -696,7 +696,7 @@ class TestWorkflowPauseIntegration:
tenant2 = Tenant(
name="Test Tenant 2",
status="normal",
status=TenantStatus.NORMAL,
)
self.session.add(tenant2)
self.session.commit()
@@ -705,7 +705,7 @@ class TestWorkflowPauseIntegration:
email="test2@example.com",
name="Test User 2",
interface_language="en-US",
status="active",
status=AccountStatus.ACTIVE,
)
self.session.add(account2)
self.session.commit()
@@ -0,0 +1,139 @@
"""Unit tests for console app import endpoints."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from controllers.console.app import app_import as app_import_module
from services.app_dsl_service import ImportStatus
def _unwrap(func):
bound_self = getattr(func, "__self__", None)
while hasattr(func, "__wrapped__"):
func = func.__wrapped__
if bound_self is not None:
return func.__get__(bound_self, bound_self.__class__)
return func
class _Result:
def __init__(self, status: ImportStatus, app_id: str | None = "app-1"):
self.status = status
self.app_id = app_id
def model_dump(self, mode: str = "json"):
return {"status": self.status, "app_id": self.app_id}
def _install_features(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
features = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=enabled))
monkeypatch.setattr(app_import_module.FeatureService, "get_system_features", lambda: features)
def _mock_session(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
fake_session = MagicMock()
fake_session.__enter__.return_value = fake_session
fake_session.__exit__.return_value = None
monkeypatch.setattr(app_import_module, "db", SimpleNamespace(engine=object()))
monkeypatch.setattr(app_import_module, "Session", lambda *_args, **_kwargs: fake_session)
return fake_session
class TestAppImportApi:
@pytest.fixture
def api(self):
return app_import_module.AppImportApi()
def test_import_post_returns_failed_status_and_rolls_back(self, api, app, monkeypatch: pytest.MonkeyPatch) -> None:
method = _unwrap(api.post)
_install_features(monkeypatch, enabled=False)
session = _mock_session(monkeypatch)
monkeypatch.setattr(
app_import_module.AppDslService,
"import_app",
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED, app_id=None),
)
monkeypatch.setattr(app_import_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
response, status = method()
session.rollback.assert_called_once_with()
session.commit.assert_not_called()
assert status == 400
assert response["status"] == ImportStatus.FAILED
def test_import_post_returns_pending_status_and_commits(self, api, app, monkeypatch: pytest.MonkeyPatch) -> None:
method = _unwrap(api.post)
_install_features(monkeypatch, enabled=False)
session = _mock_session(monkeypatch)
monkeypatch.setattr(
app_import_module.AppDslService,
"import_app",
lambda *_args, **_kwargs: _Result(ImportStatus.PENDING),
)
monkeypatch.setattr(app_import_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
response, status = method()
session.commit.assert_called_once_with()
session.rollback.assert_not_called()
assert status == 202
assert response["status"] == ImportStatus.PENDING
def test_import_post_updates_webapp_auth_when_enabled(self, api, app, monkeypatch: pytest.MonkeyPatch) -> None:
method = _unwrap(api.post)
_install_features(monkeypatch, enabled=True)
session = _mock_session(monkeypatch)
monkeypatch.setattr(
app_import_module.AppDslService,
"import_app",
lambda *_args, **_kwargs: _Result(ImportStatus.COMPLETED, app_id="app-123"),
)
update_access = MagicMock()
monkeypatch.setattr(app_import_module.EnterpriseService.WebAppAuth, "update_app_access_mode", update_access)
monkeypatch.setattr(app_import_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
response, status = method()
session.commit.assert_called_once_with()
session.rollback.assert_not_called()
update_access.assert_called_once_with("app-123", "private")
assert status == 200
assert response["status"] == ImportStatus.COMPLETED
class TestAppImportConfirmApi:
@pytest.fixture
def api(self):
return app_import_module.AppImportConfirmApi()
def test_import_confirm_returns_failed_status_and_rolls_back(
self, api, app, monkeypatch: pytest.MonkeyPatch
) -> None:
method = _unwrap(api.post)
session = _mock_session(monkeypatch)
monkeypatch.setattr(
app_import_module.AppDslService,
"confirm_import",
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED),
)
monkeypatch.setattr(app_import_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
with app.test_request_context("/console/api/apps/imports/import-1/confirm", method="POST"):
response, status = method(import_id="import-1")
session.rollback.assert_called_once_with()
session.commit.assert_not_called()
assert status == 400
assert response["status"] == ImportStatus.FAILED
@@ -138,12 +138,15 @@ def app_models(app_module):
def patch_signed_url(monkeypatch, app_module):
"""Ensure icon URL generation uses a deterministic helper for tests."""
def _fake_signed_url(key: str | None) -> str | None:
if not key:
def _fake_build_icon_url(_icon_type, key: str | None) -> str | None:
if key is None:
return None
icon_type = str(_icon_type).lower()
if icon_type != "image":
return None
return f"signed:{key}"
monkeypatch.setattr(app_module.file_helpers, "get_signed_file_url", _fake_signed_url)
monkeypatch.setattr(app_module, "build_icon_url", _fake_build_icon_url)
def _ts(hour: int = 12) -> datetime:
@@ -0,0 +1,108 @@
from __future__ import annotations
from contextlib import nullcontext
from datetime import UTC, datetime
from types import SimpleNamespace
import pytest
from graphon.variables.types import SegmentType
from pydantic import ValidationError
from controllers.console.app import conversation_variables as conversation_variables_module
def _unwrap(func):
bound_self = getattr(func, "__self__", None)
while hasattr(func, "__wrapped__"):
func = func.__wrapped__
if bound_self is not None:
return func.__get__(bound_self, bound_self.__class__)
return func
def test_get_conversation_variables_returns_paginated_response(app, monkeypatch: pytest.MonkeyPatch) -> None:
api = conversation_variables_module.ConversationVariablesApi()
method = _unwrap(api.get)
created_at = datetime(2026, 1, 1, tzinfo=UTC)
updated_at = datetime(2026, 1, 2, tzinfo=UTC)
row = SimpleNamespace(
created_at=created_at,
updated_at=updated_at,
to_variable=lambda: SimpleNamespace(
model_dump=lambda: {
"id": "var-1",
"name": "my_var",
"value_type": "string",
"value": "value",
"description": "desc",
}
),
)
session = SimpleNamespace(scalars=lambda _stmt: SimpleNamespace(all=lambda: [row]))
monkeypatch.setattr(conversation_variables_module, "db", SimpleNamespace(engine=object()))
monkeypatch.setattr(
conversation_variables_module,
"sessionmaker",
lambda *_args, **_kwargs: SimpleNamespace(begin=lambda: nullcontext(session)),
)
with app.test_request_context(
"/console/api/apps/app-1/conversation-variables",
method="GET",
query_string={"conversation_id": "conv-1"},
):
response = method(app_model=SimpleNamespace(id="app-1"))
assert response["page"] == 1
assert response["limit"] == 100
assert response["total"] == 1
assert response["has_more"] is False
assert response["data"][0]["id"] == "var-1"
assert response["data"][0]["created_at"] == int(created_at.timestamp())
assert response["data"][0]["updated_at"] == int(updated_at.timestamp())
def test_get_conversation_variables_normalizes_value_type_and_value(app, monkeypatch: pytest.MonkeyPatch) -> None:
api = conversation_variables_module.ConversationVariablesApi()
method = _unwrap(api.get)
row = SimpleNamespace(
created_at=None,
updated_at=None,
to_variable=lambda: SimpleNamespace(
model_dump=lambda: {
"id": "var-2",
"name": "my_var_2",
"value_type": SegmentType.INTEGER,
"value": 42,
"description": None,
}
),
)
session = SimpleNamespace(scalars=lambda _stmt: SimpleNamespace(all=lambda: [row]))
monkeypatch.setattr(conversation_variables_module, "db", SimpleNamespace(engine=object()))
monkeypatch.setattr(
conversation_variables_module,
"sessionmaker",
lambda *_args, **_kwargs: SimpleNamespace(begin=lambda: nullcontext(session)),
)
with app.test_request_context(
"/console/api/apps/app-1/conversation-variables",
method="GET",
query_string={"conversation_id": "conv-1"},
):
response = method(app_model=SimpleNamespace(id="app-1"))
assert response["data"][0]["value_type"] == "number"
assert response["data"][0]["value"] == "42"
def test_get_conversation_variables_requires_conversation_id(app) -> None:
api = conversation_variables_module.ConversationVariablesApi()
method = _unwrap(api.get)
with app.test_request_context("/console/api/apps/app-1/conversation-variables", method="GET"):
with pytest.raises(ValidationError):
method(app_model=SimpleNamespace(id="app-1"))
@@ -1,5 +1,7 @@
from __future__ import annotations
from datetime import UTC, datetime
import pytest
from controllers.console.app import message as message_module
@@ -120,3 +122,24 @@ def test_suggested_questions_response(app, monkeypatch: pytest.MonkeyPatch) -> N
response = message_module.SuggestedQuestionsResponse(data=["What is AI?", "How does ML work?"])
assert len(response.data) == 2
assert response.data[0] == "What is AI?"
def test_message_detail_response_normalizes_aliases_and_timestamp(app, monkeypatch: pytest.MonkeyPatch) -> None:
"""Test MessageDetailResponse normalizes alias fields and datetime timestamps."""
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
response = message_module.MessageDetailResponse.model_validate(
{
"id": "550e8400-e29b-41d4-a716-446655440000",
"conversation_id": "550e8400-e29b-41d4-a716-446655440001",
"inputs": {"foo": "bar"},
"query": "hello",
"re_sign_file_url_answer": "world",
"from_source": "user",
"status": "normal",
"created_at": created_at,
"message_metadata_dict": {"token_usage": 3},
}
)
assert response.answer == "world"
assert response.metadata == {"token_usage": 3}
assert response.created_at == int(created_at.timestamp())
@@ -258,6 +258,63 @@ def test_restore_published_workflow_to_draft_returns_400_for_invalid_structure(
assert exc.value.description == "invalid workflow graph"
def test_get_published_workflows_marshals_items_before_session_closes(app, monkeypatch: pytest.MonkeyPatch) -> None:
api = workflow_module.PublishedAllWorkflowApi()
handler = _unwrap(api.get)
session_state = {"open": False}
class _SessionContext:
def __enter__(self):
session_state["open"] = True
return object()
def __exit__(self, exc_type, exc, tb):
session_state["open"] = False
return False
class _SessionMaker:
def begin(self):
return _SessionContext()
class _Workflow:
@property
def id(self):
assert session_state["open"] is True
return "w1"
monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object()))
monkeypatch.setattr(workflow_module, "sessionmaker", lambda *_args, **_kwargs: _SessionMaker())
monkeypatch.setattr(workflow_module, "current_account_with_tenant", lambda: (SimpleNamespace(id="u1"), "t1"))
monkeypatch.setattr(
workflow_module,
"WorkflowService",
lambda: SimpleNamespace(
get_all_published_workflow=lambda **_kwargs: ([_Workflow()], False),
),
)
def _fake_marshal(items, fields):
assert session_state["open"] is True
return [{"id": item.id} for item in items]
monkeypatch.setattr(workflow_module, "marshal", _fake_marshal)
with app.test_request_context(
"/apps/app/workflows",
method="GET",
query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"},
):
response = handler(api, app_model=SimpleNamespace(id="app", workflow_id="wf-1"))
assert response == {
"items": [{"id": "w1"}],
"page": 1,
"limit": 10,
"has_more": False,
}
def test_draft_workflow_get_not_found(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
workflow_module, "WorkflowService", lambda: SimpleNamespace(get_draft_workflow=lambda **_k: None)
@@ -0,0 +1,85 @@
from __future__ import annotations
from datetime import UTC, datetime
from graphon.enums import WorkflowExecutionStatus
from controllers.console.app import workflow_app_log as workflow_app_log_module
def test_workflow_app_log_query_parses_bool_and_datetime():
query = workflow_app_log_module.WorkflowAppLogQuery.model_validate(
{
"detail": "true",
"created_at__before": "2026-01-02T03:04:05Z",
"page": "2",
"limit": "10",
}
)
assert query.detail is True
assert query.created_at__before == datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
assert query.page == 2
assert query.limit == 10
def test_workflow_app_log_pagination_response_normalizes_nested_fields():
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
response = workflow_app_log_module.WorkflowAppLogPaginationResponse.model_validate(
{
"page": 1,
"limit": 20,
"total": 1,
"has_more": False,
"data": [
{
"id": "log-1",
"workflow_run": {
"id": "run-1",
"status": WorkflowExecutionStatus.SUCCEEDED,
"created_at": created_at,
"finished_at": created_at,
},
"details": {"trigger_metadata": {}},
"created_by_account": {"id": "acc-1", "name": "acc", "email": "acc@example.com"},
"created_at": created_at,
}
],
}
).model_dump(mode="json")
assert response["data"][0]["workflow_run"]["status"] == "succeeded"
assert response["data"][0]["workflow_run"]["created_at"] == int(created_at.timestamp())
assert response["data"][0]["created_at"] == int(created_at.timestamp())
def test_workflow_archived_log_pagination_response_normalizes_nested_fields():
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
response = workflow_app_log_module.WorkflowArchivedLogPaginationResponse.model_validate(
{
"page": 1,
"limit": 20,
"total": 1,
"has_more": False,
"data": [
{
"id": "archived-1",
"workflow_run": {
"id": "run-1",
"status": WorkflowExecutionStatus.FAILED,
},
"trigger_metadata": {"type": "trigger-plugin"},
"created_by_end_user": {
"id": "eu-1",
"type": "anonymous",
"is_anonymous": True,
"session_id": "session-1",
},
"created_at": created_at,
}
],
}
).model_dump(mode="json")
assert response["data"][0]["workflow_run"]["status"] == "failed"
assert response["data"][0]["created_at"] == int(created_at.timestamp())
@@ -0,0 +1,54 @@
from __future__ import annotations
from datetime import UTC, datetime
from types import SimpleNamespace
from controllers.console.app import workflow_trigger as workflow_trigger_module
def test_parser_models_validate():
parser = workflow_trigger_module.Parser(node_id="node-1")
enable_parser = workflow_trigger_module.ParserEnable(
trigger_id="550e8400-e29b-41d4-a716-446655440000", enable_trigger=True
)
assert parser.node_id == "node-1"
assert enable_parser.enable_trigger is True
def test_workflow_trigger_response_serializes_datetime():
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
trigger = SimpleNamespace(
id="trigger-1",
trigger_type="trigger-plugin",
title="Trigger",
node_id="node-1",
provider_name="provider",
icon="https://example.com/icon",
status="enabled",
created_at=created_at,
updated_at=created_at,
)
payload = workflow_trigger_module.WorkflowTriggerResponse.model_validate(trigger, from_attributes=True).model_dump(
mode="json"
)
assert payload["id"] == "trigger-1"
assert payload["created_at"] == "2026-01-02T03:04:05Z"
assert payload["updated_at"] == "2026-01-02T03:04:05Z"
def test_webhook_trigger_response_serializes_datetime():
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
webhook = {
"id": "webhook-1",
"webhook_id": "whk-1",
"webhook_url": "https://example.com/hook",
"webhook_debug_url": "https://example.com/hook/debug",
"node_id": "node-1",
"created_at": created_at,
}
payload = workflow_trigger_module.WebhookTriggerResponse.model_validate(webhook).model_dump(mode="json")
assert payload["webhook_id"] == "whk-1"
assert payload["created_at"] == "2026-01-02T03:04:05Z"
@@ -99,6 +99,57 @@ class TestHitTestingApi:
assert "records" in result
assert result["records"] == []
def test_hit_testing_success_with_optional_record_fields(self, app, dataset, dataset_id):
api = HitTestingApi()
method = unwrap(api.post)
payload = {
"query": "what is vector search",
}
records = [
{
"segment": None,
"child_chunks": [],
"score": None,
"tsne_position": None,
"files": [],
"summary": None,
}
]
with (
app.test_request_context("/"),
patch.object(
type(console_ns),
"payload",
new_callable=PropertyMock,
return_value=payload,
),
patch.object(
HitTestingPayload,
"model_validate",
return_value=MagicMock(model_dump=lambda **_: payload),
),
patch.object(
HitTestingApi,
"get_and_validate_dataset",
return_value=dataset,
),
patch.object(
HitTestingApi,
"hit_testing_args_check",
),
patch.object(
HitTestingApi,
"perform_hit_testing",
return_value={"query": payload["query"], "records": records},
),
):
result = method(api, dataset_id)
assert result["query"] == payload["query"]
assert result["records"] == records
def test_hit_testing_dataset_not_found(self, app, dataset_id):
api = HitTestingApi()
method = unwrap(api.post)
@@ -1,6 +1,7 @@
from unittest.mock import MagicMock, patch
import controllers.console.explore.recommended_app as module
from models.model import AppMode, IconType
def unwrap(func):
@@ -90,3 +91,48 @@ class TestRecommendedAppApi:
service_mock.assert_called_once_with("11111111-1111-1111-1111-111111111111")
assert result == result_data
class TestRecommendedAppResponseModels:
def test_recommended_app_info_response_computes_icon_url(self):
with patch.object(module, "build_icon_url", return_value="https://signed/icon.png"):
payload = module.RecommendedAppInfoResponse.model_validate(
{
"id": "app-1",
"name": "App",
"mode": AppMode.CHAT,
"icon": "icon.png",
"icon_type": IconType.IMAGE,
"icon_background": "#fff",
}
).model_dump(mode="json")
assert payload["icon_url"] == "https://signed/icon.png"
def test_recommended_app_list_response_serialization(self):
response = module.RecommendedAppListResponse.model_validate(
{
"recommended_apps": [
{
"app": {
"id": "app-1",
"name": "App",
"mode": "chat",
"icon": "icon.png",
"icon_type": "emoji",
"icon_background": "#fff",
},
"app_id": "app-1",
"description": "desc",
"category": "cat",
"position": 1,
"is_listed": True,
"can_trial": False,
}
],
"categories": ["cat"],
}
).model_dump(mode="json")
assert response["recommended_apps"][0]["app_id"] == "app-1"
assert response["categories"] == ["cat"]
@@ -94,7 +94,7 @@ class TestTrialAppWorkflowRunApi:
with app.test_request_context("/"):
with pytest.raises(NotWorkflowAppError):
method(MagicMock(mode=AppMode.CHAT))
method(api, MagicMock(mode=AppMode.CHAT))
def test_success(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -106,7 +106,7 @@ class TestTrialAppWorkflowRunApi:
patch.object(module.AppGenerateService, "generate", return_value=MagicMock()),
patch.object(module.RecommendedAppService, "add_trial_app_record"),
):
result = method(trial_app_workflow)
result = method(api, trial_app_workflow)
assert result is not None
@@ -124,7 +124,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(ProviderNotInitializeError):
method(trial_app_workflow)
method(api, trial_app_workflow)
def test_workflow_quota_exceeded(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -140,7 +140,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(ProviderQuotaExceededError):
method(trial_app_workflow)
method(api, trial_app_workflow)
def test_workflow_model_not_support(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -156,7 +156,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(ProviderModelCurrentlyNotSupportError):
method(trial_app_workflow)
method(api, trial_app_workflow)
def test_workflow_invoke_error(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -172,7 +172,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(CompletionRequestError):
method(trial_app_workflow)
method(api, trial_app_workflow)
def test_workflow_rate_limit_error(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -188,7 +188,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(InvokeRateLimitHttpError):
method(trial_app_workflow)
method(api, trial_app_workflow)
def test_workflow_value_error(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -204,7 +204,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(ValueError):
method(trial_app_workflow)
method(api, trial_app_workflow)
def test_workflow_generic_exception(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowRunApi()
@@ -220,7 +220,7 @@ class TestTrialAppWorkflowRunApi:
),
):
with pytest.raises(InternalServerError):
method(trial_app_workflow)
method(api, trial_app_workflow)
class TestTrialChatApi:
@@ -566,7 +566,7 @@ class TestTrialMessageSuggestedQuestionApi:
with app.test_request_context("/"):
with pytest.raises(NotChatAppError):
method(api, MagicMock(mode="completion"), str(uuid4()))
method(MagicMock(mode="completion"), str(uuid4()))
def test_success(self, app, trial_app_chat, account):
api = module.TrialMessageSuggestedQuestionApi()
@@ -581,7 +581,7 @@ class TestTrialMessageSuggestedQuestionApi:
return_value=["q1", "q2"],
),
):
result = method(api, trial_app_chat, str(uuid4()))
result = method(trial_app_chat, str(uuid4()))
assert result == {"data": ["q1", "q2"]}
@@ -599,7 +599,7 @@ class TestTrialMessageSuggestedQuestionApi:
),
):
with pytest.raises(NotFound):
method(api, trial_app_chat, str(uuid4()))
method(trial_app_chat, str(uuid4()))
class TestTrialAppParameterApi:
@@ -931,7 +931,7 @@ class TestTrialAppWorkflowTaskStopApi:
with app.test_request_context("/"):
with pytest.raises(NotWorkflowAppError):
method(trial_app_chat, str(uuid4()))
method(api, trial_app_chat, str(uuid4()))
def test_success(self, app, trial_app_workflow, account):
api = module.TrialAppWorkflowTaskStopApi()
@@ -944,7 +944,7 @@ class TestTrialAppWorkflowTaskStopApi:
patch.object(module.AppQueueManager, "set_stop_flag_no_user_check") as mock_set_flag,
patch.object(module.GraphEngineManager, "send_stop_command") as mock_send_cmd,
):
result = method(trial_app_workflow, task_id)
result = method(api, trial_app_workflow, task_id)
assert result == {"result": "success"}
mock_set_flag.assert_called_once_with(task_id)
@@ -1,9 +1,11 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, PropertyMock, patch
import pytest
from flask import Flask
from werkzeug.exceptions import Forbidden
import controllers.console.tag.tags as module
from controllers.console import console_ns
from controllers.console.tag.tags import (
TagBindingCreateApi,
@@ -83,13 +85,20 @@ class TestTagListApi:
),
patch(
"controllers.console.tag.tags.TagService.get_tags",
return_value=[{"id": "1", "name": "tag"}],
return_value=[
SimpleNamespace(
id="1",
name="tag",
type=TagType.KNOWLEDGE,
binding_count=1,
)
],
),
):
result, status = method(api)
assert status == 200
assert isinstance(result, list)
assert result == [{"id": "1", "name": "tag", "type": "knowledge", "binding_count": "1"}]
def test_post_success(self, app, admin_user, tag, payload_patch):
api = TagListApi()
@@ -113,6 +122,7 @@ class TestTagListApi:
assert status == 200
assert result["name"] == "test-tag"
assert result["binding_count"] == "0"
def test_post_forbidden(self, app, readonly_user, payload_patch):
api = TagListApi()
@@ -158,7 +168,7 @@ class TestTagUpdateDeleteApi:
result, status = method(api, "tag-1")
assert status == 200
assert result["binding_count"] == 3
assert result["binding_count"] == "3"
def test_patch_forbidden(self, app, readonly_user, payload_patch):
api = TagUpdateDeleteApi()
@@ -277,3 +287,13 @@ class TestTagBindingDeleteApi:
):
with pytest.raises(Forbidden):
method(api)
class TestTagResponseModel:
def test_tag_response_normalizes_enum_type(self):
payload = module.TagResponse.model_validate(
{"id": "tag-1", "name": "tag", "type": TagType.KNOWLEDGE, "binding_count": 1}
).model_dump(mode="json")
assert payload["type"] == "knowledge"
assert payload["binding_count"] == "1"
@@ -11,7 +11,7 @@ from controllers.console.workspace.account import (
ChangeEmailSendEmailApi,
CheckEmailUnique,
)
from models import Account
from models import Account, AccountStatus
from services.account_service import AccountService
@@ -33,7 +33,7 @@ def _build_account(email: str, account_id: str = "acc", tenant: object | None =
account = Account(name=account_id, email=email)
account.email = email
account.id = account_id
account.status = "active"
account.status = AccountStatus.ACTIVE
account._current_tenant = tenant_obj
return account
@@ -18,6 +18,7 @@ from controllers.console.workspace.workspace import (
CustomConfigWorkspaceApi,
SwitchWorkspaceApi,
TenantApi,
TenantInfoResponse,
TenantListApi,
WebappLogoWorkspaceApi,
WorkspaceInfoApi,
@@ -435,6 +436,23 @@ class TestTenantApi:
assert status == 200
class TestTenantInfoResponse:
def test_tenant_info_response_normalizes_enum_and_datetime(self):
created_at = naive_utc_now()
payload = TenantInfoResponse.model_validate(
{
"id": "t1",
"status": TenantStatus.NORMAL,
"plan": CloudPlan.TEAM,
"created_at": created_at,
}
).model_dump(mode="json")
assert payload["status"] == "normal"
assert payload["plan"] == "team"
assert payload["created_at"] == int(created_at.timestamp())
class TestSwitchWorkspaceApi:
def test_switch_success(self, app):
api = SwitchWorkspaceApi()
@@ -18,6 +18,7 @@ from controllers.inner_api.app.dsl import (
InnerAppDSLImportPayload,
_get_active_account,
)
from models.account import AccountStatus
from services.app_dsl_service import ImportStatus
@@ -63,7 +64,7 @@ class TestGetActiveAccount:
@patch("controllers.inner_api.app.dsl.db")
def test_returns_active_account(self, mock_db):
mock_account = MagicMock()
mock_account.status = "active"
mock_account.status = AccountStatus.ACTIVE
mock_db.session.scalar.return_value = mock_account
result = _get_active_account("user@example.com")
@@ -74,7 +75,7 @@ class TestGetActiveAccount:
@patch("controllers.inner_api.app.dsl.db")
def test_returns_none_for_inactive_account(self, mock_db):
mock_account = MagicMock()
mock_account.status = "banned"
mock_account.status = AccountStatus.BANNED
mock_db.session.scalar.return_value = mock_account
result = _get_active_account("banned@example.com")
@@ -102,16 +103,16 @@ class TestEnterpriseAppDSLImport:
@pytest.fixture
def _mock_import_deps(self):
"""Patch db, sessionmaker, and AppDslService for import handler tests."""
mock_session_ctx = MagicMock()
mock_session_ctx.__enter__ = MagicMock(return_value=MagicMock())
mock_session_ctx.__exit__ = MagicMock(return_value=False)
mock_sessionmaker = MagicMock(return_value=MagicMock(begin=MagicMock(return_value=mock_session_ctx)))
"""Patch db, Session, and AppDslService for import handler tests."""
mock_session = MagicMock()
mock_session.__enter__ = MagicMock(return_value=mock_session)
mock_session.__exit__ = MagicMock(return_value=False)
with (
patch("controllers.inner_api.app.dsl.db"),
patch("controllers.inner_api.app.dsl.sessionmaker", mock_sessionmaker),
patch("controllers.inner_api.app.dsl.Session", return_value=mock_session),
patch("controllers.inner_api.app.dsl.AppDslService") as mock_dsl_cls,
):
self._mock_session = mock_session
self._mock_dsl = MagicMock()
mock_dsl_cls.return_value = self._mock_dsl
yield
@@ -147,6 +148,8 @@ class TestEnterpriseAppDSLImport:
assert status_code == 200
assert body["status"] == "completed"
mock_account.set_tenant_id.assert_called_once_with("ws-123")
self._mock_session.commit.assert_called_once_with()
self._mock_session.rollback.assert_not_called()
@pytest.mark.usefixtures("_mock_import_deps")
@patch("controllers.inner_api.app.dsl._get_active_account")
@@ -162,6 +165,8 @@ class TestEnterpriseAppDSLImport:
assert status_code == 202
assert body["status"] == "pending"
self._mock_session.commit.assert_called_once_with()
self._mock_session.rollback.assert_not_called()
@pytest.mark.usefixtures("_mock_import_deps")
@patch("controllers.inner_api.app.dsl._get_active_account")
@@ -177,6 +182,8 @@ class TestEnterpriseAppDSLImport:
assert status_code == 400
assert body["status"] == "failed"
self._mock_session.rollback.assert_called_once_with()
self._mock_session.commit.assert_not_called()
@patch("controllers.inner_api.app.dsl._get_active_account")
def test_import_account_not_found_returns_404(self, mock_get_account, api_instance, app: Flask):
@@ -2,6 +2,7 @@
Unit tests for inner_api plugin decorators
"""
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
@@ -232,11 +233,11 @@ class TestGetUserTenant:
class PluginTestPayload:
"""Simple test payload class"""
def __init__(self, data: dict):
def __init__(self, data: dict[str, Any]):
self.value = data.get("value")
@classmethod
def model_validate(cls, data: dict):
def model_validate(cls, data: dict[str, Any]):
return cls(data)
@@ -277,7 +278,7 @@ class TestPluginData:
# Arrange
class InvalidPayload:
@classmethod
def model_validate(cls, data: dict):
def model_validate(cls, data: dict[str, Any]):
raise Exception("Validation failed")
@plugin_data(payload_type=InvalidPayload)
@@ -20,6 +20,7 @@ from controllers.inner_api.workspace.workspace import (
WorkspaceCreatePayload,
WorkspaceOwnerlessPayload,
)
from models.account import TenantStatus
class TestWorkspaceCreatePayload:
@@ -98,7 +99,7 @@ class TestEnterpriseWorkspace:
mock_tenant.id = "tenant-id"
mock_tenant.name = "My Workspace"
mock_tenant.plan = "sandbox"
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_tenant.created_at = now
mock_tenant.updated_at = now
mock_tenant_svc.create_tenant.return_value = mock_tenant
@@ -162,7 +163,7 @@ class TestEnterpriseWorkspaceNoOwnerEmail:
mock_tenant.name = "My Workspace"
mock_tenant.encrypt_public_key = "pub-key"
mock_tenant.plan = "sandbox"
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_tenant.custom_config = None
mock_tenant.created_at = now
mock_tenant.updated_at = now
@@ -10,6 +10,7 @@ from flask import Flask
from controllers.service_api.app.app import AppInfoApi, AppMetaApi, AppParameterApi
from controllers.service_api.app.error import AppUnavailableError
from models.account import TenantStatus
from models.model import App, AppMode
from tests.unit_tests.conftest import setup_mock_tenant_account_query
@@ -62,7 +63,7 @@ class TestAppParameterApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
# Mock DB queries for app and tenant
mock_db.session.get.side_effect = [
@@ -110,7 +111,7 @@ class TestAppParameterApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app_model,
@@ -151,7 +152,7 @@ class TestAppParameterApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app_model,
@@ -190,7 +191,7 @@ class TestAppParameterApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app_model,
@@ -253,7 +254,7 @@ class TestAppMetaApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app_model,
@@ -321,7 +322,7 @@ class TestAppInfoApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app_model,
@@ -378,7 +379,7 @@ class TestAppInfoApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app,
@@ -424,7 +425,7 @@ class TestAppInfoApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app,
@@ -476,7 +477,7 @@ class TestAppInfoApi:
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = "normal"
mock_tenant.status = TenantStatus.NORMAL
mock_db.session.get.side_effect = [
mock_app,
@@ -15,10 +15,12 @@ Focus on:
import sys
import uuid
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
from graphon.variables.types import SegmentType
from werkzeug.exceptions import BadRequest, NotFound
import services
@@ -29,6 +31,8 @@ from controllers.service_api.app.conversation import (
ConversationRenameApi,
ConversationRenamePayload,
ConversationVariableDetailApi,
ConversationVariableInfiniteScrollPaginationResponse,
ConversationVariableResponse,
ConversationVariablesApi,
ConversationVariablesQuery,
ConversationVariableUpdatePayload,
@@ -261,6 +265,46 @@ class TestConversationVariableUpdatePayload:
assert payload.value == nested
class TestConversationVariableResponseModels:
def test_variable_response_normalizes_value_type_and_timestamps(self):
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
response = ConversationVariableResponse.model_validate(
{
"id": "550e8400-e29b-41d4-a716-446655440000",
"name": "foo",
"value_type": SegmentType.INTEGER,
"value": 1,
"description": "desc",
"created_at": created_at,
"updated_at": created_at,
}
)
assert response.value_type == "number"
assert response.value == "1"
assert response.created_at == int(created_at.timestamp())
assert response.updated_at == int(created_at.timestamp())
def test_variable_pagination_response(self):
response = ConversationVariableInfiniteScrollPaginationResponse.model_validate(
{
"limit": 1,
"has_more": False,
"data": [
{
"id": "550e8400-e29b-41d4-a716-446655440000",
"name": "foo",
"value_type": "string",
"value": "bar",
}
],
}
)
assert response.limit == 1
assert response.has_more is False
assert len(response.data) == 1
assert response.data[0].name == "foo"
class TestConversationAppModeValidation:
"""Test app mode validation for conversation endpoints."""
@@ -549,6 +593,44 @@ class TestConversationVariablesApiController:
with pytest.raises(NotFound):
handler(api, app_model=app_model, end_user=end_user, c_id="00000000-0000-0000-0000-000000000001")
def test_success_serializes_response(self, app, monkeypatch: pytest.MonkeyPatch) -> None:
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
monkeypatch.setattr(
ConversationService,
"get_conversational_variable",
lambda *_args, **_kwargs: SimpleNamespace(
limit=1,
has_more=False,
data=[
{
"id": "550e8400-e29b-41d4-a716-446655440000",
"name": "foo",
"value_type": SegmentType.INTEGER,
"value": 1,
"created_at": created_at,
"updated_at": created_at,
}
],
),
)
api = ConversationVariablesApi()
handler = _unwrap(api.get)
app_model = SimpleNamespace(mode=AppMode.CHAT.value)
end_user = SimpleNamespace()
with app.test_request_context(
"/conversations/1/variables?limit=20",
method="GET",
):
result = handler(api, app_model=app_model, end_user=end_user, c_id="00000000-0000-0000-0000-000000000001")
assert result["limit"] == 1
assert result["has_more"] is False
assert result["data"][0]["value_type"] == "number"
assert result["data"][0]["value"] == "1"
assert result["data"][0]["created_at"] == int(created_at.timestamp())
class TestConversationVariableDetailApiController:
def test_update_type_mismatch(self, app, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -602,3 +684,41 @@ class TestConversationVariableDetailApiController:
c_id="00000000-0000-0000-0000-000000000001",
variable_id="00000000-0000-0000-0000-000000000002",
)
def test_update_success_serializes_response(self, app, monkeypatch: pytest.MonkeyPatch) -> None:
created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
monkeypatch.setattr(
ConversationService,
"update_conversation_variable",
lambda *_args, **_kwargs: {
"id": "550e8400-e29b-41d4-a716-446655440000",
"name": "foo",
"value_type": SegmentType.INTEGER,
"value": 1,
"created_at": created_at,
"updated_at": created_at,
},
)
api = ConversationVariableDetailApi()
handler = _unwrap(api.put)
app_model = SimpleNamespace(mode=AppMode.CHAT.value)
end_user = SimpleNamespace()
with app.test_request_context(
"/conversations/1/variables/2",
method="PUT",
json={"value": 1},
):
result = handler(
api,
app_model=app_model,
end_user=end_user,
c_id="00000000-0000-0000-0000-000000000001",
variable_id="00000000-0000-0000-0000-000000000002",
)
assert result["id"] == "550e8400-e29b-41d4-a716-446655440000"
assert result["value_type"] == "number"
assert result["value"] == "1"
assert result["created_at"] == int(created_at.timestamp())
@@ -15,6 +15,7 @@ Focus on:
import sys
import uuid
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import Mock, patch
@@ -43,6 +44,22 @@ from services.errors.llm import InvokeRateLimitError
from services.workflow_app_service import WorkflowAppService
def _make_mock_workflow_run(run_id: str = "run-1"):
run = Mock()
run.id = run_id
run.workflow_id = "wf-1"
run.status = WorkflowExecutionStatus.SUCCEEDED
run.inputs = {"input": "value"}
run.outputs_dict = {"output": "value"}
run.error = None
run.total_steps = 1
run.total_tokens = 10
run.created_at = datetime(2026, 1, 1, tzinfo=UTC)
run.finished_at = datetime(2026, 1, 1, tzinfo=UTC)
run.elapsed_time = 0.1
return run
class TestWorkflowRunPayload:
"""Test suite for WorkflowRunPayload Pydantic model."""
@@ -359,7 +376,7 @@ class TestWorkflowRunDetailApi:
handler(api, app_model=app_model, workflow_run_id="run")
def test_success(self, monkeypatch: pytest.MonkeyPatch) -> None:
run = SimpleNamespace(id="run")
run = _make_mock_workflow_run(run_id="run")
repo = SimpleNamespace(get_workflow_run_by_id=lambda **_kwargs: run)
workflow_module = sys.modules["controllers.service_api.app.workflow"]
monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object()))
@@ -373,7 +390,10 @@ class TestWorkflowRunDetailApi:
handler = _unwrap(api.get)
app_model = SimpleNamespace(mode=AppMode.WORKFLOW.value, tenant_id="t1", id="a1")
assert handler(api, app_model=app_model, workflow_run_id="run") == run
result = handler(api, app_model=app_model, workflow_run_id="run")
assert result["id"] == "run"
assert result["workflow_id"] == "wf-1"
assert result["status"] == "succeeded"
class TestWorkflowRunApi:
@@ -490,7 +510,7 @@ class TestWorkflowAppLogApi:
monkeypatch.setattr(
WorkflowAppService,
"get_paginate_workflow_app_logs",
lambda *_args, **_kwargs: {"items": [], "total": 0},
lambda *_args, **_kwargs: {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []},
)
api = WorkflowAppLogApi()
@@ -500,7 +520,7 @@ class TestWorkflowAppLogApi:
with app.test_request_context("/workflows/logs", method="GET"):
response = handler(api, app_model=app_model)
assert response == {"items": [], "total": 0}
assert response == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []}
# =============================================================================
@@ -527,9 +547,8 @@ def mock_workflow_app():
class TestWorkflowRunDetailApiGet:
"""Test suite for WorkflowRunDetailApi.get() endpoint.
``get`` is wrapped by ``@validate_app_token`` (preserves ``__wrapped__``)
and ``@service_api_ns.marshal_with``. We call the unwrapped method
directly; ``marshal_with`` is a no-op when calling directly.
``get`` is wrapped by ``@validate_app_token`` (preserves ``__wrapped__``),
and we call the unwrapped method directly in tests.
"""
@patch("controllers.service_api.app.workflow.DifyAPIRepositoryFactory")
@@ -542,9 +561,7 @@ class TestWorkflowRunDetailApiGet:
mock_workflow_app,
):
"""Test successful workflow run detail retrieval."""
mock_run = Mock()
mock_run.id = "run-1"
mock_run.status = "succeeded"
mock_run = _make_mock_workflow_run(run_id="run-1")
mock_repo = Mock()
mock_repo.get_workflow_run_by_id.return_value = mock_run
mock_repo_factory.create_api_workflow_run_repository.return_value = mock_repo
@@ -558,7 +575,8 @@ class TestWorkflowRunDetailApiGet:
api = WorkflowRunDetailApi()
result = _unwrap(api.get)(api, app_model=mock_workflow_app, workflow_run_id=mock_run.id)
assert result == mock_run
assert result["id"] == mock_run.id
assert result["status"] == "succeeded"
@patch("controllers.service_api.app.workflow.db")
def test_get_workflow_run_wrong_app_mode(self, mock_db, app):
@@ -622,8 +640,7 @@ class TestWorkflowTaskStopApiPost:
class TestWorkflowAppLogApiGet:
"""Test suite for WorkflowAppLogApi.get() endpoint.
``get`` is wrapped by ``@validate_app_token`` and
``@service_api_ns.marshal_with``.
``get`` is wrapped by ``@validate_app_token``.
"""
@patch("controllers.service_api.app.workflow.WorkflowAppService")
@@ -637,6 +654,10 @@ class TestWorkflowAppLogApiGet:
):
"""Test successful workflow log retrieval."""
mock_pagination = Mock()
mock_pagination.page = 1
mock_pagination.limit = 20
mock_pagination.total = 0
mock_pagination.has_more = False
mock_pagination.data = []
mock_svc_instance = Mock()
mock_svc_instance.get_paginate_workflow_app_logs.return_value = mock_pagination
@@ -661,4 +682,4 @@ class TestWorkflowAppLogApiGet:
api = WorkflowAppLogApi()
result = _unwrap(api.get)(api, app_model=mock_workflow_app)
assert result == mock_pagination
assert result == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []}
@@ -1,270 +0,0 @@
"""
Unit tests for Service API Site controller
"""
import uuid
from unittest.mock import Mock, patch
import pytest
from werkzeug.exceptions import Forbidden
from controllers.service_api.app.site import AppSiteApi
from models.account import TenantStatus
from models.model import App, Site
from tests.unit_tests.conftest import setup_mock_tenant_account_query
class TestAppSiteApi:
"""Test suite for AppSiteApi"""
@pytest.fixture
def mock_app_model(self):
"""Create a mock App model with tenant."""
app = Mock(spec=App)
app.id = str(uuid.uuid4())
app.tenant_id = str(uuid.uuid4())
app.status = "normal"
app.enable_api = True
mock_tenant = Mock()
mock_tenant.id = app.tenant_id
mock_tenant.status = TenantStatus.NORMAL
app.tenant = mock_tenant
return app
@pytest.fixture
def mock_site(self):
"""Create a mock Site model."""
site = Mock(spec=Site)
site.id = str(uuid.uuid4())
site.app_id = str(uuid.uuid4())
site.title = "Test Site"
site.icon = "icon-url"
site.icon_background = "#ffffff"
site.description = "Site description"
site.copyright = "Copyright 2024"
site.privacy_policy = "Privacy policy text"
site.custom_disclaimer = "Custom disclaimer"
site.default_language = "en-US"
site.prompt_public = True
site.show_workflow_steps = True
site.use_icon_as_answer_icon = False
site.chat_color_theme = "light"
site.chat_color_theme_inverted = False
site.icon_type = "image"
site.created_at = "2024-01-01T00:00:00"
site.updated_at = "2024-01-01T00:00:00"
return site
@patch("controllers.service_api.wraps.user_logged_in")
@patch("controllers.service_api.app.site.db")
@patch("controllers.service_api.wraps.current_app")
@patch("controllers.service_api.wraps.validate_and_get_api_token")
@patch("controllers.service_api.wraps.db")
def test_get_site_success(
self,
mock_wraps_db,
mock_validate_token,
mock_current_app,
mock_db,
mock_user_logged_in,
app,
mock_app_model,
mock_site,
):
"""Test successful retrieval of site configuration."""
# Arrange
mock_current_app.login_manager = Mock()
# Mock authentication
mock_api_token = Mock()
mock_api_token.app_id = mock_app_model.id
mock_api_token.tenant_id = mock_app_model.tenant_id
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = TenantStatus.NORMAL
mock_app_model.tenant = mock_tenant
# Mock wraps.db for authentication
mock_wraps_db.session.get.side_effect = [
mock_app_model,
mock_tenant,
]
mock_account = Mock()
mock_account.current_tenant = mock_tenant
setup_mock_tenant_account_query(mock_wraps_db, mock_tenant, mock_account)
# Mock site.db for site query
mock_db.session.scalar.return_value = mock_site
# Act
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test_token"}):
api = AppSiteApi()
response = api.get()
# Assert
assert response["title"] == "Test Site"
assert response["icon"] == "icon-url"
assert response["description"] == "Site description"
mock_db.session.scalar.assert_called_once()
@patch("controllers.service_api.wraps.user_logged_in")
@patch("controllers.service_api.app.site.db")
@patch("controllers.service_api.wraps.current_app")
@patch("controllers.service_api.wraps.validate_and_get_api_token")
@patch("controllers.service_api.wraps.db")
def test_get_site_not_found(
self,
mock_wraps_db,
mock_validate_token,
mock_current_app,
mock_db,
mock_user_logged_in,
app,
mock_app_model,
):
"""Test that Forbidden is raised when site is not found."""
# Arrange
mock_current_app.login_manager = Mock()
# Mock authentication
mock_api_token = Mock()
mock_api_token.app_id = mock_app_model.id
mock_api_token.tenant_id = mock_app_model.tenant_id
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = TenantStatus.NORMAL
mock_app_model.tenant = mock_tenant
mock_wraps_db.session.get.side_effect = [
mock_app_model,
mock_tenant,
]
mock_account = Mock()
mock_account.current_tenant = mock_tenant
setup_mock_tenant_account_query(mock_wraps_db, mock_tenant, mock_account)
# Mock site query to return None
mock_db.session.scalar.return_value = None
# Act & Assert
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test_token"}):
api = AppSiteApi()
with pytest.raises(Forbidden):
api.get()
@patch("controllers.service_api.wraps.user_logged_in")
@patch("controllers.service_api.app.site.db")
@patch("controllers.service_api.wraps.current_app")
@patch("controllers.service_api.wraps.validate_and_get_api_token")
@patch("controllers.service_api.wraps.db")
def test_get_site_tenant_archived(
self,
mock_wraps_db,
mock_validate_token,
mock_current_app,
mock_db,
mock_user_logged_in,
app,
mock_app_model,
mock_site,
):
"""Test that Forbidden is raised when tenant is archived."""
# Arrange
mock_current_app.login_manager = Mock()
# Mock authentication
mock_api_token = Mock()
mock_api_token.app_id = mock_app_model.id
mock_api_token.tenant_id = mock_app_model.tenant_id
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = TenantStatus.NORMAL
mock_wraps_db.session.get.side_effect = [
mock_app_model,
mock_tenant,
]
mock_account = Mock()
mock_account.current_tenant = mock_tenant
setup_mock_tenant_account_query(mock_wraps_db, mock_tenant, mock_account)
# Mock site query
mock_db.session.scalar.return_value = mock_site
# Set tenant status to archived AFTER authentication
mock_app_model.tenant.status = TenantStatus.ARCHIVE
# Act & Assert
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test_token"}):
api = AppSiteApi()
with pytest.raises(Forbidden):
api.get()
@patch("controllers.service_api.wraps.user_logged_in")
@patch("controllers.service_api.app.site.db")
@patch("controllers.service_api.wraps.current_app")
@patch("controllers.service_api.wraps.validate_and_get_api_token")
@patch("controllers.service_api.wraps.db")
def test_get_site_queries_by_app_id(
self, mock_wraps_db, mock_validate_token, mock_current_app, mock_db, mock_user_logged_in, app, mock_app_model
):
"""Test that site is queried using the app model's id."""
# Arrange
mock_current_app.login_manager = Mock()
# Mock authentication
mock_api_token = Mock()
mock_api_token.app_id = mock_app_model.id
mock_api_token.tenant_id = mock_app_model.tenant_id
mock_validate_token.return_value = mock_api_token
mock_tenant = Mock()
mock_tenant.status = TenantStatus.NORMAL
mock_app_model.tenant = mock_tenant
mock_wraps_db.session.get.side_effect = [
mock_app_model,
mock_tenant,
]
mock_account = Mock()
mock_account.current_tenant = mock_tenant
setup_mock_tenant_account_query(mock_wraps_db, mock_tenant, mock_account)
mock_site = Mock(spec=Site)
mock_site.id = str(uuid.uuid4())
mock_site.app_id = mock_app_model.id
mock_site.title = "Test Site"
mock_site.icon = "icon-url"
mock_site.icon_background = "#ffffff"
mock_site.description = "Site description"
mock_site.copyright = "Copyright 2024"
mock_site.privacy_policy = "Privacy policy text"
mock_site.custom_disclaimer = "Custom disclaimer"
mock_site.default_language = "en-US"
mock_site.prompt_public = True
mock_site.show_workflow_steps = True
mock_site.use_icon_as_answer_icon = False
mock_site.chat_color_theme = "light"
mock_site.chat_color_theme_inverted = False
mock_site.icon_type = "image"
mock_site.created_at = "2024-01-01T00:00:00"
mock_site.updated_at = "2024-01-01T00:00:00"
mock_db.session.scalar.return_value = mock_site
# Act
with app.test_request_context("/site", method="GET", headers={"Authorization": "Bearer test_token"}):
api = AppSiteApi()
api.get()
# Assert
# The query was executed successfully (site returned), which validates the correct query was made
mock_db.session.scalar.assert_called_once()
@@ -1,3 +1,5 @@
from typing import Any
import pytest
from core.extension.extensible import ExtensionModule
@@ -12,10 +14,10 @@ class TestExternalDataTool:
# Create a concrete subclass to test init
class ConcreteTool(ExternalDataTool):
@classmethod
def validate_config(cls, tenant_id: str, config: dict):
def validate_config(cls, tenant_id: str, config: dict[str, Any]):
return super().validate_config(tenant_id, config)
def query(self, inputs: dict, query: str | None = None) -> str:
def query(self, inputs: dict[str, Any], query: str | None = None) -> str:
return super().query(inputs, query)
tool = ConcreteTool(tenant_id="tenant_1", app_id="app_1", variable="var_1", config={"key": "value"})
@@ -28,10 +30,10 @@ class TestExternalDataTool:
# Create a concrete subclass to test init
class ConcreteTool(ExternalDataTool):
@classmethod
def validate_config(cls, tenant_id: str, config: dict):
def validate_config(cls, tenant_id: str, config: dict[str, Any]):
pass
def query(self, inputs: dict, query: str | None = None) -> str:
def query(self, inputs: dict[str, Any], query: str | None = None) -> str:
return ""
tool = ConcreteTool(tenant_id="tenant_1", app_id="app_1", variable="var_1")
@@ -43,10 +45,10 @@ class TestExternalDataTool:
def test_validate_config_raises_not_implemented(self):
class ConcreteTool(ExternalDataTool):
@classmethod
def validate_config(cls, tenant_id: str, config: dict):
def validate_config(cls, tenant_id: str, config: dict[str, Any]):
return super().validate_config(tenant_id, config)
def query(self, inputs: dict, query: str | None = None) -> str:
def query(self, inputs: dict[str, Any], query: str | None = None) -> str:
return ""
with pytest.raises(NotImplementedError):
@@ -55,10 +57,10 @@ class TestExternalDataTool:
def test_query_raises_not_implemented(self):
class ConcreteTool(ExternalDataTool):
@classmethod
def validate_config(cls, tenant_id: str, config: dict):
def validate_config(cls, tenant_id: str, config: dict[str, Any]):
pass
def query(self, inputs: dict, query: str | None = None) -> str:
def query(self, inputs: dict[str, Any], query: str | None = None) -> str:
return super().query(inputs, query)
tool = ConcreteTool(tenant_id="tenant_1", app_id="app_1", variable="var_1")
@@ -10,6 +10,7 @@ This module tests all aspects of the content moderation system including:
- Configuration validation
"""
from typing import Any
from unittest.mock import MagicMock, Mock, patch
import pytest
@@ -28,7 +29,7 @@ class TestKeywordsModeration:
"""Test suite for custom keyword-based content moderation."""
@pytest.fixture
def keywords_config(self) -> dict:
def keywords_config(self) -> dict[str, Any]:
"""
Fixture providing a standard keywords moderation configuration.
@@ -48,7 +49,7 @@ class TestKeywordsModeration:
}
@pytest.fixture
def keywords_moderation(self, keywords_config: dict) -> KeywordsModeration:
def keywords_moderation(self, keywords_config: dict[str, Any]) -> KeywordsModeration:
"""
Fixture providing a KeywordsModeration instance.
@@ -64,7 +65,7 @@ class TestKeywordsModeration:
config=keywords_config,
)
def test_validate_config_success(self, keywords_config: dict):
def test_validate_config_success(self, keywords_config: dict[str, Any]):
"""Test successful validation of keywords moderation configuration."""
# Should not raise any exception
KeywordsModeration.validate_config("test-tenant", keywords_config)
@@ -274,7 +275,7 @@ class TestOpenAIModeration:
"""Test suite for OpenAI-based content moderation."""
@pytest.fixture
def openai_config(self) -> dict:
def openai_config(self) -> dict[str, Any]:
"""
Fixture providing OpenAI moderation configuration.
@@ -293,7 +294,7 @@ class TestOpenAIModeration:
}
@pytest.fixture
def openai_moderation(self, openai_config: dict) -> OpenAIModeration:
def openai_moderation(self, openai_config: dict[str, Any]) -> OpenAIModeration:
"""
Fixture providing an OpenAIModeration instance.
@@ -309,7 +310,7 @@ class TestOpenAIModeration:
config=openai_config,
)
def test_validate_config_success(self, openai_config: dict):
def test_validate_config_success(self, openai_config: dict[str, Any]):
"""Test successful validation of OpenAI moderation configuration."""
# Should not raise any exception
OpenAIModeration.validate_config("test-tenant", openai_config)
@@ -1,5 +1,6 @@
import json
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
import pytest
@@ -57,7 +58,7 @@ class _FakeSelect:
return self
def _dataset_keyword_table(data_source_type: str = "database", keyword_table_dict: dict | None = None):
def _dataset_keyword_table(data_source_type: str = "database", keyword_table_dict: dict[str, Any] | None = None):
return SimpleNamespace(
data_source_type=data_source_type,
keyword_table_dict=keyword_table_dict,
@@ -1,4 +1,5 @@
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, Mock, call, patch
from uuid import uuid4
@@ -20,7 +21,7 @@ def create_mock_document(
doc_id: str,
score: float = 0.8,
provider: str = "dify",
additional_metadata: dict | None = None,
additional_metadata: dict[str, Any] | None = None,
) -> Document:
"""
Create a mock Document object for testing.
@@ -1,4 +1,5 @@
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock, patch
import pytest
@@ -71,7 +72,9 @@ class TestParagraphIndexProcessor:
with pytest.raises(ValueError, match="No rules found in process rule"):
processor.transform([Document(page_content="text", metadata={})], process_rule={"mode": "custom"})
def test_transform_validates_segmentation(self, processor: ParagraphIndexProcessor, process_rule: dict) -> None:
def test_transform_validates_segmentation(
self, processor: ParagraphIndexProcessor, process_rule: dict[str, Any]
) -> None:
rules_without_segmentation = SimpleNamespace(segmentation=None)
with patch(
@@ -84,7 +87,9 @@ class TestParagraphIndexProcessor:
process_rule={"mode": "custom", "rules": {"enabled": True}},
)
def test_transform_builds_split_documents(self, processor: ParagraphIndexProcessor, process_rule: dict) -> None:
def test_transform_builds_split_documents(
self, processor: ParagraphIndexProcessor, process_rule: dict[str, Any]
) -> None:
source_document = Document(page_content="source", metadata={"dataset_id": "dataset-1", "document_id": "doc-1"})
splitter = Mock()
splitter.split_documents.return_value = [

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