diff --git a/api/controllers/console/app/generator.py b/api/controllers/console/app/generator.py index 89bfbabe106fe0..ad7cd9b789c5cb 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 @@ -533,7 +534,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``. """ @@ -614,7 +615,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 dcec1b65b541b0..53d394f47a2c10 100644 --- a/api/core/llm_generator/llm_generator.py +++ b/api/core/llm_generator/llm_generator.py @@ -2,7 +2,7 @@ import logging import re from collections.abc import Sequence -from typing import Any, Literal, Protocol, TypedDict, cast +from typing import Any, Protocol, TypedDict, cast import json_repair from sqlalchemy import select @@ -30,7 +30,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 @@ -64,53 +63,6 @@ def _normalize_completion_params(completion_params: dict[str, object]) -> tuple[ return normalized_parameters, stop -# ── 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, *, session: Session) -> Workflow | None: pass @@ -250,122 +202,6 @@ def generate_conversation_name( return name - @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 39c6a2f84dbc1e..985283b8de7532 100644 --- a/api/extensions/application_services/workflow.py +++ b/api/extensions/application_services/workflow.py @@ -159,3 +159,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 6614282afc6fbc..fa82127cff3a85 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -40,6 +40,7 @@ from extensions.application_services.rbac import RBACServices, build_rbac_services from extensions.application_services.resource_access_token import build_resource_access_token_service 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, @@ -235,6 +236,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 @@ -294,6 +296,7 @@ class AppScopedEndUserServices: @dataclass(frozen=True, slots=True) class ApplicationServices: + workflow_suggestions: WorkflowInstructionSuggestions annotation_commands: AnnotationCommandService annotation_imports: AnnotationImportService annotation_queries: AnnotationQuery @@ -602,6 +605,7 @@ def build_application_services( ) return ApplicationServices( + workflow_suggestions=build_workflow_suggestions(database_client), annotation_commands=AnnotationCommandService( annotations=annotations, add_index=add_annotation_to_index_task.delay, diff --git a/api/repositories/knowledge/dataset_repository.py b/api/repositories/knowledge/dataset_repository.py index f93015d9a547f1..8d831c269a9c22 100644 --- a/api/repositories/knowledge/dataset_repository.py +++ b/api/repositories/knowledge/dataset_repository.py @@ -47,6 +47,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_maintainer_id(self, workspace_id: str, resource_id: str) -> str | None: with self._session_factory() as session: return session.scalar( 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 f5aee718f09cb5..5c0d14ae563c81 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/helpers/unit/services/workflow_collaboration_fixtures.py b/api/tests/helpers/unit/services/workflow_collaboration_fixtures.py index 5d3f9bee49b7a3..099cd8002d97ce 100644 --- a/api/tests/helpers/unit/services/workflow_collaboration_fixtures.py +++ b/api/tests/helpers/unit/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, diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index 305820c88e67e8..e912e4977bf570 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -71,6 +71,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 @@ -470,3 +471,10 @@ def workflow_runtime(sqlite_session_factory: sessionmaker[Session]) -> WorkflowE 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]) -> WorkflowInstructionSuggestions: + 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.py b/api/tests/unit_tests/controllers/console/app/test_generator.py index 2456439778d7d3..6b6c7d61113547 100644 --- a/api/tests/unit_tests/controllers/console/app/test_generator.py +++ b/api/tests/unit_tests/controllers/console/app/test_generator.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from dataclasses import dataclass from inspect import unwrap as _generator_api_unwrap from unittest.mock import MagicMock @@ -26,6 +27,11 @@ from tests.helpers.unit.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", @@ -569,27 +575,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 = ( @@ -598,18 +608,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( @@ -617,13 +631,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 = _generator_api_unwrap(api.post) @@ -633,7 +650,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", @@ -648,13 +665,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 = _generator_api_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.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py index ba36023dc0c7a5..8fb35394725ff8 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,5 +1,6 @@ +"""Tests for LLM generation and database-backed instruction modification.""" + import json -from collections.abc import Iterator from datetime import datetime, timedelta from decimal import Decimal from unittest.mock import MagicMock, Mock, patch @@ -8,17 +9,14 @@ import pytest from sqlalchemy import event from sqlalchemy.exc import SQLAlchemyError -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session, sessionmaker -import core.llm_generator.llm_generator as generator_module from core.app.app_config.entities import ModelConfig from core.llm_generator import llm_generator as llm_generator_module from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload -from core.llm_generator.llm_generator import LLMGenerator, _parse_string_list +from core.llm_generator.llm_generator import LLMGenerator from core.model_manager import ModelManager from core.plugin.impl.model_runtime_factory import create_plugin_model_manager -from core.workflow.generator import tool_catalogue as tool_catalogue_module -from core.workflow.generator.tool_catalogue import ToolCatalogueEntry from graphon.enums import WorkflowNodeExecutionStatus from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage @@ -31,12 +29,13 @@ WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, ) +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.helpers.unit.core.model_fixtures import make_model_instance from tests.helpers.unit.model_factories import make_dataset -"""Tests for LLM generation and database-backed instruction modification.""" - def _llm_result(content: str) -> LLMResult: return LLMResult( @@ -153,15 +152,9 @@ def _persist_node_execution( class TestLlmGeneratorCases: @staticmethod @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 @staticmethod @pytest.fixture @@ -868,10 +861,8 @@ def _dataset(*, dataset_id: str, tenant_id: str, name: str, created_at: datetime class TestLlmGeneratorMissingCases: @staticmethod @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 class TestParseStringList: @@ -901,53 +892,58 @@ def test_filter_non_strings(self): assert _parse_string_list('["a", 1, "b", {"foo": "bar"}]') == ["a", "b"] class TestGenerateWorkflowInstructionSuggestions: - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_no_default_model(self, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): + @patch("services.workflow.generation.suggestions.ModelManager.for_tenant") + def test_no_default_model(self, mock_for_tenant, monkeypatch: pytest.MonkeyPatch, workflow_suggestions): model_manager, _, resolve_default = _model_manager(monkeypatch) resolve_default.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, monkeypatch: pytest.MonkeyPatch): + @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, monkeypatch: pytest.MonkeyPatch, workflow_suggestions + ): mock_build_context.return_value = "context" model_manager, invocation, _ = _model_manager(monkeypatch) invocation.return_value = _llm_generator_missing_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"] invocation.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, monkeypatch: pytest.MonkeyPatch): + @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, monkeypatch: pytest.MonkeyPatch, workflow_suggestions + ): mock_build_context.return_value = "context" model_manager, invocation, _ = _model_manager(monkeypatch) invocation.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 == [] invocation.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, monkeypatch: pytest.MonkeyPatch): + @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, monkeypatch: pytest.MonkeyPatch, workflow_suggestions + ): mock_build_context.return_value = "context" model_manager, invocation, _ = _model_manager(monkeypatch) invocation.return_value = _llm_generator_missing_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 == [] invocation.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( ( @@ -983,30 +979,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: def test_real_workflow_service_exposes_protocol_methods(self): 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 f75fc26a06f8c8..f811cc86a07a8f 100644 --- a/api/tests/unit_tests/core/workflow/generator/test_runner.py +++ b/api/tests/unit_tests/core/workflow/generator/test_runner.py @@ -8,16 +8,16 @@ import pytest from jinja2 import Template -from core.workflow.generator.runner import ( +from core.workflow.generator.types import GraphDict +from graphon.model_runtime.errors.invoke import InvokeError +from services.workflow.generation.runner import ( WorkflowGenerator, _find_planned_tool_entry, _stage_error_to_envelope_code, _StageJSONError, _StageSchemaError, ) -from core.workflow.generator.tool_catalogue import ToolCatalogueEntry -from core.workflow.generator.types import GraphDict -from graphon.model_runtime.errors.invoke import InvokeError +from services.workflow.generation.tool_catalogue import ToolCatalogueEntry from tests.helpers.unit.config_override import apply_config_overrides """ @@ -261,7 +261,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={}, @@ -1270,7 +1270,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" @@ -1632,7 +1632,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) @@ -1668,7 +1668,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) @@ -1688,7 +1688,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 @@ -1697,7 +1697,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) @@ -1729,7 +1729,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 @@ -1740,7 +1740,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 @@ -1748,7 +1748,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 @@ -1758,7 +1758,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 @@ -1766,7 +1766,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 @@ -3026,7 +3026,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_tool_catalogue.py b/api/tests/unit_tests/core/workflow/generator/test_tool_catalogue.py index bc0a02aa596186..e0bb70722a2e2b 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 @@ -7,7 +7,6 @@ from core.tools.builtin_tool.providers.time.time import WikiPediaProvider from core.tools.builtin_tool.providers.time.tools.current_time import CurrentTimeTool from core.tools.builtin_tool.tool import BuiltinTool -from core.tools.custom_tool.provider import ApiToolProviderController from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ( ToolDescription, @@ -18,8 +17,9 @@ ToolProviderIdentity, ) from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.tool_manager import ToolManager -from core.workflow.generator.tool_catalogue import ( +from services.tools.api.provider import ApiToolProviderController +from services.tools.tool_manager import ToolManager +from services.workflow.generation.tool_catalogue import ( MAX_ROUTED_TOOL_CANDIDATES, MAX_ROUTED_TOOLS_PER_PROVIDER, ToolCapabilityQuery, @@ -32,7 +32,6 @@ format_tool_catalogue, installed_tool_keys, select_legacy_fallback_selection, - select_legacy_fallback_tools, select_tool_candidates, ) @@ -269,19 +268,6 @@ def test_500_tool_catalogue_never_exceeds_global_candidate_limit(self) -> None: assert len(selection.entries) == 15 assert len(selection.entries) <= MAX_ROUTED_TOOL_CANDIDATES - def test_legacy_fallback_keeps_explicit_tool_beyond_first_80(self) -> None: - 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) -> None: @@ -295,10 +281,19 @@ def test_returns_selected_and_omitted_tools(self) -> None: assert selection.limit == 80 assert selection.overflow_count == 0 - def test_keeps_legacy_fallback_tools_compatible(self) -> None: + def test_keeps_explicit_tool_beyond_first_80(self) -> None: 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) -> None: entries = [_entry("provider", f"tool_{index:03d}") for index in range(85)] 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..331998a2ab584b --- /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], +) -> None: + 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, +) -> None: + 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] +) -> None: + 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] +) -> None: + 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..5431664b2e2e31 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 override 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,36 @@ 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(RedisClientWrapper): + @override + def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> Pipeline: + assert transaction is False + assert shard_hint is None + pipe = Pipeline() + pipelines.append(pipe) + return pipe + + repository = WorkflowCollaborationRepository(redis=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 99fde0bca4cdcc..477898128f8813 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.helpers.unit.config_override import config_overrides_context from tests.helpers.unit.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]: 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/test_workflow_generator_service.py b/api/tests/unit_tests/services/test_workflow_generator_service.py index 14ef9a2e6ffd73..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,13 +29,13 @@ ToolProviderIdentity, ) 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.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