From 5c01fed182eed0139397bdcef54b9de2bcbb83da Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Thu, 8 Oct 2026 18:32:10 +0800 Subject: [PATCH 1/4] refactor(api): move Workflow generation and collaboration readers --- api/controllers/console/app/generator.py | 5 +- api/controllers/console/socketio/workflow.py | 2 +- api/core/llm_generator/llm_generator.py | 166 +--------- api/core/workflow/generator/__init__.py | 21 +- .../generator/prompts/builder_prompts.py | 2 +- api/core/workflow/generator/types.py | 4 +- .../application_services/workflow.py | 11 + api/extensions/ext_application_services.py | 4 + .../knowledge/dataset_repository.py | 11 + .../workflow/collaboration_repository.py | 292 ++++++++++++++++++ .../workflow/generation}/runner.py | 26 +- .../workflow/generation/suggestions.py | 180 +++++++++++ .../workflow/generation}/tool_catalogue.py | 16 +- .../workflow_collaboration_service.py | 2 +- api/services/workflow_generator_service.py | 14 +- api/tests/unit_tests/conftest.py | 7 + .../console/app/test_generator_api.py | 54 +++- .../test_llm_generator_missing.py | 65 ++-- .../core/workflow/generator/test_runner.py | 28 +- .../workflow/generator/test_runner_missing.py | 4 +- .../workflow/generator/test_tool_catalogue.py | 57 ++-- .../knowledge/test_pipeline_documents.py | 136 ++++++++ .../test_workflow_collaboration_repository.py | 37 ++- .../test_workflow_collaboration_service.py | 2 +- 24 files changed, 819 insertions(+), 327 deletions(-) create mode 100644 api/repositories/workflow/collaboration_repository.py rename api/{core/workflow/generator => services/workflow/generation}/runner.py (99%) create mode 100644 api/services/workflow/generation/suggestions.py rename api/{core/workflow/generator => services/workflow/generation}/tool_catalogue.py (97%) create mode 100644 api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py diff --git a/api/controllers/console/app/generator.py b/api/controllers/console/app/generator.py index b2cf228c9d6377..ccd18e3a42abc7 100644 --- a/api/controllers/console/app/generator.py +++ b/api/controllers/console/app/generator.py @@ -31,6 +31,7 @@ from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload from core.llm_generator.llm_generator import LLMGenerator from core.workflow.generator.types import WorkflowGenerateErrorCode +from extensions.ext_application_services import application_services from fields.base import ResponseModel from graphon.model_runtime.entities.llm_entities import LLMMode from graphon.model_runtime.errors.invoke import InvokeError @@ -484,7 +485,7 @@ class WorkflowGenerateApi(Resource): """Generate a Workflow / Chatflow draft graph from a natural-language description. Triggered by the cmd+k `/create` slash command. Returns a graph payload - shaped exactly like ``WorkflowService.sync_draft_workflow``'s input, so the + shaped exactly like ``WorkflowDraftService.sync``'s input, so the frontend can hand it straight to ``/apps/{id}/workflows/draft``. """ @@ -559,7 +560,7 @@ class WorkflowInstructionSuggestionsApi(Resource): @with_current_tenant_id @model_validate(WorkflowInstructionSuggestionsPayload) def post(self, req_data: WorkflowInstructionSuggestionsPayload, current_tenant_id: str): - suggestions = LLMGenerator.generate_workflow_instruction_suggestions( + suggestions = application_services().workflow_suggestions.generate_workflow_instruction_suggestions( tenant_id=current_tenant_id, mode=req_data.mode, language=req_data.language, diff --git a/api/controllers/console/socketio/workflow.py b/api/controllers/console/socketio/workflow.py index 13555b1a973ab3..d231c9ac2d16b5 100644 --- a/api/controllers/console/socketio/workflow.py +++ b/api/controllers/console/socketio/workflow.py @@ -9,7 +9,7 @@ from extensions.ext_socketio import sio from libs.passport import PassportService from libs.token import extract_access_token -from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository +from repositories.workflow.collaboration_repository import WorkflowCollaborationRepository from services.workflow_collaboration_service import WorkflowCollaborationService repository = WorkflowCollaborationRepository() diff --git a/api/core/llm_generator/llm_generator.py b/api/core/llm_generator/llm_generator.py index ea066698a29ab6..0275dcd2f41605 100644 --- a/api/core/llm_generator/llm_generator.py +++ b/api/core/llm_generator/llm_generator.py @@ -3,7 +3,7 @@ import re from collections.abc import Sequence from dataclasses import dataclass -from typing import Any, Literal, NotRequired, Protocol, TypedDict, cast +from typing import Any, NotRequired, Protocol, TypedDict, cast import json_repair from sqlalchemy import select @@ -33,7 +33,6 @@ from core.prompt.utils.prompt_template_parser import PromptTemplateParser from core.telemetry import PromptGenerationEvent, TelemetryContext from core.telemetry import emit as telemetry_emit -from extensions.ext_database import db from extensions.ext_storage import storage from graphon.enums import WorkflowNodeExecutionMetadataKey from graphon.model_runtime.entities.llm_entities import LLMResult @@ -121,53 +120,6 @@ def _default_suggested_questions_model_parameters(model_instance: ModelInstance) return parameters -# ── Workflow instruction-suggestion tuning ──────────────────────────────── -# Suggestions are a soft, pre-model-pick enhancement: short, buildable example -# instructions proposed from the tenant's DEFAULT model. Every failure path -# degrades to an empty list, never an error. -_SUGGESTION_MIN_COUNT = 1 -_SUGGESTION_MAX_COUNT = 6 -_SUGGESTION_MAX_TOKENS = 512 -_SUGGESTION_TEMPERATURE = 0.8 -# Bound the grounding context so the prompt stays small regardless of how many -# knowledge bases / tools the tenant has installed. -_SUGGESTION_KB_LIMIT = 10 -_SUGGESTION_TOOL_SAMPLE_LINES = 20 - -_SUGGESTION_SYSTEM_PROMPT = ( - "You help a user start building a Dify app by proposing example build instructions. " - "Each suggestion must be a SHORT (at most 8 words), concrete, and BUILDABLE instruction " - "describing an app to generate for the given app type. Make the suggestions diverse — cover " - "different use cases. When the listed knowledge bases or installed tools fit a suggestion, " - "prefer them, but NEVER invent tools or knowledge bases that are not listed. " - "Reply with ONLY a JSON array of strings and nothing else." -) - - -def _parse_string_list(text: str) -> list[str]: - """Extract a JSON array of strings from a (possibly noisy) LLM response. - - Slices the first ``[...]`` span so surrounding prose / markdown fences are - tolerated, parses it with ``json`` and falls back to ``json_repair``, then - keeps only ``str`` items. Returns ``[]`` on any failure so callers can - treat parsing as best-effort. - """ - match = re.search(r"\[.*\]", text.strip(), re.DOTALL) - if not match: - return [] - raw = match.group(0) - try: - parsed = json.loads(raw) - except Exception: - try: - parsed = json_repair.loads(raw) - except Exception: - return [] - if not isinstance(parsed, list): - return [] - return [item for item in parsed if isinstance(item, str)] - - class WorkflowServiceInterface(Protocol): def get_draft_workflow( self, app_model: App, workflow_id: str | None = None, *, session: Session @@ -427,122 +379,6 @@ def invoke_suggested_questions_after_answer( return questions - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.WORKFLOW_INSTRUCTION_SUGGESTIONS) - def generate_workflow_instruction_suggestions( - cls, - tenant_id: str, - *, - mode: Literal["workflow", "advanced-chat"], - language: str | None = None, - count: int = 4, - ) -> list[str]: - """Propose short, buildable example instructions for the workflow generator. - - Runs BEFORE the user picks a model, so it uses the tenant's DEFAULT LLM - only. Suggestions are a soft enhancement, never a blocker: every failure - path (no default model, KB / tool lookup error, LLM error, unparseable - output) is swallowed and surfaced as an empty list — a valid result the - caller renders as "no suggestions". This method NEVER raises. - """ - count = max(_SUGGESTION_MIN_COUNT, min(count, _SUGGESTION_MAX_COUNT)) - - try: - model_instance = ModelManager.for_tenant(tenant_id=tenant_id).get_default_model_instance( - tenant_id=tenant_id, - model_type=ModelType.LLM, - ) - except Exception: - logger.info("Workflow instruction suggestions: no default model for tenant %s", tenant_id) - return [] - - context_block = cls._build_suggestion_context(tenant_id) - app_type_label = ( - "Workflow — single-shot automation" if mode == "workflow" else "Chatflow — conversational multi-turn" - ) - - user_lines = [ - f"App type: {app_type_label}", - context_block, - f"Return exactly {count} distinct ideas as a JSON array of strings.", - ] - if language: - user_lines.append(f"Write every idea in this language: {language}.") - user_prompt = "\n".join(line for line in user_lines if line) - - prompt_messages: list[PromptMessage] = [ - SystemPromptMessage(content=_SUGGESTION_SYSTEM_PROMPT), - UserPromptMessage(content=user_prompt), - ] - - try: - response: LLMResult = model_instance.invoke_llm( - prompt_messages=prompt_messages, - model_parameters={"max_tokens": _SUGGESTION_MAX_TOKENS, "temperature": _SUGGESTION_TEMPERATURE}, - stream=False, - ) - except Exception: - logger.exception("Workflow instruction suggestions: LLM invocation failed") - return [] - - raw_suggestions = _parse_string_list(response.message.get_text_content() or "") - - # Strip whitespace + surrounding quotes, drop empties, dedupe - # case-insensitively (preserving first-seen casing), cap to ``count``. - cleaned: list[str] = [] - seen: set[str] = set() - for item in raw_suggestions: - idea = item.strip().strip("\"'").strip() - if not idea: - continue - key = idea.casefold() - if key in seen: - continue - seen.add(key) - cleaned.append(idea) - if len(cleaned) >= count: - break - return cleaned - - @staticmethod - def _build_suggestion_context(tenant_id: str) -> str: - """Assemble an optional grounding block naming the tenant's KBs and tools. - - Best-effort: each section is isolated in its own try/except so a failure - enumerating one (DB hiccup, plugin daemon down) never blocks the other - or the suggestion call itself. Returns "" when nothing is available. - """ - sections: list[str] = [] - - try: - from models.dataset import Dataset - - names = db.session.scalars( - select(Dataset.name) - .where(Dataset.tenant_id == tenant_id) - .order_by(Dataset.created_at.desc()) - .limit(_SUGGESTION_KB_LIMIT) - ).all() - kb_names = [name for name in names if name] - if kb_names: - sections.append("Knowledge bases:\n" + "\n".join(f"- {name}" for name in kb_names)) - except Exception: - logger.info("Workflow instruction suggestions: failed to load knowledge bases", exc_info=True) - - try: - from core.workflow.generator.tool_catalogue import build_tool_catalogue, format_tool_catalogue - - tool_text = format_tool_catalogue(build_tool_catalogue(tenant_id)) - if tool_text: - sample = "\n".join(tool_text.splitlines()[:_SUGGESTION_TOOL_SAMPLE_LINES]) - sections.append("Installed tools:\n" + sample) - except Exception: - logger.info("Workflow instruction suggestions: failed to load tool catalogue", exc_info=True) - - if not sections: - return "" - return "\n\n".join(sections) + "\n\n" - @classmethod @with_credit_usage_created_by(CreditUsageCreatedBy.RULE_CONFIG) def generate_rule_config(cls, tenant_id: str, args: RuleGeneratePayload, *, app_id: str | None = None): diff --git a/api/core/workflow/generator/__init__.py b/api/core/workflow/generator/__init__.py index c919ea0c53ca72..ef45ae7cd52de9 100644 --- a/api/core/workflow/generator/__init__.py +++ b/api/core/workflow/generator/__init__.py @@ -1,20 +1 @@ -""" -Workflow generator package. - -Generates a Dify workflow graph (nodes, edges, viewport) from a natural-language -instruction. Intended for the cmd+k `/create` slash command's preview/apply flow. - -Pipeline: - - runner.WorkflowGenerator.generate_workflow_graph(...) - ├── planner_prompts: short LLM call → node and edge plan - ├── node_builder_prompts: bounded parallel calls → semantic node configs - └── postprocess: assemble wrappers, auto-layout, validate graph - -The runner is pure domain logic; ``WorkflowGeneratorService`` (in ``services/``) -owns the model-manager dependency and is what controllers call. -""" - -from .runner import WorkflowGenerator - -__all__ = ["WorkflowGenerator"] +"""Graph generation values, prompts and pure graph transformations.""" diff --git a/api/core/workflow/generator/prompts/builder_prompts.py b/api/core/workflow/generator/prompts/builder_prompts.py index b59f3bbae516b5..6dbc00f49b423e 100644 --- a/api/core/workflow/generator/prompts/builder_prompts.py +++ b/api/core/workflow/generator/prompts/builder_prompts.py @@ -5,7 +5,7 @@ # Each entry mirrors the production ``defaultValue`` from # ``web/app/components/workflow/nodes//default.ts`` so the generated # graph loads in Studio identically to a manually-created node and survives -# both ``WorkflowService.sync_draft_workflow``'s structural checks and the +# both ``WorkflowDraftService.sync``'s structural checks and the # runtime entity validation each node performs when the workflow runs. # # Each snippet mirrors the production node default closely enough for one diff --git a/api/core/workflow/generator/types.py b/api/core/workflow/generator/types.py index 6c7341e16ea9d9..e4d2af1965faaf 100644 --- a/api/core/workflow/generator/types.py +++ b/api/core/workflow/generator/types.py @@ -3,7 +3,7 @@ These TypedDicts describe the planner payload and the runtime graph assembled from builder LLM responses after ``json_repair`` parsing. The graph types mirror -the shape consumed by ``WorkflowService.sync_draft_workflow`` so the output can +the shape consumed by ``WorkflowDraftService.sync`` so the output can be written straight into a draft workflow. """ @@ -145,7 +145,7 @@ class GraphViewportDict(TypedDict): class GraphDict(TypedDict): - """Full graph payload — matches ``WorkflowService.sync_draft_workflow``.""" + """Full graph payload — matches ``WorkflowDraftService.sync``.""" nodes: list[GraphNodeDict] edges: list[GraphEdgeDict] diff --git a/api/extensions/application_services/workflow.py b/api/extensions/application_services/workflow.py index 7a5df690ed603c..8a982e24ecf54f 100644 --- a/api/extensions/application_services/workflow.py +++ b/api/extensions/application_services/workflow.py @@ -111,3 +111,14 @@ def build_workflow_execution_dependencies(database_client: sessionmaker[Session] DifyCoreRepositoryFactory.create_workflow_node_execution_repository, session_factory=database_client ), ) + + +def build_workflow_suggestions(database_client: sessionmaker[Session]): + from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository + from services.workflow.generation.suggestions import WorkflowInstructionSuggestions + from services.workflow.generation.tool_catalogue import build_tool_catalogue, format_tool_catalogue + + return WorkflowInstructionSuggestions( + datasets=SQLAlchemyDatasetRepository(session_factory=database_client), + tools=lambda tenant_id: format_tool_catalogue(build_tool_catalogue(tenant_id)), + ) diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index cb36a69716f901..10763ce96af197 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -36,6 +36,7 @@ build_knowledge_services, ) from extensions.application_services.trial_app import TrialAppServices, build_trial_app_services +from extensions.application_services.workflow import build_workflow_suggestions from extensions.application_services.workspace import ( WorkspaceServices, build_workspace_membership_services, @@ -185,6 +186,7 @@ from services.web_passport_service import WebPassportService from services.webapp_access_adapters import EnterpriseWebAppAccessPolicyGateway from services.webapp_access_query_service import WebAppAccessQueryService, WebAppAccessUnavailableError +from services.workflow.generation.suggestions import WorkflowInstructionSuggestions from services.workflow_app_log_query_service import WorkflowAppLogQueryService from services.workflow_run_service import WorkflowRunService from services.workflow_statistic_query_service import WorkflowStatisticQueryService @@ -238,6 +240,7 @@ class AppScopedEndUserServices: @dataclass(frozen=True, slots=True) class ApplicationServices: + workflow_suggestions: WorkflowInstructionSuggestions agent_apps: AgentAppServices advanced_prompt_templates: AdvancedPromptTemplateService credential_queries: CredentialQuery @@ -491,6 +494,7 @@ def build_application_services( invitation_tokens=invitation_tokens, ) return ApplicationServices( + workflow_suggestions=build_workflow_suggestions(database_client), accounts=account_services, apps=apps, credential_queries=CredentialQueryRepository(session_factory=database_client), diff --git a/api/repositories/knowledge/dataset_repository.py b/api/repositories/knowledge/dataset_repository.py index 447b4d8a4e846e..1de1176f458bb3 100644 --- a/api/repositories/knowledge/dataset_repository.py +++ b/api/repositories/knowledge/dataset_repository.py @@ -45,6 +45,17 @@ class SQLAlchemyDatasetRepository: def __init__(self, *, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory + def names(self, *, tenant_id: str, limit: int) -> list[str]: + with self._session_factory() as session: + return list( + session.scalars( + select(Dataset.name) + .where(Dataset.tenant_id == tenant_id) + .order_by(Dataset.created_at.desc()) + .limit(limit) + ) + ) + def get_estimate_record(self, dataset_ref: DatasetRef) -> DatasetEstimateRecord | None: with self._session_factory() as session: dataset = _get_dataset(session, dataset_ref) diff --git a/api/repositories/workflow/collaboration_repository.py b/api/repositories/workflow/collaboration_repository.py new file mode 100644 index 00000000000000..31bf8c0042e853 --- /dev/null +++ b/api/repositories/workflow/collaboration_repository.py @@ -0,0 +1,292 @@ +from __future__ import annotations + +import json +from typing import Any, NotRequired, TypedDict, override + +from redis.lock import Lock + +from extensions.ext_redis import RedisClientWrapper, redis_client +from extensions.redis_names import serialize_redis_name + +SESSION_STATE_TTL_SECONDS = 3600 +SERVER_HEARTBEAT_TTL_SECONDS = 90 +WORKFLOW_ONLINE_USERS_PREFIX = "workflow_online_users:" +WORKFLOW_LEADER_PREFIX = "workflow_leader:" +WS_SID_MAP_PREFIX = "ws_sid_map:" +WS_SERVER_HEARTBEAT_PREFIX = "ws_server_heartbeat:" +WS_SERVER_SESSIONS_PREFIX = "ws_server_sessions:" +GRAPH_VIEW_STATE_LOCK_PREFIX = "workflow_graph_view_state_lock:" +GRAPH_VIEW_STATE_LOCK_TIMEOUT_SECONDS = 15 + +_UPDATE_SESSION_GRAPH_ACTIVE_LUA = """ +local raw = redis.call('HGET', KEYS[1], ARGV[1]) +if not raw then + return 0 +end + +local decoded_ok, session_info = pcall(cjson.decode, raw) +if not decoded_ok or type(session_info) ~= 'table' then + return 0 +end + +local incoming_sequence = tonumber(ARGV[3]) +local current_sequence = tonumber(session_info.graph_active_sequence) +if current_sequence and incoming_sequence <= current_sequence then + return 0 +end + +session_info.graph_active = ARGV[2] == '1' +session_info.graph_active_sequence = incoming_sequence +redis.call('HSET', KEYS[1], ARGV[1], cjson.encode(session_info)) +return 1 +""" + + +class WorkflowSessionInfo(TypedDict): + user_id: str + username: str + avatar: str | None + sid: str + connected_at: int + server_id: NotRequired[str] + graph_active: NotRequired[bool] + + +class SidMapping(TypedDict): + workflow_id: str + user_id: str + server_id: NotRequired[str] + + +class WorkflowCollaborationRepository: + def __init__(self, *, redis: RedisClientWrapper | None = None) -> None: + self._redis = redis if redis is not None else redis_client + + def online_users(self, app_ids: list[str]) -> dict[str, list[dict[str, Any]]]: + """Batch the existing collaboration hashes without changing their TTLs.""" + results: dict[str, list[dict[str, Any]]] = {} + for start in range(0, len(app_ids), 50): + batch = app_ids[start : start + 50] + pipe = self._redis.pipeline(transaction=False) + for app_id in batch: + pipe.hgetall(self.workflow_key(app_id)) + for app_id, entries in zip(batch, pipe.execute(), strict=True): + users = [] + for raw in entries.values(): + try: + value = json.loads(raw) + except (TypeError, ValueError): + continue + if isinstance(value, dict): + users.append(value) + results[app_id] = users + return results + + @override + def __repr__(self) -> str: + return f"{self.__class__.__name__}(redis_client={self._redis})" + + @staticmethod + def workflow_key(workflow_id: str) -> str: + return f"{WORKFLOW_ONLINE_USERS_PREFIX}{workflow_id}" + + @staticmethod + def leader_key(workflow_id: str) -> str: + return f"{WORKFLOW_LEADER_PREFIX}{workflow_id}" + + @staticmethod + def sid_key(sid: str) -> str: + return f"{WS_SID_MAP_PREFIX}{sid}" + + @staticmethod + def server_key(server_id: str) -> str: + return f"{WS_SERVER_HEARTBEAT_PREFIX}{server_id}" + + @staticmethod + def server_sessions_key(server_id: str) -> str: + return f"{WS_SERVER_SESSIONS_PREFIX}{server_id}" + + @staticmethod + def _decode(value: str | bytes | None) -> str | None: + if value is None: + return None + if isinstance(value, bytes): + return value.decode("utf-8") + return value + + def refresh_session_state(self, workflow_id: str, sid: str) -> None: + workflow_key = self.workflow_key(workflow_id) + sid_key = self.sid_key(sid) + if self._redis.exists(workflow_key): + self._redis.expire(workflow_key, SESSION_STATE_TTL_SECONDS) + if self._redis.exists(sid_key): + self._redis.expire(sid_key, SESSION_STATE_TTL_SECONDS) + + def set_session_info(self, workflow_id: str, session_info: WorkflowSessionInfo) -> None: + workflow_key = self.workflow_key(workflow_id) + sid_mapping: SidMapping = {"workflow_id": workflow_id, "user_id": session_info["user_id"]} + if server_id := session_info.get("server_id"): + sid_mapping["server_id"] = server_id + + self._redis.hset(workflow_key, session_info["sid"], json.dumps(session_info)) + if server_id: + self._redis.hset(self.server_sessions_key(server_id), session_info["sid"], workflow_id) + self._redis.expire(self.server_sessions_key(server_id), SESSION_STATE_TTL_SECONDS) + self._redis.set( + self.sid_key(session_info["sid"]), + json.dumps(sid_mapping), + ex=SESSION_STATE_TTL_SECONDS, + ) + self.refresh_session_state(workflow_id, session_info["sid"]) + + def get_sid_mapping(self, sid: str) -> SidMapping | None: + raw = self._redis.get(self.sid_key(sid)) + if not raw: + return None + value = self._decode(raw) + if not value: + return None + try: + return json.loads(value) + except (TypeError, json.JSONDecodeError): + return None + + def delete_session(self, workflow_id: str, sid: str) -> None: + mapping = self.get_sid_mapping(sid) + if mapping and (server_id := mapping.get("server_id")): + self._redis.hdel(self.server_sessions_key(server_id), sid) + self._redis.hdel(self.workflow_key(workflow_id), sid) + self._redis.delete(self.sid_key(sid)) + + def session_exists(self, workflow_id: str, sid: str) -> bool: + return self._redis.hexists(self.workflow_key(workflow_id), sid) + + def sid_mapping_exists(self, sid: str) -> bool: + return self._redis.exists(self.sid_key(sid)) + + def get_session_sids(self, workflow_id: str) -> list[str]: + raw_sids = self._redis.hkeys(self.workflow_key(workflow_id)) + decoded_sids: list[str] = [] + for sid in raw_sids: + decoded = self._decode(sid) + if decoded: + decoded_sids.append(decoded) + return decoded_sids + + def list_sessions(self, workflow_id: str) -> list[WorkflowSessionInfo]: + sessions_json = self._redis.hgetall(self.workflow_key(workflow_id)) + users: list[WorkflowSessionInfo] = [] + + for session_info_json in sessions_json.values(): + value = self._decode(session_info_json) + if not value: + continue + try: + session_info = json.loads(value) + except (TypeError, json.JSONDecodeError): + continue + + if not isinstance(session_info, dict): + continue + if "user_id" not in session_info or "username" not in session_info or "sid" not in session_info: + continue + + user: WorkflowSessionInfo = { + "user_id": str(session_info["user_id"]), + "username": str(session_info["username"]), + "avatar": session_info.get("avatar"), + "sid": str(session_info["sid"]), + "connected_at": int(session_info.get("connected_at") or 0), + } + if isinstance(session_info.get("graph_active"), bool): + user["graph_active"] = session_info["graph_active"] + users.append(user) + + return users + + def get_session_info(self, workflow_id: str, sid: str) -> WorkflowSessionInfo | None: + raw = self._redis.hget(self.workflow_key(workflow_id), sid) + value = self._decode(raw) + if not value: + return None + try: + session_info = json.loads(value) + except (TypeError, json.JSONDecodeError): + return None + + if not isinstance(session_info, dict): + return None + if "user_id" not in session_info or "username" not in session_info or "sid" not in session_info: + return None + + user: WorkflowSessionInfo = { + "user_id": str(session_info["user_id"]), + "username": str(session_info["username"]), + "avatar": session_info.get("avatar"), + "sid": str(session_info["sid"]), + "connected_at": int(session_info.get("connected_at") or 0), + } + if isinstance(session_info.get("server_id"), str): + user["server_id"] = session_info["server_id"] + if isinstance(session_info.get("graph_active"), bool): + user["graph_active"] = session_info["graph_active"] + return user + + def update_session_graph_active(self, workflow_id: str, sid: str, active: bool, sequence: int) -> bool: + """Atomically apply a graph visibility update when its client sequence is newer.""" + # RedisClientWrapper prefixes regular hash calls, but eval is delegated to the raw client. + workflow_key = serialize_redis_name(self.workflow_key(workflow_id)) + result = self._redis.eval( + _UPDATE_SESSION_GRAPH_ACTIVE_LUA, + 1, + workflow_key, + sid, + "1" if active else "0", + sequence, + ) + return bool(result) + + def graph_view_state_lock(self, workflow_id: str) -> Lock: + """Serialize visibility state changes and their leader-election side effects.""" + return self._redis.lock( + f"{GRAPH_VIEW_STATE_LOCK_PREFIX}{workflow_id}", + timeout=GRAPH_VIEW_STATE_LOCK_TIMEOUT_SECONDS, + ) + + def refresh_server_heartbeat(self, server_id: str) -> None: + self._redis.set(self.server_key(server_id), "1", ex=SERVER_HEARTBEAT_TTL_SECONDS) + + def server_heartbeat_exists(self, server_id: str) -> bool: + return bool(self._redis.exists(self.server_key(server_id))) # tests assert `is True` + + def refresh_server_sessions(self, server_id: str) -> None: + """Refresh Redis TTLs for sessions owned by a live websocket worker.""" + server_sessions_key = self.server_sessions_key(server_id) + sessions = self._redis.hgetall(server_sessions_key) + for raw_sid, raw_workflow_id in sessions.items(): + sid = self._decode(raw_sid) + workflow_id = self._decode(raw_workflow_id) + if not sid or not workflow_id: + continue + if not self.sid_mapping_exists(sid): + self._redis.hdel(server_sessions_key, sid) + continue + self.refresh_session_state(workflow_id, sid) + if sessions: + self._redis.expire(server_sessions_key, SESSION_STATE_TTL_SECONDS) + + def get_current_leader(self, workflow_id: str) -> str | None: + raw = self._redis.get(self.leader_key(workflow_id)) + return self._decode(raw) + + def set_leader_if_absent(self, workflow_id: str, sid: str) -> bool: + return bool(self._redis.set(self.leader_key(workflow_id), sid, nx=True, ex=SESSION_STATE_TTL_SECONDS)) + + def set_leader(self, workflow_id: str, sid: str) -> None: + self._redis.set(self.leader_key(workflow_id), sid, ex=SESSION_STATE_TTL_SECONDS) + + def delete_leader(self, workflow_id: str) -> None: + self._redis.delete(self.leader_key(workflow_id)) + + def expire_leader(self, workflow_id: str) -> None: + self._redis.expire(self.leader_key(workflow_id), SESSION_STATE_TTL_SECONDS) diff --git a/api/core/workflow/generator/runner.py b/api/services/workflow/generation/runner.py similarity index 99% rename from api/core/workflow/generator/runner.py rename to api/services/workflow/generation/runner.py index 64f8b76e417d16..e36493d93de992 100644 --- a/api/core/workflow/generator/runner.py +++ b/api/services/workflow/generation/runner.py @@ -55,18 +55,6 @@ format_tool_catalogue_section, ) from core.workflow.generator.prompts.tool_router_prompts import TOOL_ROUTER_SYSTEM_PROMPT, TOOL_ROUTER_USER_PROMPT -from core.workflow.generator.tool_catalogue import ( - DIRECT_TOOL_INJECTION_LIMIT, - LegacyToolFallbackSelection, - ToolCapabilityQuery, - ToolCatalogueEntry, - find_tool_entry, - format_tool_builder_context, - format_tool_catalogue, - select_legacy_fallback_selection, - select_tool_candidates, - text_mentions_tool_identifier, -) from core.workflow.generator.types import ( GraphDict, GraphViewportDict, @@ -85,6 +73,18 @@ InvokeRateLimitError, InvokeServerUnavailableError, ) +from services.workflow.generation.tool_catalogue import ( + DIRECT_TOOL_INJECTION_LIMIT, + LegacyToolFallbackSelection, + ToolCapabilityQuery, + ToolCatalogueEntry, + find_tool_entry, + format_tool_builder_context, + format_tool_catalogue, + select_legacy_fallback_selection, + select_tool_candidates, + text_mentions_tool_identifier, +) logger = logging.getLogger(__name__) @@ -2463,7 +2463,7 @@ def _validate_structure( installed. Per-node config validation (model spec, prompt template shape, etc.) - is deferred to ``WorkflowService.sync_draft_workflow``; we only fail + is deferred to ``WorkflowDraftService.sync``; we only fail on structural issues the user must know about so they don't get a broken-at-runtime draft. """ diff --git a/api/services/workflow/generation/suggestions.py b/api/services/workflow/generation/suggestions.py new file mode 100644 index 00000000000000..315cfff8f98e07 --- /dev/null +++ b/api/services/workflow/generation/suggestions.py @@ -0,0 +1,180 @@ +"""Suggest workflow instructions with closed dataset reads before remote calls.""" + +import json +import logging +import re +from collections.abc import Callable, Sequence +from typing import Literal, Protocol + +import json_repair + +from core.app.entities.app_invoke_entities import CreditUsageCreatedBy +from core.model_context import with_credit_usage_created_by +from core.model_manager import ModelManager +from graphon.model_runtime.entities.llm_entities import LLMResult +from graphon.model_runtime.entities.message_entities import PromptMessage, SystemPromptMessage, UserPromptMessage +from graphon.model_runtime.entities.model_entities import ModelType + +logger = logging.getLogger(__name__) + + +class SuggestionDatasets(Protocol): + def names(self, *, tenant_id: str, limit: int) -> Sequence[str]: ... + + +# ── Workflow instruction-suggestion tuning ──────────────────────────────── +# Suggestions are a soft, pre-model-pick enhancement: short, buildable example +# instructions proposed from the tenant's DEFAULT model. Every failure path +# degrades to an empty list, never an error. +_SUGGESTION_MIN_COUNT = 1 +_SUGGESTION_MAX_COUNT = 6 +_SUGGESTION_MAX_TOKENS = 512 +_SUGGESTION_TEMPERATURE = 0.8 +# Bound the grounding context so the prompt stays small regardless of how many +# knowledge bases / tools the tenant has installed. +_SUGGESTION_KB_LIMIT = 10 +_SUGGESTION_TOOL_SAMPLE_LINES = 20 + +_SUGGESTION_SYSTEM_PROMPT = ( + "You help a user start building a Dify app by proposing example build instructions. " + "Each suggestion must be a SHORT (at most 8 words), concrete, and BUILDABLE instruction " + "describing an app to generate for the given app type. Make the suggestions diverse — cover " + "different use cases. When the listed knowledge bases or installed tools fit a suggestion, " + "prefer them, but NEVER invent tools or knowledge bases that are not listed. " + "Reply with ONLY a JSON array of strings and nothing else." +) + + +def _parse_string_list(text: str) -> list[str]: + """Extract a JSON array of strings from a (possibly noisy) LLM response. + + Slices the first ``[...]`` span so surrounding prose / markdown fences are + tolerated, parses it with ``json`` and falls back to ``json_repair``, then + keeps only ``str`` items. Returns ``[]`` on any failure so callers can + treat parsing as best-effort. + """ + match = re.search(r"\[.*\]", text.strip(), re.DOTALL) + if not match: + return [] + raw = match.group(0) + try: + parsed = json.loads(raw) + except Exception: + try: + parsed = json_repair.loads(raw) + except Exception: + return [] + if not isinstance(parsed, list): + return [] + return [item for item in parsed if isinstance(item, str)] + + +class WorkflowInstructionSuggestions: + def __init__(self, *, datasets: SuggestionDatasets, tools: Callable[[str], str]) -> None: + self._datasets = datasets + self._tools = tools + + @with_credit_usage_created_by(CreditUsageCreatedBy.WORKFLOW_INSTRUCTION_SUGGESTIONS) + def generate_workflow_instruction_suggestions( + self, + tenant_id: str, + *, + mode: Literal["workflow", "advanced-chat"], + language: str | None = None, + count: int = 4, + ) -> list[str]: + """Propose short, buildable example instructions for the workflow generator. + + Runs BEFORE the user picks a model, so it uses the tenant's DEFAULT LLM + only. Suggestions are a soft enhancement, never a blocker: every failure + path (no default model, KB / tool lookup error, LLM error, unparseable + output) is swallowed and surfaced as an empty list — a valid result the + caller renders as "no suggestions". This method NEVER raises. + """ + count = max(_SUGGESTION_MIN_COUNT, min(count, _SUGGESTION_MAX_COUNT)) + + try: + model_instance = ModelManager.for_tenant(tenant_id=tenant_id).get_default_model_instance( + tenant_id=tenant_id, + model_type=ModelType.LLM, + ) + except Exception: + logger.info("Workflow instruction suggestions: no default model for tenant %s", tenant_id) + return [] + + context_block = self._build_suggestion_context(tenant_id) + app_type_label = ( + "Workflow — single-shot automation" if mode == "workflow" else "Chatflow — conversational multi-turn" + ) + + user_lines = [ + f"App type: {app_type_label}", + context_block, + f"Return exactly {count} distinct ideas as a JSON array of strings.", + ] + if language: + user_lines.append(f"Write every idea in this language: {language}.") + user_prompt = "\n".join(line for line in user_lines if line) + + prompt_messages: list[PromptMessage] = [ + SystemPromptMessage(content=_SUGGESTION_SYSTEM_PROMPT), + UserPromptMessage(content=user_prompt), + ] + + try: + response: LLMResult = model_instance.invoke_llm( + prompt_messages=prompt_messages, + model_parameters={"max_tokens": _SUGGESTION_MAX_TOKENS, "temperature": _SUGGESTION_TEMPERATURE}, + stream=False, + ) + except Exception: + logger.exception("Workflow instruction suggestions: LLM invocation failed") + return [] + + raw_suggestions = _parse_string_list(response.message.get_text_content() or "") + + # Strip whitespace + surrounding quotes, drop empties, dedupe + # case-insensitively (preserving first-seen casing), cap to ``count``. + cleaned: list[str] = [] + seen: set[str] = set() + for item in raw_suggestions: + idea = item.strip().strip("\"'").strip() + if not idea: + continue + key = idea.casefold() + if key in seen: + continue + seen.add(key) + cleaned.append(idea) + if len(cleaned) >= count: + break + return cleaned + + def _build_suggestion_context(self, tenant_id: str) -> str: + """Assemble an optional grounding block naming the tenant's KBs and tools. + + Best-effort: each section is isolated in its own try/except so a failure + enumerating one (DB hiccup, plugin daemon down) never blocks the other + or the suggestion call itself. Returns "" when nothing is available. + """ + sections: list[str] = [] + + try: + names = self._datasets.names(tenant_id=tenant_id, limit=_SUGGESTION_KB_LIMIT) + kb_names = [name for name in names if name] + if kb_names: + sections.append("Knowledge bases:\n" + "\n".join(f"- {name}" for name in kb_names)) + except Exception: + logger.info("Workflow instruction suggestions: failed to load knowledge bases", exc_info=True) + + try: + tool_text = self._tools(tenant_id) + if tool_text: + sample = "\n".join(tool_text.splitlines()[:_SUGGESTION_TOOL_SAMPLE_LINES]) + sections.append("Installed tools:\n" + sample) + except Exception: + logger.info("Workflow instruction suggestions: failed to load tool catalogue", exc_info=True) + + if not sections: + return "" + return "\n\n".join(sections) + "\n\n" diff --git a/api/core/workflow/generator/tool_catalogue.py b/api/services/workflow/generation/tool_catalogue.py similarity index 97% rename from api/core/workflow/generator/tool_catalogue.py rename to api/services/workflow/generation/tool_catalogue.py index 55f5a0ef0a3079..7f5497e7a57148 100644 --- a/api/core/workflow/generator/tool_catalogue.py +++ b/api/services/workflow/generation/tool_catalogue.py @@ -31,7 +31,7 @@ from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolDescription from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.tool_manager import ToolManager +from services.tools.tool_manager import ToolManager logger = logging.getLogger(__name__) @@ -329,20 +329,6 @@ def select_legacy_fallback_selection( ) -def select_legacy_fallback_tools( - entries: list[ToolCatalogueEntry], - *, - explicit_text: str = "", - current_graph: dict[str, Any] | None = None, -) -> list[ToolCatalogueEntry]: - """Keep pinned tools while reproducing the legacy bounded prompt fallback.""" - return select_legacy_fallback_selection( - entries, - explicit_text=explicit_text, - current_graph=current_graph, - ).entries - - def _find_explicit_tool_keys( entries: list[ToolCatalogueEntry], text: str, diff --git a/api/services/workflow_collaboration_service.py b/api/services/workflow_collaboration_service.py index 8e1f0c8a01a162..005b9741bb0a23 100644 --- a/api/services/workflow_collaboration_service.py +++ b/api/services/workflow_collaboration_service.py @@ -16,7 +16,7 @@ from core.rbac import RBACPermission, RBACResourceScope from models.account import Account from models.model import App -from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo +from repositories.workflow.collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo from services.enterprise.rbac_service import RBACService logger = logging.getLogger(__name__) diff --git a/api/services/workflow_generator_service.py b/api/services/workflow_generator_service.py index 66c31e60401ace..78a8eae5025ae0 100644 --- a/api/services/workflow_generator_service.py +++ b/api/services/workflow_generator_service.py @@ -1,7 +1,7 @@ """ Workflow generator service. -Thin facade over ``core.workflow.generator.WorkflowGenerator`` that owns the +Application service over ``services.workflow.generation.runner.WorkflowGenerator`` that owns the model-manager / model-instance plumbing. Controllers call this; the pure domain class never touches the model registry directly. @@ -18,17 +18,17 @@ from core.app.app_config.entities import ModelConfig from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy from core.model_manager import ModelInstance, ModelManager -from core.workflow.generator import WorkflowGenerator -from core.workflow.generator.tool_catalogue import ( - ToolCatalogueEntry, - build_tool_catalogue, - installed_tool_keys, -) from core.workflow.generator.types import ( WorkflowGenerateResultDict, WorkflowGenerationModeRequest, ) from graphon.model_runtime.entities.model_entities import ModelType +from services.workflow.generation.runner import WorkflowGenerator +from services.workflow.generation.tool_catalogue import ( + ToolCatalogueEntry, + build_tool_catalogue, + installed_tool_keys, +) logger = logging.getLogger(__name__) diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index 2ef580442ac54e..20e761e6000c72 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -415,3 +415,10 @@ def workflow_runtime(sqlite_session_factory: sessionmaker[Session]): from extensions.application_services.workflow import build_workflow_execution_dependencies return build_workflow_execution_dependencies(sqlite_session_factory) + + +@pytest.fixture +def workflow_suggestions(sqlite_session_factory: sessionmaker[Session]): + from extensions.application_services.workflow import build_workflow_suggestions + + return build_workflow_suggestions(sqlite_session_factory) diff --git a/api/tests/unit_tests/controllers/console/app/test_generator_api.py b/api/tests/unit_tests/controllers/console/app/test_generator_api.py index 94371295601385..8c93257fdfcc29 100644 --- a/api/tests/unit_tests/controllers/console/app/test_generator_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_generator_api.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from dataclasses import dataclass from inspect import unwrap from unittest.mock import MagicMock @@ -24,6 +25,11 @@ from tests.unit_tests.model_factories import make_workflow +@dataclass +class SuggestionServices: + workflow_suggestions: object + + def _persist_app(session: Session, *, tenant_id: str = "t1") -> App: app_model = App( id="app-1", @@ -561,27 +567,31 @@ def _capture(**kwargs): # ─ /workflow-generate/suggestions ───────────────────────────────────────────── -def test_generate_instruction_suggestions_parses_and_cleans(monkeypatch: pytest.MonkeyPatch) -> None: +def test_generate_instruction_suggestions_parses_and_cleans( + monkeypatch: pytest.MonkeyPatch, workflow_suggestions +) -> None: """Task 1c (i): a mocked default model returning a JSON array is parsed + cleaned.""" - from core.llm_generator import llm_generator as llm_gen_module + from services.workflow.generation import suggestions as llm_gen_module instance = MagicMock() instance.invoke_llm.return_value.message.get_text_content.return_value = '["Summarize a URL", "Translate text"]' mock_manager = MagicMock() mock_manager.for_tenant.return_value.get_default_model_instance.return_value = instance monkeypatch.setattr(llm_gen_module, "ModelManager", mock_manager) - monkeypatch.setattr(llm_gen_module.LLMGenerator, "_build_suggestion_context", staticmethod(lambda _tenant: "")) - - result = llm_gen_module.LLMGenerator.generate_workflow_instruction_suggestions( - tenant_id="t1", mode="workflow", count=4 + monkeypatch.setattr( + llm_gen_module.WorkflowInstructionSuggestions, "_build_suggestion_context", staticmethod(lambda _tenant: "") ) + result = workflow_suggestions.generate_workflow_instruction_suggestions(tenant_id="t1", mode="workflow", count=4) + assert result == ["Summarize a URL", "Translate text"] -def test_generate_instruction_suggestions_dedupes_and_caps(monkeypatch: pytest.MonkeyPatch) -> None: +def test_generate_instruction_suggestions_dedupes_and_caps( + monkeypatch: pytest.MonkeyPatch, workflow_suggestions +) -> None: """Whitespace / surrounding quotes are stripped, case-insensitive dupes dropped, capped to count.""" - from core.llm_generator import llm_generator as llm_gen_module + from services.workflow.generation import suggestions as llm_gen_module instance = MagicMock() instance.invoke_llm.return_value.message.get_text_content.return_value = ( @@ -590,18 +600,22 @@ def test_generate_instruction_suggestions_dedupes_and_caps(monkeypatch: pytest.M mock_manager = MagicMock() mock_manager.for_tenant.return_value.get_default_model_instance.return_value = instance monkeypatch.setattr(llm_gen_module, "ModelManager", mock_manager) - monkeypatch.setattr(llm_gen_module.LLMGenerator, "_build_suggestion_context", staticmethod(lambda _tenant: "")) + monkeypatch.setattr( + llm_gen_module.WorkflowInstructionSuggestions, "_build_suggestion_context", staticmethod(lambda _tenant: "") + ) - result = llm_gen_module.LLMGenerator.generate_workflow_instruction_suggestions( + result = workflow_suggestions.generate_workflow_instruction_suggestions( tenant_id="t1", mode="advanced-chat", count=3 ) assert result == ["Summarize a URL", "Translate text", "Draft an email"] -def test_generate_instruction_suggestions_no_default_model_returns_empty(monkeypatch: pytest.MonkeyPatch) -> None: +def test_generate_instruction_suggestions_no_default_model_returns_empty( + monkeypatch: pytest.MonkeyPatch, workflow_suggestions +) -> None: """Task 1c (ii): a missing default model degrades to an empty list, never raising.""" - from core.llm_generator import llm_generator as llm_gen_module + from services.workflow.generation import suggestions as llm_gen_module mock_manager = MagicMock() mock_manager.for_tenant.return_value.get_default_model_instance.side_effect = ProviderTokenNotInitError( @@ -609,13 +623,16 @@ def test_generate_instruction_suggestions_no_default_model_returns_empty(monkeyp ) monkeypatch.setattr(llm_gen_module, "ModelManager", mock_manager) - result = llm_gen_module.LLMGenerator.generate_workflow_instruction_suggestions(tenant_id="t1", mode="workflow") + result = workflow_suggestions.generate_workflow_instruction_suggestions(tenant_id="t1", mode="workflow") assert result == [] -def test_workflow_instruction_suggestions_route_returns_list(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_workflow_instruction_suggestions_route_returns_list( + app: Flask, monkeypatch: pytest.MonkeyPatch, workflow_suggestions +) -> None: """Task 1c (iii): the route wraps the generator output in {"suggestions": [...]}.""" + monkeypatch.setattr(generator_module, "application_services", lambda: SuggestionServices(workflow_suggestions)) api = generator_module.WorkflowInstructionSuggestionsApi() method = unwrap(api.post) @@ -625,7 +642,7 @@ def _suggest(**kwargs): captured.update(kwargs) return ["Summarize a URL", "Translate text"] - monkeypatch.setattr(generator_module.LLMGenerator, "generate_workflow_instruction_suggestions", _suggest) + monkeypatch.setattr(workflow_suggestions, "generate_workflow_instruction_suggestions", _suggest) with app.test_request_context( "/console/api/workflow-generate/suggestions", @@ -640,13 +657,16 @@ def _suggest(**kwargs): assert captured["count"] == 3 -def test_workflow_instruction_suggestions_route_empty_is_valid_200(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_workflow_instruction_suggestions_route_empty_is_valid_200( + app: Flask, monkeypatch: pytest.MonkeyPatch, workflow_suggestions +) -> None: """Task 1c (iii): an empty list is a valid soft-fail response.""" + monkeypatch.setattr(generator_module, "application_services", lambda: SuggestionServices(workflow_suggestions)) api = generator_module.WorkflowInstructionSuggestionsApi() method = unwrap(api.post) monkeypatch.setattr( - generator_module.LLMGenerator, + workflow_suggestions, "generate_workflow_instruction_suggestions", lambda **_kwargs: [], ) diff --git a/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py index 6fbf18afbbbe87..492b1a28b38d9b 100644 --- a/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py +++ b/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py @@ -6,23 +6,20 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session, sessionmaker -import core.llm_generator.llm_generator as generator_module -from core.llm_generator.llm_generator import LLMGenerator, _parse_string_list from core.model_manager import ModelInstance, ModelManager -from core.workflow.generator import tool_catalogue as tool_catalogue_module -from core.workflow.generator.tool_catalogue import ToolCatalogueEntry from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from models.dataset import Dataset +from services.workflow.generation import tool_catalogue as tool_catalogue_module +from services.workflow.generation.suggestions import _parse_string_list +from services.workflow.generation.tool_catalogue import ToolCatalogueEntry from services.workflow_service import WorkflowService from tests.unit_tests.model_factories import make_dataset @pytest.fixture -def dataset_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind the real SQLite session to the production database extension.""" - - monkeypatch.setattr(generator_module.db, "session", sqlite_session) +def dataset_session(sqlite_session: Session) -> Session: + """Seed the database consumed by the injected suggestion repository.""" return sqlite_session @@ -83,54 +80,53 @@ def test_filter_non_strings(self): class TestGenerateWorkflowInstructionSuggestions: - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_no_default_model(self, mock_for_tenant): + @patch("services.workflow.generation.suggestions.ModelManager.for_tenant") + def test_no_default_model(self, mock_for_tenant, workflow_suggestions): model_manager, _ = _model_manager() model_manager.get_default_model_instance.side_effect = RuntimeError("no default model") mock_for_tenant.return_value = model_manager - assert LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") == [] + assert workflow_suggestions.generate_workflow_instruction_suggestions("tenant", mode="workflow") == [] - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") - def test_llm_success(self, mock_build_context, mock_for_tenant): + @patch("services.workflow.generation.suggestions.ModelManager.for_tenant") + @patch("services.workflow.generation.suggestions.WorkflowInstructionSuggestions._build_suggestion_context") + def test_llm_success(self, mock_build_context, mock_for_tenant, workflow_suggestions): mock_build_context.return_value = "context" model_manager, model_instance = _model_manager() model_instance.invoke_llm.return_value = _llm_result('["idea 1", "idea 2"]') mock_for_tenant.return_value = model_manager - result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") + result = workflow_suggestions.generate_workflow_instruction_suggestions("tenant", mode="workflow") assert result == ["idea 1", "idea 2"] model_instance.invoke_llm.assert_called_once() - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") - def test_llm_error(self, mock_build_context, mock_for_tenant): + @patch("services.workflow.generation.suggestions.ModelManager.for_tenant") + @patch("services.workflow.generation.suggestions.WorkflowInstructionSuggestions._build_suggestion_context") + def test_llm_error(self, mock_build_context, mock_for_tenant, workflow_suggestions): mock_build_context.return_value = "context" model_manager, model_instance = _model_manager() model_instance.invoke_llm.side_effect = RuntimeError("API error") mock_for_tenant.return_value = model_manager - result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") + result = workflow_suggestions.generate_workflow_instruction_suggestions("tenant", mode="workflow") assert result == [] model_instance.invoke_llm.assert_called_once() - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") - def test_llm_bad_output(self, mock_build_context, mock_for_tenant): + @patch("services.workflow.generation.suggestions.ModelManager.for_tenant") + @patch("services.workflow.generation.suggestions.WorkflowInstructionSuggestions._build_suggestion_context") + def test_llm_bad_output(self, mock_build_context, mock_for_tenant, workflow_suggestions): mock_build_context.return_value = "context" model_manager, model_instance = _model_manager() model_instance.invoke_llm.return_value = _llm_result("Not a list") mock_for_tenant.return_value = model_manager - result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") + result = workflow_suggestions.generate_workflow_instruction_suggestions("tenant", mode="workflow") assert result == [] model_instance.invoke_llm.assert_called_once() -@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True) class TestBuildSuggestionContext: - def test_both_success(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch): + def test_both_success(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch, workflow_suggestions): now = datetime.now() dataset_session.add_all( ( @@ -166,30 +162,33 @@ def build_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]: ), ] - # Keep the real module and formatter; only isolate provider/plugin discovery. - monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", build_tool_catalogue) + def tools(tenant_id): + assert workflow_suggestions._datasets._session_factory.kw["bind"].pool.checkedout() == 0 + return tool_catalogue_module.format_tool_catalogue(build_tool_catalogue(tenant_id)) + + monkeypatch.setattr(workflow_suggestions, "_tools", tools) - result = LLMGenerator._build_suggestion_context("tenant") + result = workflow_suggestions._build_suggestion_context("tenant") assert "Knowledge bases:\n- kb1\n- kb2" in result assert ( 'Installed tools:\n- provider/tool1 [provider_id="provider"; tool_name="tool1"] — First tool\n' '- provider/tool2 [provider_id="provider"; tool_name="tool2"] — Second tool' ) in result - def test_both_fail(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch): + def test_both_fail(self, monkeypatch: pytest.MonkeyPatch, workflow_suggestions): def fail_query(_orm_execute_state: object) -> None: raise SQLAlchemyError("DB error") def fail_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]: raise RuntimeError("Tool error") - event.listen(dataset_session, "do_orm_execute", fail_query) - monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", fail_tool_catalogue) + event.listen(workflow_suggestions._datasets._session_factory, "do_orm_execute", fail_query) + monkeypatch.setattr(workflow_suggestions, "_tools", fail_tool_catalogue) try: - assert LLMGenerator._build_suggestion_context("tenant") == "" + assert workflow_suggestions._build_suggestion_context("tenant") == "" finally: - event.remove(dataset_session, "do_orm_execute", fail_query) + event.remove(workflow_suggestions._datasets._session_factory, "do_orm_execute", fail_query) class TestWorkflowServiceInterface: diff --git a/api/tests/unit_tests/core/workflow/generator/test_runner.py b/api/tests/unit_tests/core/workflow/generator/test_runner.py index 89f192fa28647d..cbb24e86e23697 100644 --- a/api/tests/unit_tests/core/workflow/generator/test_runner.py +++ b/api/tests/unit_tests/core/workflow/generator/test_runner.py @@ -17,9 +17,9 @@ import pytest from jinja2 import Template -from core.workflow.generator.runner import WorkflowGenerator, _find_planned_tool_entry -from core.workflow.generator.tool_catalogue import ToolCatalogueEntry from core.workflow.generator.types import GraphDict +from services.workflow.generation.runner import WorkflowGenerator, _find_planned_tool_entry +from services.workflow.generation.tool_catalogue import ToolCatalogueEntry from tests.unit_tests.config_override import apply_config_overrides @@ -254,7 +254,7 @@ def test_router_json_parser_exception_falls_back_to_legacy_prompt(self): model.invoke_llm.return_value = _llm_result("invalid") entries = [_tool_entry("provider", f"tool_{index:03d}") for index in range(100)] - with patch("core.workflow.generator.runner.json_repair.loads", side_effect=ValueError("bad JSON")): + with patch("services.workflow.generation.runner.json_repair.loads", side_effect=ValueError("bad JSON")): text = WorkflowGenerator._resolve_prompt_tool_catalogue( model_instance=model, model_parameters={}, @@ -1263,7 +1263,7 @@ def test_auto_planner_failure_stamps_conversational_default(self): assert result["mode"] == "advanced-chat" def test_auto_with_no_terminal_node_defaults_to_conversational(self): - from core.workflow.generator.runner import _resolve_generation_mode + from services.workflow.generation.runner import _resolve_generation_mode plan = cast(Any, {"title": "x", "description": "x", "nodes": [{"node_type": "llm"}]}) assert _resolve_generation_mode("auto", plan) == "advanced-chat" @@ -1625,7 +1625,7 @@ def test_retries_transient_invoke_error_then_succeeds(self, monkeypatch: pytest. # The planner's first invoke raises a transient connection error; the # retry succeeds and the pipeline completes normally. Sleep is patched # out so the test doesn't actually wait for the backoff. - import core.workflow.generator.runner as _runner_mod + import services.workflow.generation.runner as _runner_mod from graphon.model_runtime.errors.invoke import InvokeConnectionError monkeypatch.setattr(_runner_mod.time, "sleep", lambda _s: None) @@ -1661,7 +1661,7 @@ def test_gives_up_after_exhausting_transient_retries(self, monkeypatch: pytest.M # Every attempt hits the transient error — once we exhaust the retry # budget the failure surfaces as a normal error envelope rather than # hanging or looping forever. - import core.workflow.generator.runner as _runner_mod + import services.workflow.generation.runner as _runner_mod from graphon.model_runtime.errors.invoke import InvokeServerUnavailableError monkeypatch.setattr(_runner_mod.time, "sleep", lambda _s: None) @@ -1681,7 +1681,7 @@ def test_gives_up_after_exhausting_transient_retries(self, monkeypatch: pytest.M assert result["error"] # Bounded: planner attempts == the retry budget, nothing more. - from core.workflow.generator.runner import _INVOKE_MAX_ATTEMPTS + from services.workflow.generation.runner import _INVOKE_MAX_ATTEMPTS assert model_instance.invoke_llm.call_count == _INVOKE_MAX_ATTEMPTS @@ -1690,7 +1690,7 @@ def test_does_not_retry_permanent_invoke_error(self, monkeypatch: pytest.MonkeyP # The runner must fail on the first attempt. # If the code wrongly slept here we'd want the test to still be fast; # patch sleep defensively so a regression can't hang CI. - import core.workflow.generator.runner as _runner_mod + import services.workflow.generation.runner as _runner_mod from graphon.model_runtime.errors.invoke import InvokeAuthorizationError monkeypatch.setattr(_runner_mod.time, "sleep", lambda _s: None) @@ -1722,7 +1722,7 @@ def test_clamp_for_planner_lowers_high_temperature(self): # The planner needs deterministic output — a permissive temperature # would let it ramble. The runner pins it back down for the planner # call while leaving the builder call alone. - from core.workflow.generator.runner import _clamp_for_planner + from services.workflow.generation.runner import _clamp_for_planner out = _clamp_for_planner({"temperature": 0.9, "max_tokens": 1024}) assert out["temperature"] == 0.2 @@ -1733,7 +1733,7 @@ def test_clamp_for_planner_lowers_high_temperature(self): def test_clamp_for_planner_preserves_low_temperature(self): # A user who already picked a tight temperature shouldn't have their # setting overridden — clamping only kicks in above 0.5. - from core.workflow.generator.runner import _clamp_for_planner + from services.workflow.generation.runner import _clamp_for_planner out = _clamp_for_planner({"temperature": 0.3}) assert out["temperature"] == 0.3 @@ -1741,7 +1741,7 @@ def test_clamp_for_planner_preserves_low_temperature(self): def test_clamp_for_planner_injects_default_when_missing(self): # No temperature → planner picks 0.2 so the output stays consistent # across calls. - from core.workflow.generator.runner import _clamp_for_planner + from services.workflow.generation.runner import _clamp_for_planner out = _clamp_for_planner({}) assert out["temperature"] == 0.2 @@ -1751,7 +1751,7 @@ def test_clamp_for_planner_caps_max_tokens_when_missing(self): # pin max_tokens we inject a tight default. This bounds latency/cost # and stops a model that ignores the JSON instruction from rambling # until it hits the provider's (often huge) default ceiling. - from core.workflow.generator.runner import _PLANNER_DEFAULT_MAX_TOKENS, _clamp_for_planner + from services.workflow.generation.runner import _PLANNER_DEFAULT_MAX_TOKENS, _clamp_for_planner out = _clamp_for_planner({}) assert out["max_tokens"] == _PLANNER_DEFAULT_MAX_TOKENS @@ -1759,7 +1759,7 @@ def test_clamp_for_planner_caps_max_tokens_when_missing(self): def test_clamp_for_planner_preserves_caller_max_tokens(self): # A caller who explicitly asked for a budget keeps it — we only fill # the default in when it's absent, never override an intentional value. - from core.workflow.generator.runner import _clamp_for_planner + from services.workflow.generation.runner import _clamp_for_planner out = _clamp_for_planner({"max_tokens": 8192}) assert out["max_tokens"] == 8192 @@ -3019,7 +3019,7 @@ def test_walker_matches_double_brace_placeholders_only(self): # the walker should not treat it as a reference (the LLM-at-runtime # will see the literal single-brace string and ignore it; nothing # for us to validate). - from core.workflow.generator.runner import WorkflowGenerator + from services.workflow.generation.runner import WorkflowGenerator refs: set[tuple[str, str]] = set() WorkflowGenerator._collect_refs_in_data( diff --git a/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py b/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py index 8246fae1d1b16f..1d637a7b2d03bd 100644 --- a/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py +++ b/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py @@ -1,10 +1,10 @@ -from core.workflow.generator.runner import ( +from graphon.model_runtime.errors.invoke import InvokeError +from services.workflow.generation.runner import ( WorkflowGenerator, _stage_error_to_envelope_code, _StageJSONError, _StageSchemaError, ) -from graphon.model_runtime.errors.invoke import InvokeError def test_stage_schema_error(): diff --git a/api/tests/unit_tests/core/workflow/generator/test_tool_catalogue.py b/api/tests/unit_tests/core/workflow/generator/test_tool_catalogue.py index 74043747f5c15c..468cbe2a402188 100644 --- a/api/tests/unit_tests/core/workflow/generator/test_tool_catalogue.py +++ b/api/tests/unit_tests/core/workflow/generator/test_tool_catalogue.py @@ -5,7 +5,7 @@ from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolDescription -from core.workflow.generator.tool_catalogue import ( +from services.workflow.generation.tool_catalogue import ( MAX_ROUTED_TOOL_CANDIDATES, MAX_ROUTED_TOOLS_PER_PROVIDER, ToolCapabilityQuery, @@ -18,7 +18,6 @@ format_tool_catalogue, installed_tool_keys, select_legacy_fallback_selection, - select_legacy_fallback_tools, select_tool_candidates, ) @@ -255,19 +254,6 @@ def test_500_tool_catalogue_never_exceeds_global_candidate_limit(self): assert len(selection.entries) == 15 assert len(selection.entries) <= MAX_ROUTED_TOOL_CANDIDATES - def test_legacy_fallback_keeps_explicit_tool_beyond_first_80(self): - entries = [_entry("provider", f"tool_{index:03d}") for index in range(100)] - - selected = select_legacy_fallback_tools( - entries, - explicit_text="Use provider/tool_099 exactly.", - ) - - assert len(selected) == 80 - assert selected[0]["tool_name"] == "tool_099" - assert any(entry["tool_name"] == "tool_078" for entry in selected) - assert all(entry["tool_name"] != "tool_079" for entry in selected) - class TestLegacyFallbackSelection: def test_returns_selected_and_omitted_tools(self): @@ -281,10 +267,19 @@ def test_returns_selected_and_omitted_tools(self): assert selection.limit == 80 assert selection.overflow_count == 0 - def test_keeps_legacy_fallback_tools_compatible(self): + def test_keeps_explicit_tool_beyond_first_80(self): entries = [_entry("provider", f"tool_{index:03d}") for index in range(100)] - assert select_legacy_fallback_tools(entries) == select_legacy_fallback_selection(entries).entries + selection = select_legacy_fallback_selection( + entries, + explicit_text="Use provider/tool_099 exactly.", + ) + + assert selection.entries == [entries[99], *entries[:79]] + assert selection.omitted_entries == entries[79:99] + assert selection.pinned_count == 1 + assert selection.limit == 80 + assert selection.overflow_count == 0 def test_reports_pinned_overflow_without_fabricated_omissions(self): entries = [_entry("provider", f"tool_{index:03d}") for index in range(85)] @@ -498,15 +493,15 @@ class TestBuildToolCatalogue: without standing up real plugin daemon state. """ - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_returns_empty_list_for_tenant_with_no_tools(self, mock_list, mock_isinstance): mock_list.return_value = iter([]) assert build_tool_catalogue("tenant-1") == [] - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_collects_hardcoded_and_plugin_tools(self, mock_list, mock_isinstance): # Mixed-tenant scenario: hardcoded provider plus a plugin provider, # each carrying one tool. The catalogue must include all four fields @@ -542,8 +537,8 @@ def test_collects_hardcoded_and_plugin_tools(self, mock_list, mock_isinstance): assert time_entry["provider_type"] == "builtin" assert time_entry["plugin_id"] == "" - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_skips_unknown_provider_classes(self, mock_list, mock_isinstance): # If ToolManager ever yields a provider the catalogue doesn't know how # to label, we must continue (not raise) and leave it out of the @@ -556,8 +551,8 @@ def test_skips_unknown_provider_classes(self, mock_list, mock_isinstance): assert [e["provider_name"] for e in entries] == ["time"] - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_continues_when_a_provider_get_tools_raises(self, mock_list, mock_isinstance): # A buggy plugin must not break the whole catalogue. Resilient # per-provider try/except is what keeps generation usable in tenants @@ -570,8 +565,8 @@ def test_continues_when_a_provider_get_tools_raises(self, mock_list, mock_isinst assert [e["provider_name"] for e in entries] == ["time"] - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_skips_individual_tools_when_their_metadata_is_broken(self, mock_list, mock_isinstance): # Per-tool try/except — a single mis-declared tool inside an otherwise # healthy provider gets dropped, the rest still surface. @@ -585,8 +580,8 @@ def test_skips_individual_tools_when_their_metadata_is_broken(self, mock_list, m assert [e["tool_name"] for e in entries] == ["ok"] - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_keeps_complete_inventory_for_validation_beyond_prompt_cap(self, mock_list, mock_isinstance): # Prompt formatting is capped separately. Dropping entries here would # make the validator falsely report installed tools after the cap as @@ -602,8 +597,8 @@ def test_keeps_complete_inventory_for_validation_beyond_prompt_cap(self, mock_li assert len(entries) == 200 assert ("p", "t199") in installed_tool_keys(entries) - @patch("core.workflow.generator.tool_catalogue.isinstance", side_effect=_patched_isinstance) - @patch("core.workflow.generator.tool_catalogue.ToolManager.list_builtin_providers") + @patch("services.workflow.generation.tool_catalogue.isinstance", side_effect=_patched_isinstance) + @patch("services.workflow.generation.tool_catalogue.ToolManager.list_builtin_providers") def test_defaults_plugin_id_to_empty_string_when_missing(self, mock_list, mock_isinstance): # Plugin provider whose plugin_id is None should serialise to "" so # the consumer can safely index ``e["plugin_id"]`` without a None diff --git a/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py b/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py new file mode 100644 index 00000000000000..80534595db45dc --- /dev/null +++ b/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py @@ -0,0 +1,136 @@ +from collections.abc import Iterator + +import pytest +from sqlalchemy import Connection, Engine, event, inspect, select, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session, sessionmaker + +from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipeline +from models.enums import IndexingStatus +from models.pipeline_execution import PipelineDocumentSeed +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from tests.unit_tests.repositories.knowledge.test_document_repository import _dataset, _document + + +@pytest.fixture +def documents(sqlite_session_factory: sessionmaker[Session]) -> SQLAlchemyDocumentRepository: + with sqlite_session_factory.begin() as session: + dataset = _dataset("dataset-1", "tenant-1") + dataset.pipeline_id = "pipeline-1" + dataset.chunk_structure = "text_model" + session.add(dataset) + pipeline = Pipeline(tenant_id="tenant-1", name="Pipeline", description="") + pipeline.id = "pipeline-1" + session.add(pipeline) + return SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + + +def _prepare(documents: SQLAlchemyDocumentRepository, original_document_id: str | None = None) -> list[Document]: + return documents.prepare_pipeline_documents( + tenant_id="tenant-1", + pipeline_id="pipeline-1", + dataset_id="dataset-1", + user_id="account-1", + batch="batch-1", + start_node_id="source-1", + inputs={"query": "test"}, + seeds=[ + PipelineDocumentSeed("first", "upload_file", {"name": "first"}, {}), + PipelineDocumentSeed("second", "upload_file", {"name": "second"}, {}), + ], + original_document_id=original_document_id, + ) + + +@pytest.fixture +def active_transactions(sqlite_engine: Engine) -> Iterator[set[Connection]]: + active: set[Connection] = set() + begin, finish = active.add, active.discard + event.listen(sqlite_engine, "begin", begin) + event.listen(sqlite_engine, "commit", finish) + event.listen(sqlite_engine, "rollback", finish) + try: + yield active + finally: + event.remove(sqlite_engine, "begin", begin) + event.remove(sqlite_engine, "commit", finish) + event.remove(sqlite_engine, "rollback", finish) + + +def test_preparation_commits_documents_and_logs_then_releases_connection( + documents: SQLAlchemyDocumentRepository, + sqlite_session_factory: sessionmaker[Session], + active_transactions: set[Connection], +): + rows = _prepare(documents) + assert not active_transactions + assert all(inspect(row).detached for row in rows) + assert [row.position for row in rows] == [1, 2] + with sqlite_session_factory() as reader: + assert {row.id for row in reader.scalars(select(Document))} == {row.id for row in rows} + logs = list(reader.scalars(select(DocumentPipelineExecutionLog))) + assert {log.document_id for log in logs} == {row.id for row in rows} + assert all(log.input_data == {"query": "test"} for log in logs) + assert not active_transactions + + +@pytest.mark.parametrize("changed_owner", ["dataset", "pipeline", "document_tenant", "document_dataset"]) +def test_preparation_rejects_another_owner_without_changing_document( + documents: SQLAlchemyDocumentRepository, + sqlite_session_factory: sessionmaker[Session], + changed_owner: str, +): + with sqlite_session_factory.begin() as session: + original = _document("original", workspace_id="tenant-1", dataset_id="dataset-1") + original.indexing_status = IndexingStatus.COMPLETED + if changed_owner == "document_tenant": + original.tenant_id = "other-tenant" + elif changed_owner == "document_dataset": + original.dataset_id = "other-dataset" + else: + row = session.get(Dataset if changed_owner == "dataset" else Pipeline, f"{changed_owner}-1") + assert row is not None + row.tenant_id = "other-tenant" + session.add(original) + with pytest.raises(ValueError, match="Pipeline (dataset is required|document not found)"): + _prepare(documents, "original") + with sqlite_session_factory() as reader: + original = reader.get(Document, "original") + assert original is not None + assert original.indexing_status == IndexingStatus.COMPLETED + assert reader.scalars(select(DocumentPipelineExecutionLog)).all() == [] + + +def test_preparation_updates_original_and_logs_in_one_transaction( + documents: SQLAlchemyDocumentRepository, sqlite_session_factory: sessionmaker[Session] +): + with sqlite_session_factory.begin() as session: + session.add(_document("original", workspace_id="tenant-1", dataset_id="dataset-1")) + assert _prepare(documents, "original") == [] + with sqlite_session_factory() as reader: + original = reader.get(Document, "original") + assert original is not None + assert original.indexing_status == IndexingStatus.WAITING + assert reader.scalars(select(DocumentPipelineExecutionLog.document_id)).all() == ["original", "original"] + + +def test_log_failure_rolls_back_documents_and_original_status( + documents: SQLAlchemyDocumentRepository, sqlite_session_factory: sessionmaker[Session] +): + with sqlite_session_factory.begin() as session: + session.add(_document("original", workspace_id="tenant-1", dataset_id="dataset-1")) + session.execute( + text( + "CREATE TRIGGER reject_pipeline_log BEFORE INSERT ON document_pipeline_execution_logs " + "BEGIN SELECT RAISE(ABORT, 'log write failed'); END" + ) + ) + for original_id in (None, "original"): + with pytest.raises(IntegrityError, match="log write failed"): + _prepare(documents, original_id) + with sqlite_session_factory() as reader: + assert reader.scalars(select(Document.id)).all() == ["original"] + original = reader.get(Document, "original") + assert original is not None + assert original.indexing_status == IndexingStatus.COMPLETED + assert reader.scalars(select(DocumentPipelineExecutionLog)).all() == [] diff --git a/api/tests/unit_tests/repositories/test_workflow_collaboration_repository.py b/api/tests/unit_tests/repositories/test_workflow_collaboration_repository.py index 6908f63c98a3e6..3f7c46afb7594b 100644 --- a/api/tests/unit_tests/repositories/test_workflow_collaboration_repository.py +++ b/api/tests/unit_tests/repositories/test_workflow_collaboration_repository.py @@ -1,10 +1,12 @@ import json +from typing import cast from unittest.mock import Mock import pytest -from repositories import workflow_collaboration_repository as repo_module -from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository +from extensions.ext_redis import RedisClientWrapper +from repositories.workflow import collaboration_repository as repo_module +from repositories.workflow.collaboration_repository import WorkflowCollaborationRepository class TestWorkflowCollaborationRepository: @@ -269,3 +271,34 @@ def test_refresh_server_sessions_drops_stale_sid(self, mock_redis: Mock) -> None repository.refresh_server_sessions("server-1") mock_redis.hdel.assert_called_once_with("ws_server_sessions:server-1", "sid-stale") + + +def test_online_users_batches_existing_hashes_and_ignores_corrupt_values() -> None: + pipelines = [] + + class Pipeline: + def __init__(self) -> None: + self.keys = [] + + def hgetall(self, key: str) -> None: + self.keys.append(key) + + def execute(self) -> list[dict[bytes, bytes]]: + return [ + {b"good": b'{"user_id":"u","username":"User"}', b"corrupt": b"bad", b"array": b"[]"} for _ in self.keys + ] + + class Redis: + def pipeline(self, *, transaction: bool) -> Pipeline: + assert transaction is False + pipe = Pipeline() + pipelines.append(pipe) + return pipe + + repository = WorkflowCollaborationRepository(redis=cast(RedisClientWrapper, Redis())) + app_ids = [str(index) for index in range(51)] + result = repository.online_users(app_ids) + assert [len(pipe.keys) for pipe in pipelines] == [50, 1] + assert pipelines[0].keys[0] == "workflow_online_users:0" + assert list(result) == app_ids + assert result["50"] == [{"user_id": "u", "username": "User"}] diff --git a/api/tests/unit_tests/services/test_workflow_collaboration_service.py b/api/tests/unit_tests/services/test_workflow_collaboration_service.py index c1cbe203d63b12..87e0fa4296c726 100644 --- a/api/tests/unit_tests/services/test_workflow_collaboration_service.py +++ b/api/tests/unit_tests/services/test_workflow_collaboration_service.py @@ -12,7 +12,7 @@ from models.account import Account, Tenant from models.base import TypeBase from models.model import App, AppMode -from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository +from repositories.workflow.collaboration_repository import WorkflowCollaborationRepository from services.workflow_collaboration_service import SYNC_REQUEST_TIMEOUT_SECONDS, WorkflowCollaborationService from tests.unit_tests.config_override import config_overrides_context from tests.unit_tests.model_factories import make_app From 59377952b7c4cde665d0de87bb0e48241fff20d0 Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Fri, 9 Oct 2026 12:33:33 +0800 Subject: [PATCH 2/4] test(workflow): inject current generation and collaboration dependencies --- api/tests/unit_tests/conftest.py | 3 ++- .../core/llm_generator/test_llm_generator.py | 15 ++++--------- .../knowledge/test_pipeline_documents.py | 8 +++---- .../test_workflow_collaboration_service.py | 22 ++++++++++--------- .../workflow_collaboration_fixtures.py | 2 +- 5 files changed, 23 insertions(+), 27 deletions(-) diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index a12ab4d3577467..6732798550dc49 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -64,6 +64,7 @@ from services.tag_application_service import TagApplicationService from services.tools.workflow_tools_manage_service import WorkflowToolManageService from services.workflow.console_variable_service import ConsoleWorkflowVariableService + from services.workflow.generation.suggestions import WorkflowInstructionSuggestions from services.workflow.variable_service import WorkflowVariableService @@ -451,7 +452,7 @@ def workflow_runtime(sqlite_session_factory: sessionmaker[Session]) -> WorkflowE @pytest.fixture -def workflow_suggestions(sqlite_session_factory: sessionmaker[Session]): +def workflow_suggestions(sqlite_session_factory: sessionmaker[Session]) -> WorkflowInstructionSuggestions: from extensions.application_services.workflow import build_workflow_suggestions return build_workflow_suggestions(sqlite_session_factory) diff --git a/api/tests/unit_tests/core/llm_generator/test_llm_generator.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py index c744d43335610c..4bb8008a8b8b7e 100644 --- a/api/tests/unit_tests/core/llm_generator/test_llm_generator.py +++ b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py @@ -1,7 +1,6 @@ """Tests for LLM generation and database-backed instruction modification.""" import json -from collections.abc import Iterator from datetime import datetime from decimal import Decimal from types import SimpleNamespace @@ -9,7 +8,7 @@ from uuid import uuid4 import pytest -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session, sessionmaker from core.app.app_config.entities import ModelConfig from core.app.entities.app_invoke_entities import CreditUsageCreatedBy @@ -37,15 +36,9 @@ @pytest.fixture -def database(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: - """Bind the shared SQLite session to the generator's scoped-session interface.""" - - registry = scoped_session(lambda: sqlite_session) - monkeypatch.setattr(llm_generator_module.db, "session", registry) - try: - yield sqlite_session - finally: - registry.remove() +def database(sqlite_session: Session) -> Session: + """Use the real SQLite session passed explicitly to the generator.""" + return sqlite_session @pytest.fixture diff --git a/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py b/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py index 80534595db45dc..331998a2ab584b 100644 --- a/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py +++ b/api/tests/unit_tests/repositories/knowledge/test_pipeline_documents.py @@ -61,7 +61,7 @@ def test_preparation_commits_documents_and_logs_then_releases_connection( documents: SQLAlchemyDocumentRepository, sqlite_session_factory: sessionmaker[Session], active_transactions: set[Connection], -): +) -> None: rows = _prepare(documents) assert not active_transactions assert all(inspect(row).detached for row in rows) @@ -79,7 +79,7 @@ def test_preparation_rejects_another_owner_without_changing_document( documents: SQLAlchemyDocumentRepository, sqlite_session_factory: sessionmaker[Session], changed_owner: str, -): +) -> None: with sqlite_session_factory.begin() as session: original = _document("original", workspace_id="tenant-1", dataset_id="dataset-1") original.indexing_status = IndexingStatus.COMPLETED @@ -103,7 +103,7 @@ def test_preparation_rejects_another_owner_without_changing_document( def test_preparation_updates_original_and_logs_in_one_transaction( documents: SQLAlchemyDocumentRepository, sqlite_session_factory: sessionmaker[Session] -): +) -> None: with sqlite_session_factory.begin() as session: session.add(_document("original", workspace_id="tenant-1", dataset_id="dataset-1")) assert _prepare(documents, "original") == [] @@ -116,7 +116,7 @@ def test_preparation_updates_original_and_logs_in_one_transaction( def test_log_failure_rolls_back_documents_and_original_status( documents: SQLAlchemyDocumentRepository, sqlite_session_factory: sessionmaker[Session] -): +) -> None: with sqlite_session_factory.begin() as session: session.add(_document("original", workspace_id="tenant-1", dataset_id="dataset-1")) session.execute( diff --git a/api/tests/unit_tests/services/test_workflow_collaboration_service.py b/api/tests/unit_tests/services/test_workflow_collaboration_service.py index 39b132794d62d1..8a1a964c0d9050 100644 --- a/api/tests/unit_tests/services/test_workflow_collaboration_service.py +++ b/api/tests/unit_tests/services/test_workflow_collaboration_service.py @@ -1,11 +1,11 @@ import logging -from collections.abc import Iterator +from collections.abc import Generator, Iterator from contextlib import contextmanager from unittest.mock import MagicMock, Mock, patch import pytest from socketio.exceptions import TimeoutError as SocketIOTimeoutError -from sqlalchemy import Engine +from sqlalchemy import Engine, Table from sqlalchemy.orm import Session from core.rbac import RBACPermission, RBACResourceScope @@ -13,7 +13,7 @@ from models.account import Account, Tenant from models.base import TypeBase from models.model import App, AppMode -from repositories.workflow.collaboration_repository import WorkflowCollaborationRepository +from repositories.workflow.collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo from services.workflow_collaboration_service import SYNC_REQUEST_TIMEOUT_SECONDS, WorkflowCollaborationService from tests.unit_tests.config_override import config_overrides_context from tests.unit_tests.model_factories import make_app @@ -24,18 +24,17 @@ @pytest.fixture def collaboration_redis( - redis_transport: tuple[RedisClientWrapper, MagicMock], monkeypatch: pytest.MonkeyPatch + redis_transport: tuple[RedisClientWrapper, MagicMock], ) -> RedisClientWrapper: client, commands = redis_transport commands.side_effect = RedisState().execute - monkeypatch.setattr("repositories.workflow_collaboration_repository.redis_client", client) return client @pytest.fixture def real_service(collaboration_redis: RedisClientWrapper) -> ServiceFixture: assert collaboration_redis._require_client() is not None - repository = WorkflowCollaborationRepository() + repository = WorkflowCollaborationRepository(redis=collaboration_redis) socketio = Mock() return WorkflowCollaborationService(repository, socketio, server_id="server-1"), repository, socketio @@ -44,7 +43,9 @@ def real_service(collaboration_redis: RedisClientWrapper) -> ServiceFixture: def db_session(sqlite_engine: Engine) -> Iterator[Session]: """Provide a real session for tenant-scoped workflow app access checks.""" - TypeBase.metadata.create_all(sqlite_engine, tables=[App.__table__]) + app_table = App.__table__ + assert isinstance(app_table, Table) + TypeBase.metadata.create_all(sqlite_engine, tables=[app_table]) with Session(sqlite_engine, expire_on_commit=False) as session: yield session @@ -112,7 +113,8 @@ def test_authorize_and_join_workflow_room_returns_none_when_missing_user( ) -> None: # Arrange collaboration_service, _repository, socketio = real_service - socketio.get_session.return_value = {} + empty_socket_session: dict[str, object] = {} + socketio.get_session.return_value = empty_socket_session # Act result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=db_session) @@ -704,7 +706,7 @@ def test_graph_view_state_serializes_visibility_write_with_leader_demotion( events: list[str] = [] @contextmanager - def graph_view_lock() -> Iterator[None]: + def graph_view_lock() -> Generator[None, None, None]: events.append("lock_enter") yield events.append("lock_exit") @@ -1354,7 +1356,7 @@ def test_is_session_active_rejects_remote_session_with_dead_server(self, real_se assert collaboration_service.is_session_active("wf-1", "sid-remote") is False @staticmethod - def _session(sid: str, connected_at: int, graph_active: bool) -> dict: + def _session(sid: str, connected_at: int, graph_active: bool) -> WorkflowSessionInfo: return { "user_id": f"u-{sid}", "username": sid, diff --git a/api/tests/unit_tests/services/workflow_collaboration_fixtures.py b/api/tests/unit_tests/services/workflow_collaboration_fixtures.py index 5d3f9bee49b7a3..099cd8002d97ce 100644 --- a/api/tests/unit_tests/services/workflow_collaboration_fixtures.py +++ b/api/tests/unit_tests/services/workflow_collaboration_fixtures.py @@ -6,7 +6,7 @@ from redis.lock import Lock -from repositories.workflow_collaboration_repository import ( +from repositories.workflow.collaboration_repository import ( _UPDATE_SESSION_GRAPH_ACTIVE_LUA, WorkflowCollaborationRepository, WorkflowSessionInfo, From 605801e41bed2e12397c17a4a946196efc27cdf4 Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Mon, 26 Oct 2026 20:13:36 +0800 Subject: [PATCH 3/4] Point generator service tests at the extracted workflow runner owner --- .../unit_tests/services/test_workflow_generator_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/tests/unit_tests/services/test_workflow_generator_service.py b/api/tests/unit_tests/services/test_workflow_generator_service.py index 14ef9a2e6ffd73..4f66f39cdf3446 100644 --- a/api/tests/unit_tests/services/test_workflow_generator_service.py +++ b/api/tests/unit_tests/services/test_workflow_generator_service.py @@ -30,12 +30,12 @@ ) from core.tools.plugin_tool.provider import PluginToolProviderController from core.tools.tool_manager import ToolManager -from core.workflow.generator import WorkflowGenerator from extensions.ext_redis import RedisClientWrapper from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResultChunk, LLMResultChunkDelta from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from services.workflow.generation.runner import WorkflowGenerator from services.workflow_generator_service import WorkflowGeneratorService from tests.helpers.unit.core.model_fixtures import make_model_config From 14c3fe0559a46ccf932e3e66498097f382b8cb3a Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Mon, 26 Oct 2026 20:15:51 +0800 Subject: [PATCH 4/4] Patch the active service tool manager in generator wiring tests --- .../unit_tests/services/test_workflow_generator_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/tests/unit_tests/services/test_workflow_generator_service.py b/api/tests/unit_tests/services/test_workflow_generator_service.py index 4f66f39cdf3446..c84644886a84ac 100644 --- a/api/tests/unit_tests/services/test_workflow_generator_service.py +++ b/api/tests/unit_tests/services/test_workflow_generator_service.py @@ -29,12 +29,12 @@ ToolProviderIdentity, ) from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.tool_manager import ToolManager from extensions.ext_redis import RedisClientWrapper from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResultChunk, LLMResultChunkDelta from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from services.tools.tool_manager import ToolManager from services.workflow.generation.runner import WorkflowGenerator from services.workflow_generator_service import WorkflowGeneratorService from tests.helpers.unit.core.model_fixtures import make_model_config