From b6503e9a3576472ac30abe7501084caa08499517 Mon Sep 17 00:00:00 2001 From: zhangwei-way <80504572+zhangwei-way@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:59:31 +0800 Subject: [PATCH] feat(knowledge): add per-message RAGFlow retrieval scope (#5238) * feat(knowledge): integrate RAGFlow retrieval and management * test(knowledge): cover merged listing tool * feat(knowledge): add per-message retrieval scope * chore(docs): remove unrelated document * docs(knowledge): add interaction screenshots * feat(knowledge): simplify scope selector trigger * docs(knowledge): refresh selector screenshot * feat(knowledge): defer standalone management * docs(knowledge): show chat-only scope UI * fix(knowledge): honor scope on clarification replies * fix(knowledge): harden scoped replay validation * docs(knowledge): clarify replay scope precedence * fix(knowledge): keep provider settings on tools * fix(config): preserve tools-only knowledge settings * fix(knowledge): submit custom assistant identity * refactor(knowledge): trim PR scope changes * fix(knowledge): sanitize document scope display * feat(knowledge): enable scope selection in main chat * fix(knowledge): emphasize active scope icon without button frame * fix(knowledge): close context scrubbing and refresh e2e checks * fix(knowledge): preserve idempotent canonical retries * fix(knowledge): accept promptless conversation runs * style(knowledge): format backend regression tests * chore(knowledge): trim PR scope and fix frontend format * fix(knowledge): remove shared-scope notice * fix(knowledge): remove scope persistence notice * docs(knowledge): include main chat in catalog scope * fix(knowledge): preserve scope recovery and upgrades * fix(config): preserve LightRAG knowledge upgrades --------- Co-authored-by: foreleven --- README.md | 43 ++ README_zh.md | 11 + backend/app/gateway/AGENTS.md | 10 +- backend/app/gateway/app.py | 4 + .../app/gateway/knowledge_scope_admission.py | 130 +++++ backend/app/gateway/routers/features.py | 23 + backend/app/gateway/routers/knowledge.py | 220 ++++++++ backend/app/gateway/services.py | 240 +++++++- backend/docs/CONFIGURATION.md | 21 +- .../deerflow/agents/middlewares/AGENTS.md | 2 + .../middlewares/knowledge_scope_middleware.py | 134 +++++ .../tool_error_handling_middleware.py | 2 + .../deerflow/community/ragflow/client.py | 14 + .../deerflow/community/ragflow/tools.py | 289 +++++++++- .../harness/deerflow/config/AGENTS.md | 8 + .../harness/deerflow/config/__init__.py | 2 + .../harness/deerflow/config/app_config.py | 5 + .../deerflow/config/knowledge_base_config.py | 15 + .../harness/deerflow/knowledge_scope.py | 194 +++++++ .../harness/deerflow/runtime/runs/worker.py | 5 + .../harness/deerflow/subagents/AGENTS.md | 1 + .../deerflow/subagents/batch_service.py | 1 + .../harness/deerflow/subagents/executor.py | 12 +- .../tools/builtins/batch_task_tool.py | 3 + .../deerflow/tools/builtins/task_tool.py | 5 + .../packages/harness/deerflow/tools/tools.py | 7 + backend/tests/test_batch_task_tool.py | 10 + backend/tests/test_config_version.py | 270 +++++++++ backend/tests/test_features_router.py | 58 ++ .../test_gateway_knowledge_scope_admission.py | 212 ++++++++ backend/tests/test_gateway_services.py | 309 ++++++++++- backend/tests/test_knowledge_router.py | 285 ++++++++++ backend/tests/test_knowledge_scope.py | 239 ++++++++ .../tests/test_knowledge_scope_middleware.py | 181 ++++++ backend/tests/test_lightrag_tools.py | 1 + backend/tests/test_ragflow_client.py | 33 ++ backend/tests/test_ragflow_tools.py | 326 ++++++++++- backend/tests/test_subagent_batch_service.py | 10 + backend/tests/test_task_tool_core_logic.py | 18 + .../test_tool_error_handling_middleware.py | 77 ++- .../test_tool_output_budget_middleware.py | 11 +- config.example.yaml | 17 +- deploy/helm/deer-flow/README.md | 15 +- deploy/helm/deer-flow/values.yaml | 11 +- .../[agent_name]/chats/[thread_id]/page.tsx | 106 +++- .../components/workspace/chats/chat-page.tsx | 91 +++- .../src/components/workspace/input-box.tsx | 4 + .../workspace/knowledge-scope-selector.tsx | 514 ++++++++++++++++++ .../messages/knowledge-scope-summary.tsx | 66 +++ .../workspace/messages/message-list-item.tsx | 6 + frontend/src/core/features/api.ts | 12 + frontend/src/core/features/hooks.ts | 15 + frontend/src/core/i18n/locales/en-US.ts | 38 ++ frontend/src/core/i18n/locales/types.ts | 34 ++ frontend/src/core/i18n/locales/zh-CN.ts | 35 ++ frontend/src/core/knowledge/index.ts | 2 + frontend/src/core/knowledge/scope-api.ts | 65 +++ frontend/src/core/knowledge/scope.ts | 316 +++++++++++ frontend/src/core/threads/hooks.ts | 31 +- frontend/tests/e2e/agent-chat.spec.ts | 5 +- frontend/tests/e2e/knowledge-scope.spec.ts | 216 ++++++++ frontend/tests/e2e/utils/mock-api.ts | 7 + .../knowledge-scope-selector.dom.test.tsx | 121 +++++ .../tests/unit/core/agents/features.test.ts | 26 + .../tests/unit/core/i18n/translations.test.ts | 1 + .../tests/unit/core/knowledge/scope.test.ts | 159 ++++++ .../unit/core/threads/stream-options.test.ts | 9 +- scripts/config-upgrade.sh | 81 ++- 68 files changed, 5361 insertions(+), 83 deletions(-) create mode 100644 backend/app/gateway/knowledge_scope_admission.py create mode 100644 backend/app/gateway/routers/knowledge.py create mode 100644 backend/packages/harness/deerflow/agents/middlewares/knowledge_scope_middleware.py create mode 100644 backend/packages/harness/deerflow/config/knowledge_base_config.py create mode 100644 backend/packages/harness/deerflow/knowledge_scope.py create mode 100644 backend/tests/test_gateway_knowledge_scope_admission.py create mode 100644 backend/tests/test_knowledge_router.py create mode 100644 backend/tests/test_knowledge_scope.py create mode 100644 backend/tests/test_knowledge_scope_middleware.py create mode 100644 frontend/src/components/workspace/knowledge-scope-selector.tsx create mode 100644 frontend/src/components/workspace/messages/knowledge-scope-summary.tsx create mode 100644 frontend/src/core/knowledge/index.ts create mode 100644 frontend/src/core/knowledge/scope-api.ts create mode 100644 frontend/src/core/knowledge/scope.ts create mode 100644 frontend/tests/e2e/knowledge-scope.spec.ts create mode 100644 frontend/tests/unit/components/workspace/knowledge-scope-selector.dom.test.tsx create mode 100644 frontend/tests/unit/core/knowledge/scope.test.ts diff --git a/README.md b/README.md index 3cbb0e86f..91119d6b5 100644 --- a/README.md +++ b/README.md @@ -1095,6 +1095,49 @@ reuse the search entry's key, so search can use a different provider. If you previously configured a shared Tavily key only under `web_search`, also set it under `web_fetch` or use `TAVILY_API_KEY` for both. +### Private Knowledge Retrieval (RAGFlow) + +DeerFlow can optionally connect to a tenant-scoped RAGFlow deployment. The +`knowledge_search` Agent tool resolves the configured dataset scope, groups +datasets by embedding model, and retrieves those groups in parallel so mixed +embedding models do not cause a provider error. Dataset IDs and API keys are +never exposed to the model. The optional `list_knowledge_bases` tool returns +names only. + +Main and custom-agent chats can optionally expose a page-local, icon-only +**Knowledge** selector beside the mode control. Its persistent highlight +indicates that knowledge retrieval is active; the neutral state means retrieval +is off. Set `knowledge_base.scope_selection_enabled: true` in `config.yaml` +while using the built-in RAGFlow `knowledge_search` provider to allow all +permitted datasets, selected datasets/files, or no retrieval for a turn. The +same config flag controls both chat types; when disabled, neither composer +shows the selector or submits a scope. The choice resets to all when the page +is refreshed or another conversation is opened; each sent human message keeps +an immutable scope snapshot for replay and history. The Gateway validates +every snapshot, intersects it with the operator's dataset allowlist, propagates +the execution-only scope to native and durable subagents, and removes it from +model inputs and external traces. Client-supplied internal runtime controls +and credentials are also stripped from run context before execution or +checkpoint persistence. Idempotent retries accept both canonical snapshots and +legacy raw run inputs, preserving retry compatibility across upgrades. +The `knowledge_base` block is provider-neutral and only controls whether the +knowledge capability and selector are enabled. RAGFlow connection, dataset +allowlist, and retrieval parameters (`base_url`, `api_key`, `datasets`, +`page_size`, thresholds, and output limits) must be configured on the +`tools[].name: knowledge_search` entry; they are never read from +`knowledge_base`. +Custom-agent chat requests carry the selected agent name as both `assistant_id` +and `context.agent_name`, so Gateway scope admission and runtime agent loading +use the same identity. Main chat requests use `lead_agent`; both identities are +admitted only when the shared configuration enables the RAGFlow provider. +When answering a pending clarification, an explicitly submitted current +selector snapshot wins; clients that omit it inherit the prior turn's accepted +scope. Edit-and-regenerate follows the same fallback, and the file catalog is +loaded only after a dataset is switched from all files to selected files. +This release does not add an independent Knowledge item to the workspace +sidebar or a DeerFlow knowledge-management page; create, upload, parse, and +delete datasets and documents directly in RAGFlow. + Advanced deployments can enable pluggable authorization with `authorization.enabled` in `config.yaml`. A configured `AuthorizationProvider` filters denied tools before they reach the model or deferred-tool catalog, then the same provider is checked again before every business-tool execution through the existing guardrail middleware. Gateway `threads:*` and `runs:*` route permissions are derived from the same provider, while existing owner checks and admin-only management gates remain in force. Every HTTP route that starts or enables a future Agent run requires `runs:create`: this includes the stateless `POST /api/runs/stream` and `POST /api/runs/wait` endpoints plus scheduled-task create, update, resume, and manual-trigger mutations. Scheduled-task mutations retain their existing `threads:write` requirement, and the stateless routes separately enforce ownership when the optional thread ID is supplied in the request body. A generated `tool_search` may bypass the second tool check only when it fronts the current build's already-filtered deferred catalog. Model access follows the same provider: the Gateway `models` list is filtered per principal, `model:use` is enforced on model detail requests and again when the runtime resolves the agent's model, and a denied default model falls back to the first remaining candidate that also passes `model:use`. The built-in RBAC provider supports per-role `tools`, `routes`, `models`, `skills`, and `sandbox` allow/deny policies and validates that `default_role` names a configured role; authorization is disabled by default. See `config.example.yaml` and the [authorization RFC](docs/plans/2026-07-10-pluggable-authorization-rfc.md). Advanced deployments can also extend the agent runtime itself by declaring `AgentMiddleware` classes under `extensions.middlewares` in `config.yaml` or `extensions_config.json`. Each entry is a `module.path:ClassName` string (zero-argument constructor) or an object `{class, kwargs}` whose `kwargs` are passed to the constructor. `kwargs` values must be JSON types (object, array, string, number, boolean, or null); YAML dates and timestamps are coerced to ISO strings so they match JSON. DeerFlow loads the same configured list into the lead-agent and subagent pipelines after their built-in runtime middlewares and loop/token guards, but before the terminal-response/safety/clarification tail, so enterprise forks can add domain guardrails, tool-call governance, or observability hooks without patching the built-in middleware builders. Missing packages, invalid classes, broken modules, and constructor errors fail loudly at agent creation. Treat `config.yaml` and `extensions_config.json` as trusted operator-controlled files: middleware paths are code execution, just like custom tool, model, sandbox, guardrail, MCP server, and MCP interceptor declarations. Gateway skill/MCP toggle endpoints preserve this field but do not expose an API write path for `extensions.middlewares`. Separate lead-only/subagent-only middleware lists are not supported yet. diff --git a/README_zh.md b/README_zh.md index cf9aa4100..981b40c98 100644 --- a/README_zh.md +++ b/README_zh.md @@ -637,6 +637,17 @@ Skills 采用按需渐进加载，不会一次性把所有内容都塞进上下 Tools 也是同样的思路。DeerFlow 自带一组核心工具：网页搜索、网页抓取、网页渲染截图、文件操作、bash 执行；同时也支持通过 MCP Server 和 Python 函数扩展自定义工具。你可以替换任何一项，也可以继续往里加。 +### 私有知识检索（RAGFlow） + +DeerFlow 可连接租户级 RAGFlow，并通过 `knowledge_search` 按 embedding 模型分组并行召回运维允许的知识库；dataset ID 与 API key 不会暴露给模型。 + +使用内置 RAGFlow `knowledge_search` provider 时，可在 `config.yaml` 中设置 `knowledge_base.scope_selection_enabled: true`，为主智能体和自定义智能体聊天开放模式选择器右侧的纯图标“知识库”按钮。图标持续高亮表示知识检索已启用，普通状态表示本轮检索已关闭。用户可选择全部允许知识库、指定知识库/文件或关闭本轮检索。同一个配置开关统一控制两类聊天；关闭时两类输入框都不显示、也不提交知识范围。选择仅保存在当前页面内，刷新或切换对话后恢复“全部”；每条已发送的人类消息保留不可变的范围快照，用于历史回显、重试和恢复。回复待处理的澄清问题或编辑后重新生成时，若提交了当前选择器快照则以该新范围为准，未提交时继承来源轮次已接纳的范围；知识库仍处于“全部可检索文件”时，展开文件区域不会加载目录，切换为“指定文件”后才加载。Gateway 会校验快照、与运维 allowlist 取交集，把仅含执行字段的范围传递给 native/durable 子智能体，并在模型输入和外部 trace 中清除完整范围。`knowledge_base` 是与 provider 无关的能力开关，只控制知识能力和选择器是否启用；RAGFlow 的连接、dataset allowlist 和检索参数（`base_url`、`api_key`、`datasets`、`page_size`、阈值及输出上限）必须配置在 `tools[].name: knowledge_search` 条目中，`knowledge_base` 中的这些字段不会被读取。 + +自定义智能体聊天请求会同时携带该智能体名称作为 `assistant_id` 和 +`context.agent_name`，确保 Gateway 的范围校验与运行时加载的是同一个智能体；主智能体聊天使用 `lead_agent`，两者都只有在共享配置启用 RAGFlow provider 时才会提交知识范围。 + +本版不在工作区侧边栏增加独立的“知识库”入口，也不提供 DeerFlow 知识库管理页面；知识库和文件的创建、上传、解析与删除仍直接在 RAGFlow 中完成。 + Gateway 生成后续建议时，现在会先把普通字符串输出和 block/list 风格的富文本内容统一归一化，再去解析 JSON 数组响应，因此不同 provider 的内容包装方式不会再悄悄把建议吞掉。 Web UI 支持从已完成的 assistant 回复分叉出一个新的主对话。自动继承的分叉标题会使用下一个空闲的数字后缀（`标题 (2)`、`标题 (3)`……）；显式指定或手动重命名得到的同名后缀也会占号，即使它没有生成序号 metadata，后续自动分叉也不会与它重名。API 调用方显式提供的标题保持不变；重命名会清除旧的生成序号，因此从新标题继续自动分叉时会重新从 `(2)` 开始。最近对话列表还会把已加载的分叉直接排列在已加载的父对话下方，并显示低干扰的树形连接线。父对话尚未加载、谱系数据错误或成环、父子置顶状态不一致时，分叉会安全地保留在顶层，不会被隐藏或跨越置顶边界移动。新 thread 会保留该轮回复的 checkpoint 以及用户消息之前的重放 checkpoint，因此分叉后可以立即重新生成该回复。对于缺少 checkpoint 父链接的旧历史或导入历史，Gateway 会进行有界的时间顺序查找；如果不存在更早的重放 checkpoint，分叉仍会按旧版单-checkpoint 形态成功创建，但无法重新生成继承的回复。已有的单-checkpoint 分叉会保持不变，不会通过不安全的 checkpoint 复制尝试修复。只有从最新回合分叉时才会尽力复制当前 thread 的工作区文件；从历史回合分叉不会带入后续时间线创建的文件。 diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index a75d093da..bd620c336 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -72,7 +72,8 @@ owner-scoped assistant version selection remains enabled. | Router | Endpoints | |--------|-----------| | **Models** (`/api/models`) | `GET /` - list models; `GET /{name}` - model details | -| **Features** (`/api/features`) | `GET /` - UI capabilities: hot-reloaded agents, guarded browser, startup MCP tasks, separate batch repository/worker states so history stays readable without a worker, and `conversation_references` (whether `read_conversation` is configured, plus the per-run reference cap) | +| **Features** (`/api/features`) | `GET /` - UI capabilities: hot-reloaded agents, guarded browser, startup MCP tasks, separate batch repository/worker states so history stays readable without a worker, `conversation_references` (whether `read_conversation` is configured, plus the per-run reference cap), and knowledge scope selection | +| **Knowledge** (`/api/knowledge/retrieval-catalog`) | Authenticated, allowlist-safe, read-only dataset/document catalog used only by main and custom-agent chat scope selection; knowledge management remains in RAGFlow | | **Console** (`/api/console`) | Read-only cross-thread observability for the current user (the data layer for an operations dashboard or external monitoring): `GET /stats` - headline counters (runs/threads/agents/tokens/cost); `GET /runs` - paginated run history joined with thread titles (per-run cost); `GET /usage` - zero-filled daily token series + per-model breakdown with spend. Queries `runs`/`threads_meta` directly as a reporting layer (no new `RunStore` methods); requires a SQL database backend — returns 503 on `database.backend: memory`. Real-cost estimation reads optional `models[*].pricing` (`currency`, `input_per_million`, `output_per_million`, `input_cache_hit_per_million`; `ModelConfig` is `extra="allow"`, so no schema change) and prices each run from its `token_usage_by_model` input/output split. Pricing is **cache-aware**: `RunJournal` accumulates prompt-cache hits from `usage_metadata.input_token_details.cache_read` into a sparse `cache_read_tokens` bucket key (also threaded through `SubagentTokenCollector` → `record_external_llm_usage_records`), and cache-hit input tokens are billed at `input_cache_hit_per_million` (omitted → billed at the miss price, a conservative upper bound). All priced models must use one currency; mixed currencies disable cost reporting and leave cost/currency fields null instead of producing invalid aggregates. Legacy rows fall back to run-level totals at `model_name`; unpriced models yield `cost: null` and cost fields are null when no pricing is configured | | **MCP** (`/api/mcp`) | GET /config - raw/masked; PUT /config - bulk; PATCH /config - toggle; POST /config/servers - add; PUT /config/server - replace; DELETE /config/servers/{server_name:path} - bodyless. Validate expanded, save raw; reload/reset; invalid -> 400. | | **MCP Tasks** (`/api/threads/{id}/mcp-tasks`) | `GET /` - current user's durable tasks for one owned thread; `GET /{task_id}` - bounded result/input/status-error/cancellation-error detail, including cancellation attempt count, without remote task IDs or driver configuration | @@ -91,6 +92,13 @@ owner-scoped assistant version selection remains enabled. | **GitHub Webhooks** (`/api/webhooks/github`) | `POST /` - receive GitHub App / repo webhook deliveries. Verifies `X-Hub-Signature-256` against `GITHUB_WEBHOOK_SECRET`; exempt from auth + CSRF because authenticity is enforced by HMAC. The route is fail-closed: mounted only when `GITHUB_WEBHOOK_SECRET` is set, or when explicit dev opt-in `DEER_FLOW_ALLOW_UNVERIFIED_GITHUB_WEBHOOKS=1` is set. Recognized events include `ping`, `issues`, `issue_comment`, `pull_request`, `pull_request_review`, and `pull_request_review_comment`; unknown events return 200 with `handled=false`. Fan-out runtime failures return 503, keeping the delivery recorded as failed for manual/API/scripted redelivery (GitHub does not automatically retry any failed delivery, 5xx included); permanent/non-retryable conditions such as `channels.github.enabled: false`, unknown events, malformed payloads, or unavailable channel service return 200 with a skipped/handled response. | | **GitHub Event-Driven Agents** | Custom agents can declare a `github:` block in their `config.yaml` to bind to repos and event triggers. Webhook fan-out publishes one `InboundMessage` per matching binding to the channel bus; `GitHubChannel` routes those messages through `ChannelManager`. The response `dispatch` summarizes matched/fired/skipped agents. | +Edit-replay and messages-input clarification may replace checkpoint knowledge +scope with a validated client snapshot; regenerate and `Command` resume always +recover checkpoint scope from the originating genuine user turn, skipping +server-generated hidden context, and ignore client values. Scope admission +treats `input.messages: null` as an empty message collection so +conversation-reference runs without a prompt remain valid. + Thread identifiers use the shared `deerflow.utils.thread_id` contract `^[A-Za-z0-9_-]{1,64}$`. Caller-provided opaque IDs remain supported; UUIDs are generated only for `None`, while explicit empty strings fail validation. diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 4f8a6ebcc..07c3669f8 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -28,6 +28,7 @@ from app.gateway.routers import ( github_webhooks, input_polish, integrations, + knowledge, mcp, mcp_tasks, memory, @@ -906,6 +907,9 @@ This gateway provides runtime endpoints for agent runs plus custom endpoints for # First-party integrations API is mounted at /api/integrations app.include_router(integrations.router) + # Read-only RAGFlow catalog for chat knowledge-scope selection. + app.include_router(knowledge.router) + # Artifacts API is mounted at /api/threads/{thread_id}/artifacts app.include_router(artifacts.router) diff --git a/backend/app/gateway/knowledge_scope_admission.py b/backend/app/gateway/knowledge_scope_admission.py new file mode 100644 index 000000000..c275e4296 --- /dev/null +++ b/backend/app/gateway/knowledge_scope_admission.py @@ -0,0 +1,130 @@ +"""Gateway trust-boundary checks for per-message knowledge scope.""" + +from __future__ import annotations + +from typing import Any + +from fastapi import HTTPException +from langchain_core.messages import BaseMessage, HumanMessage +from pydantic import ValidationError + +from deerflow.knowledge_scope import ( + KNOWLEDGE_SCOPE_KEY, + canonicalize_knowledge_scope, + execution_scope, +) + +RAGFLOW_KNOWLEDGE_SEARCH_PROVIDER = "deerflow.community.ragflow.tools:knowledge_search_tool" + + +def assistant_supports_knowledge_scope( + *, + assistant_id: str | None, + app_config: Any, + agent_config: Any | None, +) -> bool: + """Return whether this exact assistant/provider pairing is supported.""" + if not assistant_id: + return False + knowledge_base = getattr(app_config, "knowledge_base", None) + if not getattr(knowledge_base, "enabled", False): + return False + get_tool_config = getattr(app_config, "get_tool_config", None) + tool = get_tool_config("knowledge_search") if callable(get_tool_config) else None + if getattr(tool, "use", None) != RAGFLOW_KNOWLEDGE_SEARCH_PROVIDER: + return False + # The main assistant has no custom-agent config row. Its knowledge tool is + # controlled solely by the app-level provider configuration. + if assistant_id == "lead_agent": + return True + if agent_config is None: + return False + tool_groups = getattr(agent_config, "tool_groups", None) + return tool_groups is None or "knowledge" in tool_groups + + +def _replace_scope(message: HumanMessage, scope: dict[str, Any] | None) -> HumanMessage: + additional_kwargs = dict(message.additional_kwargs or {}) + if scope is None: + additional_kwargs.pop(KNOWLEDGE_SCOPE_KEY, None) + else: + additional_kwargs[KNOWLEDGE_SCOPE_KEY] = scope + return message.model_copy(update={"additional_kwargs": additional_kwargs}) + + +def admit_message_knowledge_scope( + graph_input: dict[str, Any], + *, + assistant_id: str | None, + app_config: Any, + agent_config: Any | None, + recovery_scope: object | None = None, + recovery: bool = False, +) -> dict[str, Any] | None: + """Canonicalize the sole eligible HumanMessage and return execution scope. + + During regenerate/resume recovery, the server-resolved source snapshot is + authoritative and replaces any client-supplied value. + """ + messages = graph_input.get("messages") + if not isinstance(messages, list): + return None + + scoped_indexes: list[int] = [] + for index, message in enumerate(messages): + if not isinstance(message, BaseMessage): + continue + additional_kwargs = message.additional_kwargs + if KNOWLEDGE_SCOPE_KEY not in additional_kwargs: + continue + if not isinstance(message, HumanMessage): + raise HTTPException( + status_code=422, + detail="knowledge_scope is allowed only on the current HumanMessage", + ) + scoped_indexes.append(index) + if len(scoped_indexes) > 1: + raise HTTPException( + status_code=422, + detail="knowledge_scope is allowed on only one new HumanMessage", + ) + + target_indexes = [index for index, message in enumerate(messages) if isinstance(message, HumanMessage)] + target_index = target_indexes[-1] if target_indexes else None + if scoped_indexes and scoped_indexes[0] != target_index: + raise HTTPException( + status_code=422, + detail="knowledge_scope is allowed only on the current HumanMessage", + ) + if recovery: + raw_scope = recovery_scope + elif scoped_indexes: + target_index = scoped_indexes[0] + raw_scope = messages[target_index].additional_kwargs[KNOWLEDGE_SCOPE_KEY] + else: + return None + + canonical: dict[str, Any] | None = None + if raw_scope is not None: + if not assistant_supports_knowledge_scope( + assistant_id=assistant_id, + app_config=app_config, + agent_config=agent_config, + ): + raise HTTPException( + status_code=422, + detail="knowledge_scope is not supported by this assistant or knowledge provider", + ) + try: + canonical = canonicalize_knowledge_scope(raw_scope) + except ValidationError as exc: + raise HTTPException( + status_code=422, + detail=f"Invalid knowledge_scope: {exc.errors()[0]['msg']}", + ) from exc + elif scoped_indexes and not recovery: + raise HTTPException(status_code=422, detail="knowledge_scope must be an object") + + if target_index is not None: + messages[target_index] = _replace_scope(messages[target_index], canonical) + return execution_scope(canonical) if canonical is not None else None diff --git a/backend/app/gateway/routers/features.py b/backend/app/gateway/routers/features.py index a57b63a3c..51d2b15e4 100644 --- a/backend/app/gateway/routers/features.py +++ b/backend/app/gateway/routers/features.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, Field from app.gateway.browser_capability import browser_capability from app.gateway.conversation_access import conversation_references_enabled from app.gateway.deps import get_config +from app.gateway.knowledge_scope_admission import RAGFLOW_KNOWLEDGE_SEARCH_PROVIDER from app.gateway.run_models import MAX_CONVERSATION_REFERENCES from deerflow.config.app_config import AppConfig from deerflow.subagents.capacity import configured_subagent_max_running @@ -54,6 +55,15 @@ class ConversationReferencesFeature(BaseModel): max_references: int = Field(..., description="Maximum conversation references accepted on one run request") +class KnowledgeBaseFeature(BaseModel): + """Availability of RAGFlow retrieval scope selection in chat.""" + + scope_selection_enabled: bool = Field( + ..., + description="Whether chat may select a per-message RAGFlow retrieval scope", + ) + + class FeaturesResponse(BaseModel): """Frontend-facing feature availability flags.""" @@ -62,6 +72,7 @@ class FeaturesResponse(BaseModel): mcp_tasks: McpTasksFeature subagent_batches: SubagentBatchesFeature conversation_references: ConversationReferencesFeature + knowledge_base: KnowledgeBaseFeature @router.get( @@ -97,4 +108,16 @@ async def list_features(request: Request, config: AppConfig = Depends(get_config enabled=conversation_references_enabled(config), max_references=MAX_CONVERSATION_REFERENCES, ), + knowledge_base=KnowledgeBaseFeature( + scope_selection_enabled=_knowledge_scope_selection_enabled(config), + ), ) + + +def _knowledge_scope_selection_enabled(config: AppConfig) -> bool: + """Fail closed unless the effective knowledge_search entry is RAGFlow.""" + settings = config.knowledge_base + if not settings.enabled or not settings.scope_selection_enabled: + return False + tool = config.get_tool_config("knowledge_search") + return tool is not None and tool.use == RAGFLOW_KNOWLEDGE_SEARCH_PROVIDER diff --git a/backend/app/gateway/routers/knowledge.py b/backend/app/gateway/routers/knowledge.py new file mode 100644 index 000000000..bae5ac5ef --- /dev/null +++ b/backend/app/gateway/routers/knowledge.py @@ -0,0 +1,220 @@ +"""Authenticated, read-only RAGFlow catalog for chat retrieval scope.""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable +from typing import Annotated, Any + +from fastapi import APIRouter, Depends, HTTPException, Path, Query, Request + +from app.gateway.authz import require_permission +from app.gateway.deps import get_config +from app.gateway.knowledge_scope_admission import assistant_supports_knowledge_scope +from deerflow.community.ragflow.client import ( + RAGFlowAPIError, + RAGFlowConnectionError, + RAGFlowProtocolError, +) +from deerflow.community.ragflow.tools import ( + build_ragflow_retrieval_client as _build_retrieval_client, +) +from deerflow.community.ragflow.tools import ( + resolve_ragflow_datasets, + resolve_ragflow_retrieval_settings, +) +from deerflow.config.agents_config import load_agent_config +from deerflow.config.app_config import AppConfig +from deerflow.runtime.user_context import get_effective_user_id + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/api/knowledge", tags=["knowledge"]) + +_DatasetId = Annotated[ + str, + Path(min_length=1, max_length=256, pattern=r"^[A-Za-z0-9_-]+$"), +] +_SCOPE_UNAVAILABLE_DETAIL = "Knowledge scope selection is unavailable for this assistant." + + +async def _catalog_result[Result](operation: Awaitable[Result]) -> Result: + """Normalize provider failures without exposing credentials or raw payloads.""" + try: + return await operation + except RAGFlowAPIError as exc: + logger.warning("RAGFlow rejected a retrieval-catalog request (code=%s)", exc.code) + raise HTTPException( + status_code=502, + detail="RAGFlow rejected the retrieval-catalog request.", + ) from None + except RAGFlowConnectionError: + logger.warning("RAGFlow retrieval catalog could not connect") + raise HTTPException(status_code=502, detail="Unable to connect to RAGFlow.") from None + except RAGFlowProtocolError: + logger.warning("RAGFlow returned an invalid retrieval-catalog response") + raise HTTPException( + status_code=502, + detail="RAGFlow returned an invalid retrieval-catalog response.", + ) from None + except Exception as exc: + logger.error("Unexpected RAGFlow retrieval-catalog failure (%s)", type(exc).__name__) + raise HTTPException(status_code=502, detail="RAGFlow request failed.") from None + + +def _scope_catalog(config: AppConfig, agent_name: str): + knowledge_base = config.knowledge_base + agent_config = None + if agent_name != "lead_agent": + try: + agent_config = load_agent_config( + agent_name, + user_id=get_effective_user_id(), + ) + except (FileNotFoundError, ValueError): + raise HTTPException(status_code=404, detail="Custom agent not found.") from None + if ( + not knowledge_base.enabled + or not knowledge_base.scope_selection_enabled + or not assistant_supports_knowledge_scope( + assistant_id=agent_name, + app_config=config, + agent_config=agent_config, + ) + ): + raise HTTPException(status_code=409, detail=_SCOPE_UNAVAILABLE_DETAIL) + settings, error = resolve_ragflow_retrieval_settings(config) + if settings is None: + logger.warning( + "RAGFlow retrieval catalog settings are unavailable (%s)", + error, + ) + raise HTTPException( + status_code=503, + detail="Knowledge retrieval is not configured.", + ) + return settings + + +def _catalog_page( + items: list[dict[str, Any]], + *, + page: int, + page_size: int, +) -> tuple[list[dict[str, Any]], int]: + total = len(items) + start = (page - 1) * page_size + return items[start : start + page_size], total + + +@router.get("/retrieval-catalog/datasets") +@require_permission("threads", "read") +async def list_retrieval_catalog_datasets( + request: Request, + agent_name: Annotated[str, Query(min_length=1, max_length=128)], + page: Annotated[int, Query(ge=1)] = 1, + page_size: Annotated[int, Query(ge=1, le=100)] = 20, + search: Annotated[str, Query(max_length=256)] = "", + config: AppConfig = Depends(get_config), +) -> dict[str, Any]: + """Return only datasets that the operator permits this agent to retrieve.""" + settings = _scope_catalog(config, agent_name) + client = _build_retrieval_client(settings) + datasets, error = await _catalog_result( + resolve_ragflow_datasets(client, settings), + ) + if datasets is None: + logger.warning( + "RAGFlow retrieval catalog could not resolve operator scope (%s)", + error, + ) + raise HTTPException( + status_code=409, + detail="The configured knowledge-base scope is unavailable.", + ) + + needle = search.strip().casefold() + entries = [ + { + "id": dataset.dataset_id, + "name": dataset.name, + "selectable": bool(dataset.embedding_model) and dataset.chunk_count != 0, + } + for dataset in datasets + if not needle or needle in dataset.name.casefold() + ] + selected, total = _catalog_page(entries, page=page, page_size=page_size) + return { + "items": selected, + "page": page, + "page_size": page_size, + "total": total, + } + + +@router.get("/retrieval-catalog/datasets/{dataset_id}/documents") +@require_permission("threads", "read") +async def list_retrieval_catalog_documents( + dataset_id: _DatasetId, + request: Request, + agent_name: Annotated[str, Query(min_length=1, max_length=128)], + page: Annotated[int, Query(ge=1)] = 1, + page_size: Annotated[int, Query(ge=1, le=100)] = 20, + search: Annotated[str, Query(max_length=256)] = "", + config: AppConfig = Depends(get_config), +) -> dict[str, Any]: + """Return a normalized, read-only document page inside operator scope.""" + settings = _scope_catalog(config, agent_name) + if settings.datasets is not None and dataset_id not in set(settings.datasets): + raise HTTPException(status_code=404, detail="Knowledge base not found.") + client = _build_retrieval_client(settings) + resolved, error = await _catalog_result( + resolve_ragflow_datasets(client, settings, [dataset_id]), + ) + if not resolved: + logger.warning( + "RAGFlow retrieval catalog dataset is unavailable (dataset_id=%s, reason=%s)", + dataset_id, + error, + ) + raise HTTPException(status_code=404, detail="Knowledge base not found.") + + params = [("page", str(page)), ("page_size", str(page_size))] + if search.strip(): + params.append(("keywords", search.strip())) + payload = await _catalog_result( + client.list_documents(dataset_id, params=params), + ) + data = payload.get("data") + docs = data.get("docs") if isinstance(data, dict) else None + if not isinstance(docs, list): + raise HTTPException( + status_code=502, + detail="RAGFlow returned an invalid document list.", + ) + items = [] + for document in docs: + if not isinstance(document, dict): + continue + document_id = document.get("id") + if not isinstance(document_id, str) or not document_id.strip(): + continue + chunk_count = document.get("chunk_count") + searchable = isinstance(chunk_count, int) and not isinstance(chunk_count, bool) and chunk_count > 0 + parsed = document.get("run") == "DONE" + items.append( + { + "id": document_id.strip(), + "name": str(document.get("name") or "Unnamed document"), + "selectable": parsed and searchable, + } + ) + total = data.get("total") + if not isinstance(total, int) or isinstance(total, bool) or total < 0: + total = len(items) + return { + "items": items, + "page": page, + "page_size": page_size, + "total": total, + } diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 63829d553..16583f796 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -33,18 +33,21 @@ from app.gateway.internal_auth import ( get_internal_user, get_trusted_internal_owner_user_id, ) +from app.gateway.knowledge_scope_admission import admit_message_knowledge_scope from app.gateway.run_models import RunCreateRequest from app.gateway.utils import sanitize_log_param from app.mcp_tasks.errors import PermanentNotificationError from deerflow.agents.human_input import read_human_input_response from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY from deerflow.agents.middlewares.input_sanitization_middleware import frame_untrusted_text -from deerflow.agents.middlewares.message_utils import _SUMMARY_MESSAGE_NAME +from deerflow.agents.middlewares.message_utils import _SUMMARY_MESSAGE_NAME, is_genuine_user_message from deerflow.agents.middlewares.tool_receipt import TOOL_RECEIPT_KEY, TOOL_RECEIPT_LEDGER_KEY from deerflow.agents.middlewares.tool_transform_meta import TOOL_TRANSFORMS_KEY from deerflow.agents.middlewares.view_image_middleware import _IMAGE_CONTEXT_MESSAGE_MARKER_KEY +from deerflow.config.agents_config import load_agent_config from deerflow.config.app_config import get_app_config from deerflow.config.database_config import resolve_checkpoint_graph_cache_max +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY from deerflow.projects.context import PROJECT_CONTEXT_MESSAGE_MARKER, resolve_project_context from deerflow.runtime import ( END_SENTINEL, @@ -127,6 +130,7 @@ _SERVER_OWNED_MESSAGE_METADATA_KEYS = ( _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY, _IMAGE_CONTEXT_MESSAGE_MARKER_KEY, + KNOWLEDGE_SCOPE_RUNTIME_KEY, TOOL_RECEIPT_KEY, TOOL_RECEIPT_LEDGER_KEY, TOOL_TRANSFORMS_KEY, @@ -443,13 +447,13 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool of bubbling up as a 500. The gateway is a system boundary, so per-entry validation errors are the right shape for clients to retry against. - ``original_user_content``, dynamic-context reminder markers, the - transient view-image context marker, tool receipts, delegated receipt - metadata/verdicts, and ``untrusted_input`` are server-owned. External callers - cannot supply them; trusted internal channel calls may preserve metadata they - added before invoking this boundary. The same applies to the ``delegations`` - channel: a caller-supplied ledger entry's ``receipt_verdict`` is a forgery and - is stripped before the graph runs. + ``original_user_content``, dynamic-context reminder markers, the transient + view-image context marker, the execution-only knowledge-scope marker, tool + receipts, delegated receipt metadata/verdicts, and ``untrusted_input`` are + server-owned. External callers cannot supply them; trusted internal channel + calls may preserve metadata they added before invoking this boundary. The + same applies to the ``delegations`` channel: a caller-supplied ledger entry's + ``receipt_verdict`` is a forgery and is stripped before the graph runs. ``hide_from_ui`` and a human ``summary`` name are the exception: they stay caller-owned and are deliberately preserved, because ``hide_from_ui`` is also @@ -494,6 +498,25 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool return result +def _canonical_run_record_input( + raw_input: dict[str, Any] | None, + graph_input: object, +) -> dict[str, Any] | None: + """Persist the same normalized messages that cross run admission. + + The run record is a client-visible audit surface. Keeping the original raw + message there would preserve a non-canonical scope even though the graph + receives the validated form. + """ + if not isinstance(graph_input, dict): + return raw_input + canonical = dict(raw_input or {}) + messages = graph_input.get("messages") + if isinstance(messages, list): + canonical["messages"] = [message.model_dump(mode="json") if isinstance(message, BaseMessage) else message for message in messages] + return canonical + + _DEFAULT_ASSISTANT_ID = "lead_agent" @@ -550,6 +573,8 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: frozenset[str] = ( # at admission from threads_meta; a client-supplied value must # never survive in either run-config section. PROJECT_CONTEXT_KEY, + KNOWLEDGE_SCOPE_KEY, + KNOWLEDGE_SCOPE_RUNTIME_KEY, } ) | SANDBOX_SERVER_OWNED_CONTEXT_KEYS @@ -604,7 +629,7 @@ def strip_internal_context_keys(config: dict[str, Any]) -> None: for section in ("context", "configurable"): value = config.get(section) if isinstance(value, dict): - for key in _INTERNAL_ONLY_CONTEXT_KEYS: + for key in _INTERNAL_ONLY_CONTEXT_KEYS | _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: value.pop(key, None) @@ -1451,6 +1476,140 @@ async def ensure_checkpoint_history_seeded( logger.info("Seeded %d checkpoint-history events for thread %s", len(events), thread_id) +def _message_identifier(message: Any) -> str | None: + if isinstance(message, BaseMessage): + return str(message.id) if message.id else None + if isinstance(message, Mapping): + value = message.get("id") + return str(value) if value else None + return None + + +def _message_additional_kwargs(message: Any) -> Mapping[str, Any]: + if isinstance(message, BaseMessage): + return message.additional_kwargs + if isinstance(message, Mapping): + value = message.get("additional_kwargs") + return value if isinstance(value, Mapping) else {} + return {} + + +def _is_scope_source_human_message(message: Any) -> bool: + """Return whether a checkpoint message can originate a recovered scope.""" + if isinstance(message, HumanMessage): + return is_genuine_user_message(message) + if not isinstance(message, Mapping): + return False + if message.get("type") != "human" and message.get("role") not in {"human", "user"}: + return False + return not _skips_input_guardrail(dict(_message_additional_kwargs(message)), message.get("name")) + + +async def _recover_run_knowledge_scope( + request: Request, + *, + thread_id: str, + target_message_id: str | None, +) -> object | None: + """Resolve one replay/resume scope from the authoritative latest checkpoint.""" + accessor, config = await build_thread_checkpoint_state_accessor( + request, + thread_id=thread_id, + ) + try: + snapshot = await accessor.aget(config) + except Exception as exc: + logger.exception("Failed to recover knowledge scope for thread %s", sanitize_log_param(thread_id)) + raise HTTPException(status_code=500, detail="Failed to recover knowledge scope") from exc + values = getattr(snapshot, "values", None) + messages = values.get("messages") if isinstance(values, Mapping) else None + if not isinstance(messages, list): + messages = [] + + source: Any | None = None + if target_message_id: + target_index = next( + (index for index, message in enumerate(messages) if _message_identifier(message) == target_message_id), + None, + ) + if target_index is not None: + source = next( + (message for message in reversed(messages[:target_index]) if _is_scope_source_human_message(message)), + None, + ) + else: + # Interrupted assistant output may never reach a checkpoint. Its + # source is still the terminal HumanMessage of the latest state. + source = next( + (message for message in reversed(messages) if _is_scope_source_human_message(message)), + None, + ) + if source is None: + raise HTTPException( + status_code=409, + detail="Could not recover the source HumanMessage knowledge_scope", + ) + else: + source = next( + (message for message in reversed(messages) if _is_scope_source_human_message(message)), + None, + ) + if source is None: + return None + additional_kwargs = _message_additional_kwargs(source) + return additional_kwargs.get(KNOWLEDGE_SCOPE_KEY) + + +def _current_human_message(graph_input: object) -> HumanMessage | None: + if not isinstance(graph_input, Mapping): + return None + messages = graph_input.get("messages") + if not isinstance(messages, list): + return None + return next( + (message for message in reversed(messages) if isinstance(message, HumanMessage)), + None, + ) + + +async def _load_scope_agent_config( + *, + assistant_id: str | None, + user_id: str | None, +) -> Any | None: + if not assistant_id or assistant_id == _DEFAULT_ASSISTANT_ID: + return None + normalized = assistant_id.strip().lower().replace("_", "-") + try: + return await asyncio.to_thread( + load_agent_config, + normalized, + user_id=user_id, + ) + except (FileNotFoundError, ValueError) as exc: + raise HTTPException( + status_code=422, + detail="knowledge_scope assistant configuration could not be resolved", + ) from exc + + +async def _validate_scope_thread_binding( + run_ctx: RunContext, + *, + thread_id: str, + assistant_id: str | None, +) -> None: + existing = await run_ctx.thread_store.get(thread_id) + if not isinstance(existing, Mapping): + return + bound = existing.get("assistant_id") + if isinstance(bound, str) and bound and assistant_id and bound != assistant_id: + raise HTTPException( + status_code=409, + detail="Thread assistant does not match knowledge_scope assistant", + ) + + # --------------------------------------------------------------------------- # Run lifecycle # --------------------------------------------------------------------------- @@ -1585,6 +1744,54 @@ async def start_run( config = build_run_config(thread_id, body.config, run_metadata, assistant_id=body.assistant_id) await apply_checkpoint_to_run_config(config, body=body, thread_id=thread_id, request=request) + replay_kind = run_metadata.get("replay_kind") + target_message_id = run_metadata.get("regenerate_from_message_id") + scope_graph_input = graph_input if isinstance(graph_input, dict) else {"messages": []} + scope_messages = scope_graph_input.get("messages") + candidate_has_scope = isinstance(scope_messages, list) and any(isinstance(message, BaseMessage) and KNOWLEDGE_SCOPE_KEY in message.additional_kwargs for message in scope_messages) + current_human_message = _current_human_message(graph_input) + current_message_has_scope = current_human_message is not None and KNOWLEDGE_SCOPE_KEY in current_human_message.additional_kwargs + replay_requires_scope_recovery = isinstance(graph_input, Command) or (isinstance(target_message_id, str) and bool(target_message_id) and (replay_kind != "edit" or not current_message_has_scope)) + is_human_input_response = current_human_message is not None and "human_input_response" in current_human_message.additional_kwargs + # Clarification and edit-replay messages may intentionally replace the + # source scope. If either client omits its current selector snapshot, + # inherit the source turn's authoritative scope instead of widening the + # run to every operator-approved dataset. Other replay paths always use + # server recovery regardless of client input. + is_scope_recovery = replay_requires_scope_recovery or (is_human_input_response and not current_message_has_scope) + recovery_scope = ( + await _recover_run_knowledge_scope( + request, + thread_id=thread_id, + target_message_id=(target_message_id if isinstance(target_message_id, str) else None), + ) + if is_scope_recovery + else None + ) + agent_config = ( + await _load_scope_agent_config( + assistant_id=body.assistant_id, + user_id=owner_user_id or (str(user.id) if user is not None else None), + ) + if candidate_has_scope or recovery_scope is not None + else None + ) + admitted_knowledge_scope = admit_message_knowledge_scope( + scope_graph_input, + assistant_id=body.assistant_id, + app_config=run_ctx.app_config or get_app_config(), + agent_config=agent_config, + recovery_scope=recovery_scope, + recovery=is_scope_recovery, + ) + if admitted_knowledge_scope is not None: + await _validate_scope_thread_binding( + run_ctx, + thread_id=thread_id, + assistant_id=body.assistant_id, + ) + run_record_input = _canonical_run_record_input(body.input, graph_input) + # Merge DeerFlow-specific context overrides into both ``configurable`` and ``context``. # The ``context`` field is a custom extension for the langgraph-compat layer # that carries agent configuration (model_name, thinking_enabled, etc.). @@ -1724,6 +1931,7 @@ async def start_run( stream_subgraphs=body.stream_subgraphs, interrupt_before=body.interrupt_before, interrupt_after=body.interrupt_after, + knowledge_scope=admitted_knowledge_scope, ) try: @@ -1750,7 +1958,11 @@ async def start_run( # written to runs.kwargs_json and echoed by the run API, so a # request-scoped secret (#3861) must not ride along. The live # config built above keeps the secrets for the actual run. - kwargs={"input": body.input, "config": redact_config_secrets(body.config), **({"conversation_references": conversation_references} if conversation_references else {})}, + kwargs={ + "input": run_record_input, + "config": redact_config_secrets(body.config), + **({"conversation_references": conversation_references} if conversation_references else {}), + }, multitask_strategy=body.multitask_strategy, model_name=model_name, user_id=owner_user_id, @@ -1759,7 +1971,13 @@ async def start_run( if record.idempotency_reused: stored = record.kwargs or {} - if stored.get("input") != body.input or record.assistant_id != body.assistant_id or stored.get("conversation_references", []) != conversation_references: + stored_input = stored.get("input") + # New runs persist the admitted, canonical message snapshot + # so a scope display cannot be rewritten through the run + # record. Accept the raw request as well for records written + # by older Gateway versions, while comparing canonical + # retries to the same representation as the stored record. + if (stored_input != body.input and stored_input != run_record_input) or record.assistant_id != body.assistant_id or stored.get("conversation_references", []) != conversation_references: raise HTTPException( status_code=409, detail="Idempotency-Key already used with a different request", diff --git a/backend/docs/CONFIGURATION.md b/backend/docs/CONFIGURATION.md index ac3171d0f..a6c339117 100644 --- a/backend/docs/CONFIGURATION.md +++ b/backend/docs/CONFIGURATION.md @@ -337,7 +337,26 @@ reachable from the Gateway container or Pod; `localhost` refers to that container or Pod, not the host machine. This integration is retrieval-only. Dataset creation, uploads, parsing, and -deletion remain in RAGFlow and are not exposed as Agent tools or DeerFlow APIs. +deletion remain in RAGFlow and are not exposed as Agent tools, workspace pages, +or DeerFlow APIs. The authenticated `/api/knowledge/retrieval-catalog` routes +exist only to populate the custom-agent chat selector. The provider-neutral +`knowledge_base` block only gates DeerFlow's knowledge capability and selector; +configure the RAGFlow connection and retrieval defaults on the +`tools[].name: knowledge_search` entry shown above: + +```yaml +knowledge_base: + enabled: true + scope_selection_enabled: true +``` + +When enabled, include the `list_knowledge_bases` tool entry shown above if the +model should be able to discover configured dataset names. The frontend uses +`GET /api/features -> knowledge_base` only to gate the custom-agent chat +selector. RAGFlow API keys and dataset UUIDs are never returned to the browser +or model. Do not put RAGFlow-specific connection, allowlist, or retrieval +parameters in `knowledge_base`; they are read only from the provider tool +entry, so different knowledge providers can use their own settings. ### LightRAG Knowledge Retrieval diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index d4927cd27..3131e3a55 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -52,6 +52,8 @@ strict providers reject. **Shared runtime base** (`build_lead_runtime_middlewares`; subagents reuse most of this via `build_subagent_runtime_middlewares`): 1. **InputSanitizationMiddleware** - First, so it is the outermost `wrap_model_call` wrapper; every inner middleware (including LLM retries) sees sanitized messages. `additional_kwargs.original_user_content` is server-owned provenance: Gateway strips caller-supplied values for non-internal run requests, trusted IM calls may carry the string they captured before adding transport/file context, and the middleware replaces any non-string value before wrapping. Uploads and sanitization retain first-writer-wins only for validated strings. Caller markers are marked `untrusted_input`, never stripped; scope is every turn. + + **KnowledgeScopeMiddleware** follows input sanitization: it exposes only Gateway-admitted execution scope, removes scope/display data from model messages, and blocks `knowledge_search` when disabled without reading storage or RAGFlow. 2. **ToolOutputBudgetMiddleware** - Caps model-bound tool output per app config. Externalizes oversized results to `tool_output.storage_subdir` (default `.tool-results`, constant `TOOL_RESULTS_DIRNAME`) under thread outputs, leaving a typed synopsis + `read_file` reference. These process-feedback files are excluded from workspace-change scans and delivery verification. `wrap_model_call` elides successful `write_file` content only in model-bound requests (#5328) after a later successful same-path `read_file`/`write_file`/`str_replace`; the on-disk file becomes the reference. Preserves the newest `keep_recent_writes` writes. Pairs call occurrences via `tool_call_args.pair_tool_call_results` and rewrites through shared `tool_call_args` helpers; controls: `elide_superseded_writes`, `superseded_write_min_chars`. 3. **ToolResultSanitizationMiddleware** - Neutralizes framework/injection tags (e.g. ``) and boundary markers in *remote-content* tool results (`web_fetch`/`web_search`/`image_search`/`web_capture`) so attacker-controlled fetched pages cannot forge trusted framework context. Mirrors `InputSanitizationMiddleware`'s user-input guardrail for the other untrusted-content entry point; sits inner of `ToolOutputBudgetMiddleware` (neutralizes the raw output, then the budget truncates). Local tool output (bash/read_file) is left untouched. Scope is a name-based allowlist for the first-party web tools, plus every MCP-sourced tool via its `deerflow_mcp` metadata tag, so an MCP server naming its fetcher `fetch_url` is still covered diff --git a/backend/packages/harness/deerflow/agents/middlewares/knowledge_scope_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/knowledge_scope_middleware.py new file mode 100644 index 000000000..f61e539c7 --- /dev/null +++ b/backend/packages/harness/deerflow/agents/middlewares/knowledge_scope_middleware.py @@ -0,0 +1,134 @@ +"""Enforce per-message knowledge scope at model and tool boundaries.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Any, override + +from langchain.agents import AgentState +from langchain.agents.middleware import AgentMiddleware +from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse +from langchain_core.messages import HumanMessage, ToolMessage +from langgraph.prebuilt.tool_node import ToolCallRequest +from langgraph.runtime import Runtime +from langgraph.types import Command + +from deerflow.knowledge_scope import ( + KNOWLEDGE_SCOPE_KEY, + KNOWLEDGE_SCOPE_RUNTIME_KEY, + canonicalize_knowledge_scope, + execution_scope, + strip_message_knowledge_scope, +) +from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY + +_KNOWLEDGE_SEARCH_TOOL_NAME = "knowledge_search" + + +def _runtime_context(value: object) -> dict[str, Any] | None: + context = getattr(value, "context", None) + return context if isinstance(context, dict) else None + + +def _scope_from_runtime(value: object) -> dict[str, Any] | None: + context = _runtime_context(value) + if context is None or KNOWLEDGE_SCOPE_RUNTIME_KEY not in context: + return None + return execution_scope(canonicalize_knowledge_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY])) + + +class KnowledgeScopeMiddleware(AgentMiddleware[AgentState]): + """Project execution scope, redact message snapshots, and enforce disabled.""" + + @override + def before_agent(self, state: AgentState, runtime: Runtime) -> None: + context = _runtime_context(runtime) + if context is None: + return + + admitted = _scope_from_runtime(runtime) + if admitted is not None: + context[KNOWLEDGE_SCOPE_RUNTIME_KEY] = admitted + return + + messages = list((state or {}).get("messages") or []) + raw_boundary = context.get(CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY) + if isinstance(raw_boundary, (set, frozenset, list, tuple)): + pre_existing_ids = {str(message_id) for message_id in raw_boundary if message_id} + candidates = [message for message in messages if str(getattr(message, "id", "") or "") not in pre_existing_ids] + else: + # Standalone harness callers do not always expose a checkpoint + # boundary. Only the terminal input message is eligible; never + # search backwards through arbitrary history for a scope. + candidates = messages[-1:] + + scoped = [] + for message in candidates: + additional_kwargs = getattr(message, "additional_kwargs", None) + if isinstance(additional_kwargs, Mapping) and KNOWLEDGE_SCOPE_KEY in additional_kwargs: + if not isinstance(message, HumanMessage): + raise ValueError("knowledge_scope is allowed only on a current HumanMessage") + scoped.append(additional_kwargs[KNOWLEDGE_SCOPE_KEY]) + if len(scoped) > 1: + raise ValueError("only one current HumanMessage may carry knowledge_scope") + if scoped: + context[KNOWLEDGE_SCOPE_RUNTIME_KEY] = execution_scope(canonicalize_knowledge_scope(scoped[0])) + + @staticmethod + def _prepare_model_request(request: ModelRequest) -> ModelRequest: + messages = [strip_message_knowledge_scope(message) for message in request.messages] + tools = list(request.tools) + scope = _scope_from_runtime(request.runtime) + if scope is not None and scope["mode"] == "disabled": + tools = [tool for tool in tools if getattr(tool, "name", None) != _KNOWLEDGE_SEARCH_TOOL_NAME] + if messages == list(request.messages) and tools == list(request.tools): + return request + return request.override(messages=messages, tools=tools) + + @staticmethod + def _disabled_tool_message(request: ToolCallRequest) -> ToolMessage | None: + if str(request.tool_call.get("name") or "") != _KNOWLEDGE_SEARCH_TOOL_NAME: + return None + scope = _scope_from_runtime(getattr(request, "runtime", None)) + if scope is None or scope["mode"] != "disabled": + return None + return ToolMessage( + content="Error: Knowledge search is disabled for this turn.", + tool_call_id=str(request.tool_call.get("id") or "missing_tool_call_id"), + name=_KNOWLEDGE_SEARCH_TOOL_NAME, + status="error", + ) + + @override + def wrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], ModelResponse], + ) -> ModelCallResult: + return handler(self._prepare_model_request(request)) + + @override + async def awrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], Awaitable[ModelResponse]], + ) -> ModelCallResult: + return await handler(self._prepare_model_request(request)) + + @override + def wrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], ToolMessage | Command], + ) -> ToolMessage | Command: + blocked = self._disabled_tool_message(request) + return blocked if blocked is not None else handler(request) + + @override + async def awrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]], + ) -> ToolMessage | Command: + blocked = self._disabled_tool_message(request) + return blocked if blocked is not None else await handler(request) diff --git a/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py index 79b8fae4e..e1c43f6db 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py @@ -172,6 +172,7 @@ def _build_runtime_middlewares( ) -> list[AgentMiddleware]: """Build shared base middlewares for agent execution.""" from deerflow.agents.middlewares.input_sanitization_middleware import InputSanitizationMiddleware + from deerflow.agents.middlewares.knowledge_scope_middleware import KnowledgeScopeMiddleware from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware from deerflow.agents.middlewares.thread_data_middleware import ThreadDataMiddleware from deerflow.agents.middlewares.tool_output_budget_middleware import ToolOutputBudgetMiddleware @@ -189,6 +190,7 @@ def _build_runtime_middlewares( # neutralized text. outer_wrappers: list[AgentMiddleware] = [ InputSanitizationMiddleware(), + KnowledgeScopeMiddleware(), ToolOutputBudgetMiddleware.from_app_config(app_config), ToolResultSanitizationMiddleware(), ] diff --git a/backend/packages/harness/deerflow/community/ragflow/client.py b/backend/packages/harness/deerflow/community/ragflow/client.py index 5d9d69eec..454c0c1f2 100644 --- a/backend/packages/harness/deerflow/community/ragflow/client.py +++ b/backend/packages/harness/deerflow/community/ragflow/client.py @@ -157,11 +157,21 @@ class RAGFlowClient: raise RAGFlowProtocolError(f"RAGFlow dataset listing exceeded {_MAX_DATASET_PAGES} pages.") + async def list_documents( + self, + dataset_id: str, + *, + params: list[tuple[str, str]], + ) -> dict[str, Any]: + """Proxy one document-list request for a dataset.""" + return await self._request("GET", f"/datasets/{dataset_id}/documents", params=params) + async def retrieve( self, query: str, *, dataset_ids: list[str], + document_ids: list[str] | None = None, page_size: int = 8, similarity_threshold: float = 0.2, vector_similarity_weight: float = 0.3, @@ -170,6 +180,8 @@ class RAGFlowClient: """Retrieve chunks from an explicit, non-empty dataset allowlist.""" if not dataset_ids or not all(isinstance(dataset_id, str) and dataset_id.strip() for dataset_id in dataset_ids): raise ValueError("dataset_ids must contain at least one dataset ID") + if document_ids is not None and (not document_ids or not all(isinstance(document_id, str) and document_id.strip() for document_id in document_ids)): + raise ValueError("document_ids must be omitted or non-empty") request_body: dict[str, object] = { "question": query, @@ -179,6 +191,8 @@ class RAGFlowClient: "vector_similarity_weight": vector_similarity_weight, "top_k": top_k, } + if document_ids is not None: + request_body["document_ids"] = document_ids payload = await self._request("POST", "/retrieval", json=request_body) data = payload.get("data") diff --git a/backend/packages/harness/deerflow/community/ragflow/tools.py b/backend/packages/harness/deerflow/community/ragflow/tools.py index df41a15ca..31e6bf76f 100644 --- a/backend/packages/harness/deerflow/community/ragflow/tools.py +++ b/backend/packages/harness/deerflow/community/ragflow/tools.py @@ -5,14 +5,20 @@ from __future__ import annotations import asyncio import logging import re -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass -from typing import Any +from typing import Any, cast from langchain_core.tools import StructuredTool from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, SecretStr, ValidationError, field_validator from deerflow.config import get_app_config +from deerflow.knowledge_scope import ( + KNOWLEDGE_SCOPE_RUNTIME_KEY, + canonicalize_knowledge_scope, + execution_scope, +) +from deerflow.tools.types import Runtime from .client import RAGFlowAPIError, RAGFlowClient, RAGFlowConnectionError, RAGFlowProtocolError from .formatting import format_retrieval_result @@ -21,7 +27,7 @@ logger = logging.getLogger(__name__) _warned: set[str] = set() _RAGFLOW_UUID_PATTERN = re.compile(r"(? _RAGFlowRetrievalSettin return _RAGFlowRetrievalSettings.model_validate(dict(extra)) -def _settings_or_error() -> tuple[_RAGFlowRetrievalSettings | None, str | None]: - tool_config = get_app_config().get_tool_config("knowledge_search") +def _settings_or_error(app_config: Any | None = None) -> tuple[_RAGFlowRetrievalSettings | None, str | None]: + app_config = app_config or get_app_config() + get_tool_config = getattr(app_config, "get_tool_config", lambda _name: None) + tool_config = get_tool_config("knowledge_search") if tool_config is None: return None, "Error: knowledge_search is not configured; add its RAGFlow settings to the tools list in config.yaml." + tool_values = dict(getattr(tool_config, "model_extra", None) or {}) try: - settings = _settings_from_extra(tool_config.model_extra or {}) + settings = _settings_from_extra(tool_values) except ValidationError: logger.warning("RAGFlow knowledge_search tool configuration is invalid") return None, "Error: Invalid RAGFlow settings for knowledge_search; check config.yaml." @@ -205,10 +221,66 @@ def _log_missing_dataset(*, position: int, dataset_id: str, code: object = None) ) +async def _bounded_gather[InputT, ResultT]( + items: list[InputT], + operation: Callable[[InputT], Awaitable[ResultT]], +) -> list[ResultT]: + """Run provider requests concurrently while preserving input order.""" + semaphore = asyncio.Semaphore(_MAX_PARALLEL_RAGFLOW_REQUESTS) + + async def run(item: InputT) -> ResultT: + async with semaphore: + return await operation(item) + + raw_results = await asyncio.gather( + *(run(item) for item in items), + return_exceptions=True, + ) + results: list[ResultT] = [] + for result in raw_results: + if isinstance(result, BaseException): + raise result + results.append(cast(ResultT, result)) + return results + + +async def _list_datasets_by_id( + client: RAGFlowClient, + dataset_ids: list[str], +) -> list[list[dict]]: + async def list_dataset(dataset_id: str) -> list[dict]: + return await client.list_datasets(dataset_id=dataset_id) + + return await _bounded_gather(dataset_ids, list_dataset) + + async def _resolve_datasets( client: RAGFlowClient, settings: _RAGFlowRetrievalSettings, + requested_dataset_ids: list[str] | None = None, ) -> tuple[list[_ResolvedDataset] | None, str | None]: + if requested_dataset_ids is not None: + if settings.datasets is not None: + operator_allowlist = set(settings.datasets) + if any(dataset_id not in operator_allowlist for dataset_id in requested_dataset_ids): + logger.warning("Selected RAGFlow scope contains a dataset outside the operator allowlist") + return ( + None, + "Error: The selected knowledge scope is no longer available; choose the knowledge bases again.", + ) + dataset_results = await _list_datasets_by_id(client, requested_dataset_ids) + resolved_datasets: list[_ResolvedDataset] = [] + for dataset_id, datasets in zip(requested_dataset_ids, dataset_results, strict=True): + resolved = _current_dataset(datasets, dataset_id) + if resolved is None: + logger.warning("Selected RAGFlow dataset is missing or inaccessible (dataset_id=%s)", dataset_id) + return ( + None, + "Error: The selected knowledge scope is no longer available; choose the knowledge bases again.", + ) + resolved_datasets.append(resolved) + return resolved_datasets, None + if settings.datasets is None: datasets = await client.list_datasets() resolved_by_id: dict[str, _ResolvedDataset] = {} @@ -225,9 +297,12 @@ async def _resolve_datasets( ) return list(resolved_by_id.values()), None + dataset_results = await _list_datasets_by_id(client, settings.datasets) resolved_datasets: list[_ResolvedDataset] = [] - for position, bound_id in enumerate(settings.datasets, start=1): - datasets = await client.list_datasets(dataset_id=bound_id) + for position, (bound_id, datasets) in enumerate( + zip(settings.datasets, dataset_results, strict=True), + start=1, + ): resolved = _current_dataset(datasets, bound_id) if resolved is None: _log_missing_dataset(position=position, dataset_id=bound_id) @@ -237,6 +312,28 @@ async def _resolve_datasets( return resolved_datasets, None +def resolve_ragflow_retrieval_settings( + app_config: Any, +) -> tuple[_RAGFlowRetrievalSettings | None, str | None]: + """Resolve the provider settings shared by tools and safe catalog APIs.""" + return _settings_or_error(app_config) + + +def build_ragflow_retrieval_client( + settings: _RAGFlowRetrievalSettings, +) -> RAGFlowClient: + return _build_client(settings) + + +async def resolve_ragflow_datasets( + client: RAGFlowClient, + settings: _RAGFlowRetrievalSettings, + requested_dataset_ids: list[str] | None = None, +) -> tuple[list[_ResolvedDataset] | None, str | None]: + """Apply the operator allowlist and live RAGFlow dataset resolution.""" + return await _resolve_datasets(client, settings, requested_dataset_ids) + + def _group_searchable_datasets(datasets: list[_ResolvedDataset]) -> list[tuple[str, list[str]]]: groups: dict[str, list[str]] = {} for dataset in datasets: @@ -246,6 +343,73 @@ def _group_searchable_datasets(datasets: list[_ResolvedDataset]) -> list[tuple[s return sorted(groups.items()) +async def _validate_document_filters( + client: RAGFlowClient, + document_filters: list[dict[str, Any]], +) -> tuple[dict[str, list[str]] | None, str | None]: + async def list_documents(item: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Mapping[str, object]]]: + dataset_id = item["dataset_id"] + document_ids = list(item["document_ids"]) + params = [ + ("page", "1"), + ("page_size", str(len(document_ids))), + *(("ids", document_id) for document_id in document_ids), + ] + payload = await client.list_documents(dataset_id, params=params) + data = payload.get("data") + documents = data.get("docs") if isinstance(data, Mapping) else None + if not isinstance(documents, list): + raise RAGFlowProtocolError("RAGFlow returned an invalid document list.") + by_id = {str(document.get("id")): document for document in documents if isinstance(document, Mapping) and document.get("id") is not None} + return item, by_id + + document_results = await _bounded_gather(document_filters, list_documents) + validated: dict[str, list[str]] = {} + for item, by_id in document_results: + dataset_id = item["dataset_id"] + document_ids = list(item["document_ids"]) + for document_id in document_ids: + document = by_id.get(document_id) + chunk_count = document.get("chunk_count") if document is not None else None + run = str(document.get("run") or "").upper() if document is not None else "" + if document is None or run != "DONE" or not isinstance(chunk_count, int) or isinstance(chunk_count, bool) or chunk_count <= 0: + logger.warning( + "Selected RAGFlow document is missing or not searchable (dataset_id=%s, document_id=%s)", + dataset_id, + document_id, + ) + return ( + None, + "Error: The selected knowledge scope is no longer available; choose the knowledge bases or files again.", + ) + validated[dataset_id] = document_ids + return validated, None + + +def _group_scoped_datasets( + datasets: list[_ResolvedDataset], + document_filters: Mapping[str, list[str]], +) -> list[_RetrievalGroup]: + grouped: dict[tuple[str, bool], _RetrievalGroup] = {} + for dataset in datasets: + if dataset.chunk_count == 0: + continue + document_ids = document_filters.get(dataset.dataset_id) + key = (dataset.embedding_model, document_ids is not None) + existing = grouped.get(key) + if existing is None: + grouped[key] = _RetrievalGroup( + embedding_model=dataset.embedding_model, + dataset_ids=[dataset.dataset_id], + document_ids=list(document_ids) if document_ids is not None else None, + ) + continue + existing.dataset_ids.append(dataset.dataset_id) + if document_ids is not None and existing.document_ids is not None: + existing.document_ids.extend(document_ids) + return [grouped[key] for key in sorted(grouped)] + + def _result_chunks(result: Mapping[str, Any]) -> list[Mapping[str, Any]]: chunks = result.get("chunks") if not isinstance(chunks, list): @@ -313,22 +477,24 @@ async def _retrieve_dataset_groups( client: RAGFlowClient, settings: _RAGFlowRetrievalSettings, query: str, - groups: list[tuple[str, list[str]]], + groups: list[_RetrievalGroup], ) -> dict[str, Any]: - semaphore = asyncio.Semaphore(_MAX_PARALLEL_RETRIEVAL_GROUPS) + semaphore = asyncio.Semaphore(_MAX_PARALLEL_RAGFLOW_REQUESTS) - async def retrieve_group(dataset_ids: list[str]) -> dict[str, Any]: + async def retrieve_group(group: _RetrievalGroup) -> dict[str, Any]: async with semaphore: - return await client.retrieve( - query, - dataset_ids=dataset_ids, - page_size=settings.page_size, - similarity_threshold=settings.similarity_threshold, - vector_similarity_weight=settings.vector_similarity_weight, - top_k=settings.top_k, - ) + kwargs: dict[str, Any] = { + "dataset_ids": group.dataset_ids, + "page_size": settings.page_size, + "similarity_threshold": settings.similarity_threshold, + "vector_similarity_weight": settings.vector_similarity_weight, + "top_k": settings.top_k, + } + if group.document_ids is not None: + kwargs["document_ids"] = group.document_ids + return await client.retrieve(query, **kwargs) - results = await asyncio.gather(*(retrieve_group(dataset_ids) for _, dataset_ids in groups), return_exceptions=True) + results = await asyncio.gather(*(retrieve_group(group) for group in groups), return_exceptions=True) successful_results: list[dict[str, Any]] = [] for result in results: if isinstance(result, BaseException): @@ -338,25 +504,61 @@ async def _retrieve_dataset_groups( return _merge_group_results(successful_results, page_size=settings.page_size) -async def knowledge_search(query: str) -> str: +def _runtime_knowledge_scope(runtime: Runtime | None) -> object | None: + context = runtime.context if runtime is not None else None + if not isinstance(context, Mapping): + return None + return context.get(KNOWLEDGE_SCOPE_RUNTIME_KEY) + + +async def knowledge_search( + query: str, + *, + knowledge_scope: object | None = None, + runtime: Runtime | None = None, +) -> str: """Search the configured RAGFlow scope, defaulting to every accessible dataset.""" query = query.strip() if not query: return "Error: query must not be empty." + scope_value = knowledge_scope if knowledge_scope is not None else _runtime_knowledge_scope(runtime) + resolved_scope: dict[str, Any] | None = None + if scope_value is not None: + try: + resolved_scope = execution_scope(canonicalize_knowledge_scope(scope_value)) + except ValidationError: + return "Error: Invalid knowledge scope for this turn." + if resolved_scope["mode"] == "disabled": + return "Error: Knowledge search is disabled for this turn." + settings, error = _settings_or_error() if settings is None: return error or "Error: Invalid RAGFlow settings for knowledge_search; check config.yaml." client = _build_client(settings) try: - datasets, resolution_error = await _resolve_datasets(client, settings) + selected_dataset_ids = resolved_scope.get("dataset_ids") if resolved_scope is not None and resolved_scope["mode"] == "selected" else None + datasets, resolution_error = await _resolve_datasets( + client, + settings, + requested_dataset_ids=selected_dataset_ids, + ) if resolution_error is not None: return resolution_error if not datasets: # Defensive; both resolution paths return a non-empty scope. return "Error: No RAGFlow datasets could be resolved; check knowledge_search in config.yaml." - groups = _group_searchable_datasets(datasets) + document_filters: dict[str, list[str]] = {} + if resolved_scope is not None and resolved_scope["mode"] == "selected": + validated_filters, filter_error = await _validate_document_filters( + client, + list(resolved_scope.get("document_filters") or []), + ) + if filter_error is not None: + return filter_error + document_filters = validated_filters or {} + groups = _group_scoped_datasets(datasets, document_filters) if not groups: return _NO_RELEVANT_CONTENT @@ -375,18 +577,39 @@ async def knowledge_search(query: str) -> str: return _tool_error(exc, settings) +async def list_knowledge_bases() -> str: + """List accessible RAGFlow knowledge-base names without exposing UUIDs.""" + settings, error = _settings_or_error() + if settings is None: + return error or "Error: Invalid RAGFlow settings for knowledge_search; check config.yaml." + + client = _build_client(settings) + try: + datasets, resolution_error = await _resolve_datasets(client, settings) + if resolution_error is not None: + return resolution_error + if not datasets: + return "No accessible RAGFlow datasets were found." + lines = ["Available knowledge bases:"] + for dataset in datasets: + lines.append(f"- {dataset.name}") + return _redact_api_key("\n".join(lines), _api_key(settings)) + except Exception as exc: + return _tool_error(exc, settings) + + def _tool_description() -> str: base = "Search the operator-approved RAGFlow datasets and return compact, citation-numbered source chunks." return f"{base} If knowledge_search.datasets is omitted, all datasets accessible to the configured RAGFlow API key are searched. Dataset IDs are never shown to the model." -async def _knowledge_search_entrypoint(query: str) -> str: +async def _knowledge_search_entrypoint(query: str, runtime: Runtime) -> str: """Search the configured RAGFlow datasets, or every accessible dataset by default. Args: query: Specific question or search terms to retrieve from the configured private documents. """ - return await knowledge_search(query) + return await knowledge_search(query, runtime=runtime) knowledge_search_tool = StructuredTool.from_function( @@ -395,3 +618,19 @@ knowledge_search_tool = StructuredTool.from_function( description=_tool_description(), parse_docstring=True, ) + + +async def _list_knowledge_bases_entrypoint() -> str: + """List the operator-approved RAGFlow knowledge bases by name. + + Dataset UUIDs and other provider metadata are intentionally omitted. + """ + return await list_knowledge_bases() + + +list_knowledge_bases_tool = StructuredTool.from_function( + coroutine=_list_knowledge_bases_entrypoint, + name="list_knowledge_bases", + description="List the names of accessible RAGFlow knowledge bases without exposing dataset IDs.", + parse_docstring=True, +) diff --git a/backend/packages/harness/deerflow/config/AGENTS.md b/backend/packages/harness/deerflow/config/AGENTS.md index 8b7958e03..d592c6670 100644 --- a/backend/packages/harness/deerflow/config/AGENTS.md +++ b/backend/packages/harness/deerflow/config/AGENTS.md @@ -24,6 +24,13 @@ Setup: Copy `config.example.yaml` to `config.yaml` in the **project root** direc **Config Versioning**: `config.example.yaml` has a `config_version` field. On startup, `AppConfig.from_file()` compares user version vs example version and emits a warning if outdated. Missing `config_version` = version 0. Run `make config-upgrade` to auto-merge missing fields. When changing the config schema, bump `config_version` in `config.example.yaml`. +The v46 upgrade treats an existing `tools[]` entry in the `knowledge` group as +pre-gate enablement and sets `knowledge_base.enabled: true` only when that flag +was absent. Explicit `true` or `false` values remain authoritative. Moving +legacy provider settings out of `knowledge_base` remains specific to the +RAGFlow `knowledge_search` tool; LightRAG and other knowledge providers keep +their tool-local settings unchanged. + Top-level `recursion_limit` and `max_recursion_limit` are hot-reloaded per Gateway run. The former supplies the default when a request omits or provides an invalid value; the latter caps both configured and client-provided budgets. **Config Caching**: `get_app_config()` caches the parsed config, but automatically reloads it when the resolved config path or file content signature changes. The signature includes file metadata and a content digest, so Gateway and LangGraph reads stay aligned with `config.yaml` edits even on object-store or network mounts where mtime can remain stale. @@ -89,6 +96,7 @@ Extensions are optional only in the fallback *search* mode (priority 3-4 above): - `subagent_runtime` - Startup-only shared process admission (`max_running`, bounded async wait queue, queue/reject policy, and queue timeout) for ordinary and durable-batch native subagents - `subagent_batches` - Startup-only explicit durable batch scheduler limits (disabled by default), including separate total, live, and running dimensions plus leases/retries/result bounds - `memory` - Memory system (enabled, storage_path, debounce_seconds, shutdown_flush_timeout_seconds, model_name, max_facts, fact_confidence_threshold, injection_enabled, max_injection_tokens, staleness_review_enabled, staleness_age_days, staleness_min_candidates, staleness_max_removals_per_cycle, staleness_protected_categories, staleness_max_lifetime_multiplier, staleness_max_extension_days) +- `knowledge_base` - Hot-reloadable, provider-agnostic knowledge capability and custom-agent scope-selector flags. It gates the read-only Agent tools and selector; provider connection, allowlist, and retrieval defaults belong to the matching entry in `tools[]` (for example, the RAGFlow `knowledge_search` tool). **`extensions_config.json`**: - `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`, `session_init_timeout`). `routing.mode="prefer"` emits `` prompt guidance; if `tool_search` defers the hinted tool, `McpRoutingMiddleware` can also auto-promote matching deferred schemas before the model call. It does not hard-disable other tools. `session_init_timeout` (default `DEFAULT_MCP_SESSION_INIT_TIMEOUT` = 60s, `null` to disable) bounds server bring-up: tool discovery and persistent stdio session initialization, so a hung server cannot block agent construction indefinitely; durable HTTP/SSE task calls use it for their ephemeral session initialization too. `tool_call_timeout` bounds individual stdio calls and durable-task calls on every transport; other HTTP/SSE tools use transport-level timeouts. diff --git a/backend/packages/harness/deerflow/config/__init__.py b/backend/packages/harness/deerflow/config/__init__.py index 76751936b..0d1169056 100644 --- a/backend/packages/harness/deerflow/config/__init__.py +++ b/backend/packages/harness/deerflow/config/__init__.py @@ -1,5 +1,6 @@ from .app_config import get_app_config from .extensions_config import ExtensionsConfig, get_extensions_config +from .knowledge_base_config import KnowledgeBaseConfig from .loop_detection_config import LoopDetectionConfig from .memory_config import MemoryConfig, get_memory_config from .paths import Paths, get_paths @@ -21,6 +22,7 @@ __all__ = [ "get_paths", "SkillsConfig", "ExtensionsConfig", + "KnowledgeBaseConfig", "get_extensions_config", "LoopDetectionConfig", "MemoryConfig", diff --git a/backend/packages/harness/deerflow/config/app_config.py b/backend/packages/harness/deerflow/config/app_config.py index 84b38c6ec..fc86ef628 100644 --- a/backend/packages/harness/deerflow/config/app_config.py +++ b/backend/packages/harness/deerflow/config/app_config.py @@ -23,6 +23,7 @@ from deerflow.config.file_signature import ConfigSignature as _ConfigSignature from deerflow.config.file_signature import get_config_signature as _get_config_signature from deerflow.config.guardrails_config import GuardrailsConfig, load_guardrails_config_from_dict from deerflow.config.input_polish_config import InputPolishConfig +from deerflow.config.knowledge_base_config import KnowledgeBaseConfig from deerflow.config.loop_detection_config import LoopDetectionConfig from deerflow.config.mcp_tasks_config import McpTasksConfig from deerflow.config.memory_config import MemoryConfig, load_memory_config_from_dict @@ -245,6 +246,10 @@ class AppConfig(BaseModel): summarization: SummarizationConfig = Field(default_factory=SummarizationConfig, description="Conversation summarization configuration") task_continuity: TaskContinuityConfig = Field(default_factory=TaskContinuityConfig, description="Thread-local notes and compacted-source recall") memory: MemoryConfig = Field(default_factory=MemoryConfig, description="Memory subsystem configuration") + knowledge_base: KnowledgeBaseConfig = Field( + default_factory=KnowledgeBaseConfig, + description="Provider-agnostic knowledge capability and custom-agent scope-selection configuration", + ) agents_api: AgentsApiConfig = Field(default_factory=AgentsApiConfig, description="Custom-agent management API configuration") acp_agents: dict[str, ACPAgentConfig] = Field(default_factory=dict, description="ACP-compatible agent configuration") subagents: SubagentsAppConfig = Field(default_factory=SubagentsAppConfig, description="Subagent runtime configuration") diff --git a/backend/packages/harness/deerflow/config/knowledge_base_config.py b/backend/packages/harness/deerflow/config/knowledge_base_config.py new file mode 100644 index 000000000..9402415dd --- /dev/null +++ b/backend/packages/harness/deerflow/config/knowledge_base_config.py @@ -0,0 +1,15 @@ +from pydantic import BaseModel, ConfigDict, Field + + +class KnowledgeBaseConfig(BaseModel): + """Hot-reloadable DeerFlow knowledge capability settings. + + Provider connection and retrieval options belong to the provider tool + entry (for example ``tools[].use: ...ragflow...``), not this generic + capability block. + """ + + model_config = ConfigDict(validate_default=True) + + enabled: bool = Field(default=False) + scope_selection_enabled: bool = Field(default=False) diff --git a/backend/packages/harness/deerflow/knowledge_scope.py b/backend/packages/harness/deerflow/knowledge_scope.py new file mode 100644 index 000000000..28a6c1ac4 --- /dev/null +++ b/backend/packages/harness/deerflow/knowledge_scope.py @@ -0,0 +1,194 @@ +"""Versioned per-message knowledge-retrieval scope contract. + +The message snapshot may contain an untrusted display block for historical +UI rendering. Runtime consumers must use execution_scope, which projects only +the fields that can constrain retrieval. +""" + +from __future__ import annotations + +import json +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +KNOWLEDGE_SCOPE_KEY = "knowledge_scope" +KNOWLEDGE_SCOPE_RUNTIME_KEY = "__knowledge_scope_execution" +KNOWLEDGE_SCOPE_VERSION = 1 +MAX_KNOWLEDGE_SCOPE_BYTES = 64 * 1024 +MAX_DATASET_IDS = 100 +MAX_DOCUMENT_IDS = 1000 +MAX_DISPLAY_DATASETS = 20 +MAX_DISPLAY_DOCUMENTS = 50 +MAX_ID_CODEPOINTS = 256 +MAX_DISPLAY_NAME_CODEPOINTS = 256 + + +class _StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +def _clean_id(value: str) -> str: + if not isinstance(value, str): + raise ValueError("knowledge scope IDs must be strings") + cleaned = value.strip() + if not cleaned or len(cleaned) > MAX_ID_CODEPOINTS: + raise ValueError(f"knowledge scope IDs must contain between 1 and {MAX_ID_CODEPOINTS} characters") + return cleaned + + +def _stable_unique_ids(values: list[str]) -> list[str]: + normalized: list[str] = [] + seen: set[str] = set() + for value in values: + cleaned = _clean_id(value) + if cleaned not in seen: + normalized.append(cleaned) + seen.add(cleaned) + return normalized + + +class KnowledgeDocumentFilter(_StrictModel): + dataset_id: str + document_ids: list[str] = Field(min_length=1, max_length=MAX_DOCUMENT_IDS) + + @model_validator(mode="after") + def _normalize(self) -> KnowledgeDocumentFilter: + object.__setattr__(self, "dataset_id", _clean_id(self.dataset_id)) + object.__setattr__(self, "document_ids", _stable_unique_ids(self.document_ids)) + return self + + +class KnowledgeDisplayDocument(_StrictModel): + id: str + name: str = Field(min_length=1, max_length=MAX_DISPLAY_NAME_CODEPOINTS) + + @model_validator(mode="after") + def _normalize(self) -> KnowledgeDisplayDocument: + object.__setattr__(self, "id", _clean_id(self.id)) + if not self.name.strip(): + raise ValueError("display names must not be blank") + return self + + +class KnowledgeDisplayDataset(_StrictModel): + id: str + name: str = Field(min_length=1, max_length=MAX_DISPLAY_NAME_CODEPOINTS) + documents: list[KnowledgeDisplayDocument] | None = Field(default=None, max_length=MAX_DISPLAY_DOCUMENTS) + + @model_validator(mode="after") + def _normalize(self) -> KnowledgeDisplayDataset: + object.__setattr__(self, "id", _clean_id(self.id)) + if not self.name.strip(): + raise ValueError("display names must not be blank") + if self.documents is not None: + document_ids = [item.id for item in self.documents] + if len(document_ids) != len(set(document_ids)): + raise ValueError("display document IDs must not be duplicated") + return self + + +class KnowledgeScopeDisplay(_StrictModel): + datasets: list[KnowledgeDisplayDataset] = Field(max_length=MAX_DISPLAY_DATASETS) + + +class KnowledgeScope(_StrictModel): + """Canonical message snapshot for one user turn.""" + + version: Literal[1] + mode: Literal["all", "selected", "disabled"] + dataset_ids: list[str] | None = Field(default=None, max_length=MAX_DATASET_IDS) + document_filters: list[KnowledgeDocumentFilter] | None = Field(default=None, max_length=MAX_DATASET_IDS) + display: KnowledgeScopeDisplay | None = None + + @model_validator(mode="after") + def _validate_and_normalize(self) -> KnowledgeScope: + dataset_ids = _stable_unique_ids(self.dataset_ids or []) + filters = self.document_filters or [] + + if self.mode in {"all", "disabled"}: + if dataset_ids or filters or self.display is not None: + raise ValueError(f"{self.mode} knowledge scope must not contain selections or display") + object.__setattr__(self, "dataset_ids", None) + object.__setattr__(self, "document_filters", None) + return self._validate_size() + + if not dataset_ids: + raise ValueError("selected knowledge scope requires at least one dataset ID") + object.__setattr__(self, "dataset_ids", dataset_ids) + + allowed_datasets = set(dataset_ids) + filter_ids = [item.dataset_id for item in filters] + if len(filter_ids) != len(set(filter_ids)): + raise ValueError("each dataset may have at most one document filter") + if any(dataset_id not in allowed_datasets for dataset_id in filter_ids): + raise ValueError("document filters must belong to selected datasets") + if sum(len(item.document_ids) for item in filters) > MAX_DOCUMENT_IDS: + raise ValueError(f"knowledge scope may contain at most {MAX_DOCUMENT_IDS} document IDs") + object.__setattr__(self, "document_filters", filters or None) + + if self.display is not None: + display_dataset_ids = [item.id for item in self.display.datasets] + if len(display_dataset_ids) != len(set(display_dataset_ids)): + raise ValueError("display dataset IDs must not be duplicated") + if any(dataset_id not in allowed_datasets for dataset_id in display_dataset_ids): + raise ValueError("display datasets must belong to selected datasets") + filters_by_dataset = {item.dataset_id: set(item.document_ids) for item in filters} + display_document_count = 0 + for dataset in self.display.datasets: + documents = dataset.documents or [] + display_document_count += len(documents) + allowed_documents = filters_by_dataset.get(dataset.id) + if documents and allowed_documents is None: + raise ValueError("display documents require an explicit document filter") + if allowed_documents is not None and any(document.id not in allowed_documents for document in documents): + raise ValueError("display documents must belong to the corresponding document filter") + if display_document_count > MAX_DISPLAY_DOCUMENTS: + raise ValueError(f"display may contain at most {MAX_DISPLAY_DOCUMENTS} document names") + + return self._validate_size() + + def _canonical_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = {"version": self.version, "mode": self.mode} + if self.mode == "selected": + payload["dataset_ids"] = list(self.dataset_ids or []) + if self.document_filters: + payload["document_filters"] = [item.model_dump() for item in self.document_filters] + if self.display is not None: + payload["display"] = self.display.model_dump(exclude_none=True) + return payload + + def _validate_size(self) -> KnowledgeScope: + raw = json.dumps( + self._canonical_dict(), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + if len(raw) > MAX_KNOWLEDGE_SCOPE_BYTES: + raise ValueError(f"knowledge scope must not exceed {MAX_KNOWLEDGE_SCOPE_BYTES} UTF-8 JSON bytes") + return self + + +def canonicalize_knowledge_scope(value: object) -> dict[str, Any]: + """Validate and return the stable JSON-safe message representation.""" + scope = value if isinstance(value, KnowledgeScope) else KnowledgeScope.model_validate(value) + return scope._canonical_dict() + + +def execution_scope(value: object) -> dict[str, Any]: + """Return only the execution fields, excluding the untrusted display block.""" + canonical = canonicalize_knowledge_scope(value) + canonical.pop("display", None) + return canonical + + +def strip_message_knowledge_scope(message: Any) -> Any: + """Copy a LangChain message without its knowledge-scope snapshot.""" + additional_kwargs = getattr(message, "additional_kwargs", None) + if not isinstance(additional_kwargs, dict) or not ({KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY} & additional_kwargs.keys()): + return message + cleaned = dict(additional_kwargs) + cleaned.pop(KNOWLEDGE_SCOPE_KEY, None) + cleaned.pop(KNOWLEDGE_SCOPE_RUNTIME_KEY, None) + return message.model_copy(update={"additional_kwargs": cleaned}) diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 8ca5feba2..5288f2dc6 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -41,6 +41,7 @@ from deerflow.agents.middlewares.input_sanitization_middleware import neutralize from deerflow.config.app_config import AppConfig from deerflow.config.database_config import CheckpointChannelMode from deerflow.constants import CONVERSATION_READER_CONTEXT_KEY, TOOL_RESULTS_DIRNAME +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope from deerflow.runtime.checkpoint_mode import ( aensure_checkpoint_mode_compatible, inject_checkpoint_mode, @@ -522,6 +523,7 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = ( # admission (spec §7.1); a caller-supplied value in # ``config['context']`` must never be merged (§12). PROJECT_CONTEXT_KEY, + KNOWLEDGE_SCOPE_RUNTIME_KEY, } ) | SANDBOX_SERVER_OWNED_CONTEXT_KEYS @@ -808,6 +810,7 @@ async def run_agent( stream_subgraphs: bool = False, interrupt_before: list[str] | Literal["*"] | None = None, interrupt_after: list[str] | Literal["*"] | None = None, + knowledge_scope: dict[str, Any] | None = None, ) -> None: """Execute an agent in the background, publishing events to *bridge*.""" @@ -1088,6 +1091,8 @@ async def run_agent( config["metadata"] = checkpoint_metadata checkpoint_metadata[CHECKPOINT_AGENT_NAME_METADATA_KEY] = DEFAULT_AGENT_NAME_METADATA_VALUE if checkpoint_agent_name is None else checkpoint_agent_name _pin_admission_project_context(config, runtime_ctx) + if knowledge_scope is not None: + runtime_ctx[KNOWLEDGE_SCOPE_RUNTIME_KEY] = execution_scope(knowledge_scope) deerflow_trace_id = _bind_trace_id(config, runtime_ctx) # Expose the run-scoped journal under a sentinel key so middleware can # write audit events (e.g. SafetyFinishReasonMiddleware recording diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 33cd0d490..04e1cfdc9 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -21,6 +21,7 @@ executions are not checked, and acceptance never changes automatic retry policy. **Benefit-based routing policy**: Enabling subagents exposes delegation as an optimization, not a default response to complexity. The lead prompt defaults to direct execution and permits `task` only when parallel latency, specialist capability, or context-isolation benefit clearly exceeds startup, duplicate-discovery, synthesis, state-conflict, and side-effect costs. Inter-agent output dependencies and overlapping mutable state are hard vetoes for parallel dispatch, while duplicate discovery and a cheap direct path remain costs rather than categorical vetoes; a bounded sequential chain may run in one subagent when specialist or context-isolation benefit clearly wins. Parallel scopes must be independent and non-overlapping, the lead uses the fewest useful subagents, and every later batch is re-evaluated while retaining any within-batch parallel benefit. When the enforced per-response limit is 1, the rendered prompt removes parallel and multi-batch benefit guidance and permits delegation only for material specialist or context-isolation benefit. Keep this policy aligned across `lead_agent/prompt.py`, the `task` tool description, and both built-in role descriptions; routing regressions are pinned in `tests/test_subagent_routing_prompt.py`, `tests/test_subagent_prompt_security.py`, and `tests/test_lead_agent_prompt.py`. **User-scoped Skills**: Subagents resolve their configured skills through `get_or_new_user_skill_storage(user_id)` using the parent runtime identity, with `DEFAULT_USER_ID` only when no identity is available. This keeps custom-skill shadowing and visibility aligned with the lead agent instead of reading the global-only catalog. **Upload-state boundary**: Ordinary `task` delegation snapshots a valid parent `ThreadState.uploaded_files` list at dispatch, deep-copies it across the isolated-loop boundary, seeds it into the child's fresh state, and only then makes `list_uploaded_files` eligible for normal tool-policy filtering. An explicit empty list is valid and must be preserved because it means every upload in the thread is historical for this run. Missing or malformed state fails closed with the tool disabled. Durable `batch_task` execution intentionally keeps the tool disabled: delayed and recovered items have no valid parent-run upload boundary, and supporting that case requires a separate persisted-state contract. +**Knowledge-scope boundary**: `task` and `batch_task` inherit the parent's canonical execution scope, never display data. Batch specs persist it across lease recovery; malformed scope fails. Model input cannot set or broaden it, and shared middleware/retrieval enforce `disabled`. **Date context (#4781)**: Every built-in subagent execution registers `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` hook adds a hidden framework-owned `SystemMessage` containing only `` before the first model call; it does not read `AppConfig.memory`, call the memory manager, rewrite the task `HumanMessage`, or inherit the lead agent's frozen-conversation/midnight lifecycle. The coalescer merges that reminder with the subagent's static prompt so strict providers still receive exactly one leading `SystemMessage`. The lead-only `DynamicContextMiddleware` registration and its date, optional-memory, and midnight-update behavior remain unchanged. **Execution**: Ordinary and durable-batch native subagents submit coroutines directly to one persistent isolated event loop. Gateway/embedded startup installs one process-wide async FIFO admission controller (default 3 running, bounded queue). Direct `create_deerflow_agent` callers can instead pass a caller-owned `SubagentRuntime`; reuse the same instance across graphs so its bound `task`, optional batch tools/service, middleware limits, and `SubagentExecutor` all share one controller without reading global YAML. An owned batch service must be started before graph construction and stopped at application shutdown. Waiters hold no scheduler thread, and cancellation/timeout release queue/slot ownership. **Shared sandbox execution lifecycle** (#5128): every admitted subagent run carries a stable task-derived `sandbox_lease_owner_id` and matching `sandbox_command_scope_id` in its runtime context. Sandbox middleware retains that execution against the lead thread's active provider client, so one child finishing cannot close the sandbox while siblings still run; the final holder performs any pending provider release. A rollback/fork-restored child reusing the parent's live client binds a non-releasing holder: it fences parent cleanup and owns its command scope without requesting a park itself; a parent's earlier park request waits for the child, while a missing inherited client falls through to a normal fresh acquire. On AIO, the command scope selects one explicit persistent shell session per subagent, allowing independent scopes to run concurrently while preserving in-order shell state within one child. Sync sandbox tool bodies offloaded with `asyncio.to_thread` are shielded and drained across repeated cancellation before the outer execution can clean its holder; a cancelled worker can therefore neither re-admit an already-released owner nor run after subagent terminalization. `SubagentExecutor` invokes active LangGraph stream cleanup before releasing its sandbox lease or notifying task stop. Slow cooperative cleanup keeps the result non-terminal and retains its sandbox lease and capacity slot; its warning is diagnostic, not a safe hard-timeout boundary. Middleware performs the normal release, and `SubagentExecutor` repeats it idempotently in `finally` so exceptions, cooperative cancellation, and timeout unwind paths cannot leak a lease or scoped session. diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index d7438e63b..ba49706f4 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -237,6 +237,7 @@ class SubagentBatchService: channel_user_id=spec.get("channel_user_id"), is_internal=spec.get("is_internal") is True, authz_attributes=spec.get("authz_attributes"), + knowledge_scope=spec.get("knowledge_scope"), execution_capacity=self._execution_capacity, acceptance_criteria=item.get("acceptance_criteria"), ) diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index 8a2bdf2dc..af3005f41 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -36,6 +36,7 @@ from deerflow.agents.thread_state import SandboxState, ThreadDataState, ThreadSt from deerflow.authz.principal import normalize_authz_attributes from deerflow.config import get_app_config from deerflow.config.app_config import AppConfig +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope from deerflow.models import create_chat_model from deerflow.runtime.runs.stream_cleanup import close_agent_stream from deerflow.runtime.user_context import DEFAULT_USER_ID @@ -56,7 +57,10 @@ from deerflow.subagents.step_events import capture_new_step_messages from deerflow.subagents.token_collector import SubagentTokenCollector from deerflow.subagents.turn_budget import find_jumping_hooks, resolve_recursion_limit from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_context, resolve_trace_id -from deerflow.tracing import build_tracing_callbacks, inject_langfuse_metadata +from deerflow.tracing import ( + build_tracing_callbacks, + inject_langfuse_metadata, +) from deerflow.utils.messages import message_content_to_text if TYPE_CHECKING: @@ -794,6 +798,7 @@ class SubagentExecutor: is_internal: bool = False, authz_attributes: Mapping[str, Any] | None = None, deerflow_trace_id: str | None = None, + knowledge_scope: dict[str, Any] | None = None, extensions: Any | None = None, execution_capacity: SubagentExecutionCapacity | None = None, acceptance_criteria: list[str] | None = None, @@ -830,6 +835,8 @@ class SubagentExecutor: from the parent run for Langfuse metadata correlation. Falls back to the ambient trace so the attribute is always a real id, never ``None``. + knowledge_scope: Canonical execution-only knowledge scope inherited + from the parent turn. Display labels are never propagated. extensions: The parent run's immutable ``LoadedExtensions`` snapshot, captured at ``task_tool`` dispatch. When None (embedded client, standalone LangGraph Server), ``_aexecute`` falls back to the @@ -893,6 +900,7 @@ class SubagentExecutor: # trace contract, and ``_aexecute`` rebinds it because a subagent runs # on the isolated loop thread where the parent ContextVar may be gone. self.deerflow_trace_id = resolve_trace_id(deerflow_trace_id) + self.knowledge_scope = execution_scope(knowledge_scope) if knowledge_scope is not None else None # Parent run's extension snapshot. Binding it here (rather than reading # the singleton at execution time) is what keeps one run on a single # extension generation: a concurrent ``set_loaded_extensions()`` between @@ -1562,6 +1570,8 @@ class SubagentExecutor: context["is_internal"] = self.is_internal context["authz_attributes"] = dict(self.authz_attributes) context[DEERFLOW_TRACE_METADATA_KEY] = self.deerflow_trace_id + if self.knowledge_scope is not None: + context[KNOWLEDGE_SCOPE_RUNTIME_KEY] = dict(self.knowledge_scope) context["is_subagent"] = True context[_SANDBOX_LEASE_OWNER_CONTEXT_KEY] = sandbox_lease_owner_id context[_SANDBOX_COMMAND_SCOPE_CONTEXT_KEY] = sandbox_lease_owner_id diff --git a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py index 6ba9288ba..489d4c2bc 100644 --- a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py @@ -14,6 +14,7 @@ from langgraph.types import Command from pydantic import BaseModel, Field from deerflow.authz.principal import normalize_authz_attributes +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.subagents.batch_runtime import ( BatchItemInput, @@ -212,6 +213,8 @@ async def batch_task( "is_internal": context.get("is_internal") is True, "authz_attributes": normalize_authz_attributes(context.get("authz_attributes")), } + if KNOWLEDGE_SCOPE_RUNTIME_KEY in context: + execution_spec["knowledge_scope"] = execution_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY]) try: batch = await submitter.submit( BatchSubmitRequest( diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index 13951b40c..be5990b75 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -21,6 +21,7 @@ from deerflow.agents.middlewares.receipt_verification import verify_receipt_cita from deerflow.authz.principal import normalize_authz_attributes from deerflow.config import get_app_config from deerflow.extensions import resolve_run_extensions +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.sandbox.security import LOCAL_BASH_SUBAGENT_DISABLED_MESSAGE, is_host_bash_allowed from deerflow.subagents import SubagentExecutor, get_available_subagent_names, get_subagent_config @@ -867,6 +868,9 @@ async def task_tool( # runtime context is authoritative (worker._bind_trace_id always fills it); # the ambient fallback covers tools invoked outside a Gateway run. deerflow_trace_id = resolve_trace_id(parent_context.get(DEERFLOW_TRACE_METADATA_KEY)) + knowledge_scope = None + if KNOWLEDGE_SCOPE_RUNTIME_KEY in parent_context: + knowledge_scope = execution_scope(parent_context[KNOWLEDGE_SCOPE_RUNTIME_KEY]) parent_available_skills = metadata.get("available_skills") if parent_available_skills is not None: @@ -922,6 +926,7 @@ async def task_tool( "is_internal": is_internal, "authz_attributes": authz_attributes, "deerflow_trace_id": deerflow_trace_id, + "knowledge_scope": knowledge_scope, # RFC #4651 PR3: lead-supplied acceptance criteria are handed to the # executor, which appends them to the subagent's task HumanMessage as # untrusted data (sanitized and boundary-framed by diff --git a/backend/packages/harness/deerflow/tools/tools.py b/backend/packages/harness/deerflow/tools/tools.py index 11a791185..892c9c204 100644 --- a/backend/packages/harness/deerflow/tools/tools.py +++ b/backend/packages/harness/deerflow/tools/tools.py @@ -106,6 +106,13 @@ def get_available_tools( if not include_conversation_reader: tool_configs = [tool for tool in tool_configs if tool.use != CONVERSATION_TOOL_USE] + # Knowledge tools are opt-in as a group. Provider connection and retrieval + # settings live on each tool entry; the generic capability flag controls + # whether the group is exposed at all. + knowledge_base_config = getattr(config, "knowledge_base", None) + if not getattr(knowledge_base_config, "enabled", False): + tool_configs = [tool for tool in tool_configs if tool.group != "knowledge"] + # Do not expose host bash by default when LocalSandboxProvider is active. if not is_host_bash_allowed(config): tool_configs = [tool for tool in tool_configs if not _is_host_bash_tool(tool)] diff --git a/backend/tests/test_batch_task_tool.py b/backend/tests/test_batch_task_tool.py index 75d31cba3..40dd99272 100644 --- a/backend/tests/test_batch_task_tool.py +++ b/backend/tests/test_batch_task_tool.py @@ -20,6 +20,11 @@ def _runtime(): "run_id": "run-1", "user_id": "user-1", "user_role": "member", + "__knowledge_scope_execution": { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + }, }, config={ "metadata": { @@ -81,6 +86,11 @@ async def test_batch_task_is_explicit_idempotent_submission(monkeypatch) -> None assert [item["key"] for item in request.items] == ["record-1", "record-2"] assert request.max_live_items == 20 assert request.max_running_items == 5 + assert request.execution_spec["knowledge_scope"] == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + } assert message.additional_kwargs["subagent_batch_id"] == "subagent-batch-1" assert "running independently" in message.content diff --git a/backend/tests/test_config_version.py b/backend/tests/test_config_version.py index 7c4f734f5..2b31b7c15 100644 --- a/backend/tests/test_config_version.py +++ b/backend/tests/test_config_version.py @@ -12,6 +12,22 @@ import yaml from support.shell import find_script_bash from deerflow.config.app_config import AppConfig +from deerflow.config.knowledge_base_config import KnowledgeBaseConfig +from deerflow.tools.tools import get_available_tools + + +def test_knowledge_base_config_is_provider_agnostic() -> None: + assert set(KnowledgeBaseConfig.model_fields) == {"enabled", "scope_selection_enabled"} + config = KnowledgeBaseConfig.model_validate( + { + "enabled": True, + "scope_selection_enabled": True, + "base_url": "http://legacy-ragflow.test", + "api_key": "legacy-secret", + } + ) + assert config.model_dump() == {"enabled": True, "scope_selection_enabled": True} + # Only the upgrade-script test shells out; it needs Git Bash on Windows (the # WSL launcher and Store alias stubs cannot run the repo scripts). @@ -184,6 +200,260 @@ def test_version_26_config_upgrades_to_checkpoint_channel_mode(tmp_path, caplog) assert upgraded["verification"]["judge_model_name"] is None +def test_version_41_config_moves_legacy_ragflow_settings_to_tool(tmp_path): + """The v46 migration keeps provider settings on the RAGFlow tool entry.""" + import subprocess + + repo_root = Path(__file__).resolve().parents[2] + example_src = repo_root / "config.example.yaml" + expected_version = yaml.safe_load(example_src.read_text(encoding="utf-8"))["config_version"] + assert expected_version >= 46 + + config_path = tmp_path / "config.yaml" + legacy = { + "config_version": 41, + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "knowledge_base": { + "enabled": True, + "scope_selection_enabled": True, + "base_url": "http://legacy-ragflow:9380", + "api_key": "$LEGACY_RAGFLOW_API_KEY", + "page_size": 12, + }, + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.ragflow.tools:knowledge_search_tool", + # Explicit tool values win over the legacy global value. + "api_key": "$CURRENT_RAGFLOW_API_KEY", + } + ], + } + config_path.write_text(yaml.dump(legacy), encoding="utf-8") + + env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)} + result = subprocess.run( + ["bash", str(repo_root / "scripts" / "config-upgrade.sh")], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + + upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert upgraded["config_version"] == expected_version, result.stdout + result.stderr + assert upgraded["knowledge_base"] == { + "enabled": True, + "scope_selection_enabled": True, + } + tool = upgraded["tools"][0] + assert tool["base_url"] == "http://legacy-ragflow:9380" + assert tool["page_size"] == 12 + assert tool["api_key"] == "$CURRENT_RAGFLOW_API_KEY" + + +def test_version_41_tools_only_ragflow_config_enables_knowledge_capability(tmp_path): + """Tools-only legacy configs must not be disabled by the new capability gate.""" + import subprocess + + repo_root = Path(__file__).resolve().parents[2] + example_src = repo_root / "config.example.yaml" + expected_version = yaml.safe_load(example_src.read_text(encoding="utf-8"))["config_version"] + assert expected_version >= 46 + + config_path = tmp_path / "config.yaml" + legacy = { + "config_version": 41, + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + # This was the documented enablement path before knowledge_base existed. + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.ragflow.tools:knowledge_search_tool", + "base_url": "http://legacy-ragflow:9380", + "api_key": "$RAGFLOW_API_KEY", + } + ], + } + config_path.write_text(yaml.dump(legacy), encoding="utf-8") + + env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)} + result = subprocess.run( + ["bash", str(repo_root / "scripts" / "config-upgrade.sh")], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + assert "knowledge_base.enabled set to true" in result.stdout + + upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert upgraded["config_version"] == expected_version + assert upgraded["knowledge_base"] == { + "enabled": True, + "scope_selection_enabled": False, + } + assert upgraded["tools"][0]["base_url"] == "http://legacy-ragflow:9380" + + +def test_version_45_tools_only_ragflow_config_runs_knowledge_migration(tmp_path): + """The knowledge migration must run for configs at the former base version.""" + import subprocess + + repo_root = Path(__file__).resolve().parents[2] + example_src = repo_root / "config.example.yaml" + expected_version = yaml.safe_load(example_src.read_text(encoding="utf-8"))["config_version"] + assert expected_version > 45 + + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.dump( + { + "config_version": 45, + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.ragflow.tools:knowledge_search_tool", + "base_url": "http://legacy-ragflow:9380", + "api_key": "$RAGFLOW_API_KEY", + } + ], + } + ), + encoding="utf-8", + ) + + env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)} + result = subprocess.run( + [SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + assert "knowledge_base.enabled set to true" in result.stdout + + upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert upgraded["config_version"] == expected_version + assert upgraded["knowledge_base"] == { + "enabled": True, + "scope_selection_enabled": False, + } + assert upgraded["tools"][0]["base_url"] == "http://legacy-ragflow:9380" + + +def test_version_45_tools_only_lightrag_config_keeps_knowledge_tool_available(tmp_path): + """Upgrading a configured LightRAG provider must enable the new knowledge gate.""" + import subprocess + + repo_root = Path(__file__).resolve().parents[2] + example_src = repo_root / "config.example.yaml" + expected_version = yaml.safe_load(example_src.read_text(encoding="utf-8"))["config_version"] + assert expected_version > 45 + + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.dump( + { + "config_version": 45, + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.lightrag.tools:knowledge_search_tool", + "base_url": "http://legacy-lightrag:9621", + "api_key": "$LIGHTRAG_API_KEY", + } + ], + } + ), + encoding="utf-8", + ) + + env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)} + result = subprocess.run( + [SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + assert "knowledge_base.enabled set to true" in result.stdout + + upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert upgraded["config_version"] == expected_version + assert upgraded["knowledge_base"] == { + "enabled": True, + "scope_selection_enabled": False, + } + + app_config = AppConfig.model_validate(upgraded) + tools = get_available_tools( + groups=["knowledge"], + include_mcp=False, + include_upload_tool=False, + app_config=app_config, + ) + assert "knowledge_search" in {tool.name for tool in tools} + + +def test_version_45_lightrag_config_preserves_explicit_disabled_gate(tmp_path): + """Migration must not override an operator's explicit knowledge gate value.""" + import subprocess + + repo_root = Path(__file__).resolve().parents[2] + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.dump( + { + "config_version": 45, + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "knowledge_base": {"enabled": False}, + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.lightrag.tools:knowledge_search_tool", + "base_url": "http://legacy-lightrag:9621", + } + ], + } + ), + encoding="utf-8", + ) + + env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)} + result = subprocess.run( + [SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + + upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert upgraded["knowledge_base"]["enabled"] is False + + app_config = AppConfig.model_validate(upgraded) + tools = get_available_tools( + groups=["knowledge"], + include_mcp=False, + include_upload_tool=False, + app_config=app_config, + ) + assert "knowledge_search" not in {tool.name for tool in tools} + + def _load_repo_example() -> dict: """Load the real repo config.example.yaml (first-run template).""" example_path = Path(__file__).resolve().parents[2] / "config.example.yaml" diff --git a/backend/tests/test_features_router.py b/backend/tests/test_features_router.py index 6b228f812..70a33d0e0 100644 --- a/backend/tests/test_features_router.py +++ b/backend/tests/test_features_router.py @@ -1,6 +1,7 @@ from types import SimpleNamespace from unittest.mock import patch +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient @@ -17,6 +18,9 @@ def _app_with_config( subagent_batches_available: bool = False, subagent_batch_repo_available: bool | None = None, conversation_references_enabled: bool = False, + knowledge_base_enabled: bool = False, + scope_selection_enabled: bool = False, + knowledge_search_provider: str | None = None, ) -> FastAPI: app = FastAPI() app.state.mcp_tasks_available = mcp_tasks_available @@ -34,7 +38,13 @@ def _app_with_config( agents_api=SimpleNamespace(enabled=agents_api_enabled), tools=tools, subagent_runtime=SimpleNamespace(max_running=3), + knowledge_base=SimpleNamespace( + enabled=knowledge_base_enabled, + scope_selection_enabled=scope_selection_enabled, + ), ) + search_tool = SimpleNamespace(use=knowledge_search_provider) if knowledge_search_provider is not None else None + fake_config.get_tool_config = lambda name: search_tool if name == "knowledge_search" else None app.dependency_overrides[get_config] = lambda: fake_config return app @@ -54,6 +64,9 @@ def test_features_reports_agents_api_enabled() -> None: "max_running": 3, }, "conversation_references": {"enabled": False, "max_references": 3}, + "knowledge_base": { + "scope_selection_enabled": False, + }, } @@ -72,6 +85,9 @@ def test_features_reports_agents_api_disabled() -> None: "max_running": 3, }, "conversation_references": {"enabled": False, "max_references": 3}, + "knowledge_base": { + "scope_selection_enabled": False, + }, } @@ -82,6 +98,48 @@ def test_features_reports_conversation_references_when_the_tool_is_configured() assert response.json()["conversation_references"] == {"enabled": True, "max_references": 3} +def test_features_enables_scope_selection_only_for_exact_ragflow_provider() -> None: + with TestClient( + _app_with_config( + agents_api_enabled=True, + knowledge_base_enabled=True, + scope_selection_enabled=True, + knowledge_search_provider=("deerflow.community.ragflow.tools:knowledge_search_tool"), + ) + ) as client: + response = client.get("/api/features") + + assert response.status_code == 200 + assert response.json()["knowledge_base"]["scope_selection_enabled"] is True + + +@pytest.mark.parametrize( + ("knowledge_base_enabled", "provider"), + [ + (False, "deerflow.community.ragflow.tools:knowledge_search_tool"), + (True, "deerflow.community.lightrag.tools:knowledge_search_tool"), + (True, "custom.provider:knowledge_search_tool"), + (True, None), + ], +) +def test_features_scope_selection_fails_closed( + knowledge_base_enabled: bool, + provider: str | None, +) -> None: + with TestClient( + _app_with_config( + agents_api_enabled=True, + knowledge_base_enabled=knowledge_base_enabled, + scope_selection_enabled=True, + knowledge_search_provider=provider, + ) + ) as client: + response = client.get("/api/features") + + assert response.status_code == 200 + assert response.json()["knowledge_base"]["scope_selection_enabled"] is False + + def test_features_reports_mcp_tasks_startup_capability() -> None: with TestClient(_app_with_config(agents_api_enabled=True, mcp_tasks_available=True)) as client: response = client.get("/api/features") diff --git a/backend/tests/test_gateway_knowledge_scope_admission.py b/backend/tests/test_gateway_knowledge_scope_admission.py new file mode 100644 index 000000000..4367efc96 --- /dev/null +++ b/backend/tests/test_gateway_knowledge_scope_admission.py @@ -0,0 +1,212 @@ +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException +from langchain_core.messages import AIMessage, HumanMessage + +from app.gateway.knowledge_scope_admission import admit_message_knowledge_scope +from app.gateway.services import strip_internal_context_keys +from deerflow.config.tool_config import ToolConfig +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY + + +def _app_config(provider: str = "deerflow.community.ragflow.tools:knowledge_search_tool"): + tool_config = ToolConfig( + name="knowledge_search", + group="knowledge", + use=provider, + ) + return SimpleNamespace( + knowledge_base=SimpleNamespace(enabled=True), + get_tool_config=lambda name: tool_config if name == "knowledge_search" else None, + ) + + +def _agent_config(tool_groups=None): + return SimpleNamespace(tool_groups=tool_groups) + + +def _input(message): + return {"messages": [message]} + + +def test_admission_canonicalizes_custom_agent_human_message() -> None: + graph_input = _input( + HumanMessage( + content="question", + additional_kwargs={ + KNOWLEDGE_SCOPE_KEY: { + "version": 1, + "mode": "selected", + "dataset_ids": [" dataset-a ", "dataset-a"], + } + }, + ) + ) + + admitted = admit_message_knowledge_scope( + graph_input, + assistant_id="agriculture-agent", + app_config=_app_config(), + agent_config=_agent_config(None), + ) + + assert admitted == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + } + assert graph_input["messages"][0].additional_kwargs[KNOWLEDGE_SCOPE_KEY] == admitted + + +@pytest.mark.parametrize( + ("assistant_id", "provider", "tool_groups"), + [ + (None, "deerflow.community.ragflow.tools:knowledge_search_tool", None), + ("agent", "deerflow.community.lightrag.tools:knowledge_search_tool", None), + ("agent", "deerflow.community.ragflow.tools:knowledge_search_tool", []), + ("agent", "deerflow.community.ragflow.tools:knowledge_search_tool", ["web"]), + ], +) +def test_scope_is_rejected_outside_supported_custom_agent( + assistant_id: str | None, + provider: str, + tool_groups: list[str] | None, +) -> None: + graph_input = _input( + HumanMessage( + content="question", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ) + ) + + with pytest.raises(HTTPException) as exc_info: + admit_message_knowledge_scope( + graph_input, + assistant_id=assistant_id, + app_config=_app_config(provider), + agent_config=_agent_config(tool_groups) if assistant_id not in {None, "lead_agent"} else None, + ) + + assert exc_info.value.status_code == 422 + + +def test_admission_accepts_main_agent_with_configured_ragflow_provider() -> None: + graph_input = _input( + HumanMessage( + content="question", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ) + ) + + admitted = admit_message_knowledge_scope( + graph_input, + assistant_id="lead_agent", + app_config=_app_config(), + agent_config=None, + ) + + assert admitted == {"version": 1, "mode": "all"} + + +def test_scope_on_non_human_or_multiple_humans_is_rejected() -> None: + invalid_ai = _input( + AIMessage( + content="answer", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ) + ) + with pytest.raises(HTTPException, match="HumanMessage"): + admit_message_knowledge_scope( + invalid_ai, + assistant_id="agent", + app_config=_app_config(), + agent_config=_agent_config(), + ) + + duplicate = { + "messages": [ + HumanMessage( + content="one", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ), + HumanMessage( + content="two", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ), + ] + } + with pytest.raises(HTTPException, match="one new HumanMessage"): + admit_message_knowledge_scope( + duplicate, + assistant_id="agent", + app_config=_app_config(), + agent_config=_agent_config(), + ) + + +def test_scope_on_an_earlier_human_message_is_rejected() -> None: + graph_input = { + "messages": [ + HumanMessage( + content="historical", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ), + HumanMessage(content="current"), + ] + } + + with pytest.raises(HTTPException, match="current HumanMessage"): + admit_message_knowledge_scope( + graph_input, + assistant_id="agent", + app_config=_app_config(), + agent_config=_agent_config(), + ) + + +def test_recovery_scope_replaces_client_forgery_and_legacy_removes_it() -> None: + graph_input = _input( + HumanMessage( + content="question", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}}, + ) + ) + recovered = admit_message_knowledge_scope( + graph_input, + assistant_id="agent", + app_config=_app_config(), + agent_config=_agent_config(), + recovery_scope={"version": 1, "mode": "disabled"}, + recovery=True, + ) + assert recovered == {"version": 1, "mode": "disabled"} + assert graph_input["messages"][0].additional_kwargs[KNOWLEDGE_SCOPE_KEY] == recovered + + admitted = admit_message_knowledge_scope( + graph_input, + assistant_id="agent", + app_config=_app_config(), + agent_config=_agent_config(), + recovery_scope=None, + recovery=True, + ) + assert admitted is None + assert KNOWLEDGE_SCOPE_KEY not in graph_input["messages"][0].additional_kwargs + + +def test_free_form_runtime_scope_fields_are_scrubbed() -> None: + config = { + "context": { + KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}, + KNOWLEDGE_SCOPE_RUNTIME_KEY: {"version": 1, "mode": "disabled"}, + }, + "configurable": { + KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}, + KNOWLEDGE_SCOPE_RUNTIME_KEY: {"version": 1, "mode": "disabled"}, + }, + } + + strip_internal_context_keys(config) + + assert config == {"context": {}, "configurable": {}} diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index faf0addce..3207634c4 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -259,6 +259,47 @@ def test_normalize_input_preserves_additional_kwargs_and_id(): assert msg.additional_kwargs == {"files": files, "custom": "keep-me"} +def test_canonical_run_record_input_uses_admitted_message_snapshot(): + from langchain_core.messages import HumanMessage + + from app.gateway.services import _canonical_run_record_input + + raw = { + "messages": [ + { + "role": "user", + "content": "Search", + "additional_kwargs": { + "knowledge_scope": { + "version": 1, + "mode": "selected", + "dataset_ids": [" dataset-1 ", "dataset-1"], + } + }, + } + ] + } + admitted = { + "messages": [ + HumanMessage( + content="Search", + additional_kwargs={ + "knowledge_scope": { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + } + }, + ) + ] + } + + stored = _canonical_run_record_input(raw, admitted) + + assert stored is not None + assert stored["messages"][0]["additional_kwargs"]["knowledge_scope"]["dataset_ids"] == ["dataset-1"] + + @pytest.mark.parametrize( "forged_original", ["spoofed audit text", [{"type": "text", "text": "spoofed audit text"}]], @@ -295,6 +336,7 @@ def test_normalize_input_strips_external_dynamic_context_metadata(): """ from app.gateway.services import normalize_input from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY + from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY result = normalize_input( { @@ -307,6 +349,8 @@ def test_normalize_input_strips_external_dynamic_context_metadata(): "hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True, _REMINDER_DATE_KEY: "2099-01-01, Thursday", + KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}, + KNOWLEDGE_SCOPE_RUNTIME_KEY: {"version": 1, "mode": "disabled"}, "custom": "keep-me", }, } @@ -317,7 +361,12 @@ def test_normalize_input_strips_external_dynamic_context_metadata(): from deerflow.utils.messages import UNTRUSTED_INPUT_KEY assert result["messages"][0].id == "known-checkpoint-id__memory" - assert result["messages"][0].additional_kwargs == {"hide_from_ui": True, "custom": "keep-me", UNTRUSTED_INPUT_KEY: True} + assert result["messages"][0].additional_kwargs == { + "hide_from_ui": True, + KNOWLEDGE_SCOPE_KEY: {"version": 1, "mode": "all"}, + "custom": "keep-me", + UNTRUSTED_INPUT_KEY: True, + } def test_normalize_input_strips_external_view_image_context_marker(): @@ -2076,6 +2125,264 @@ async def _capture_start_run_graph_input(body, *, auth_source=None): return captured["graph_input"] +@pytest.mark.parametrize( + "target_message_id", + ["assistant-answer", "missing-assistant", None], + ids=["regenerate", "interrupted-regenerate-fallback", "resume-fallback"], +) +@pytest.mark.asyncio +async def test_recover_knowledge_scope_skips_hidden_conversation_reference_message(target_message_id): + from unittest.mock import AsyncMock, patch + + from langchain_core.messages import AIMessage, HumanMessage + + from app.gateway.services import _recover_run_knowledge_scope + + source_scope = { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-source"], + } + messages = [ + HumanMessage( + id="user-source", + content="Search only the selected dataset", + additional_kwargs={"knowledge_scope": source_scope}, + ), + HumanMessage( + id="conversation-references", + content='Read-only conversation references for this run: ["thread-source"]', + additional_kwargs={"hide_from_ui": True}, + ), + AIMessage(id="assistant-answer", content="Scoped answer"), + ] + accessor = SimpleNamespace( + aget=AsyncMock(return_value=SimpleNamespace(values={"messages": messages})), + ) + + with patch( + "app.gateway.services.build_thread_checkpoint_state_accessor", + new=AsyncMock(return_value=(accessor, {})), + ): + recovered = await _recover_run_knowledge_scope( + SimpleNamespace(), + thread_id="thread-scope-recovery", + target_message_id=target_message_id, + ) + + assert recovered == source_scope + + +@pytest.mark.parametrize( + ("include_current_scope", "expected_scope", "recovery_calls"), + [ + ( + True, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-current"], + }, + 0, + ), + (False, {"version": 1, "mode": "disabled"}, 1), + ], + ids=["current-selection-wins", "omitted-selection-recovers"], +) +@pytest.mark.asyncio +async def test_clarification_reply_scope_uses_current_selection_or_recovers_when_omitted( + _stub_app_config, + include_current_scope, + expected_scope, + recovery_calls, +): + from unittest.mock import AsyncMock, patch + + from app.gateway.routers.thread_runs import RunCreateRequest + from app.gateway.services import start_run + from deerflow.runtime import RunManager + from deerflow.runtime.runs.store.memory import MemoryRunStore + + set_app_config( + AppConfig.model_validate( + { + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "knowledge_base": {"enabled": True, "scope_selection_enabled": True}, + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.ragflow.tools:knowledge_search_tool", + } + ], + } + ) + ) + response_metadata = { + "version": 1, + "kind": "human_input_response", + "source": "ask_clarification", + "request_id": "clarification:call-scope", + "response_kind": "text", + "value": "Use the current dataset", + } + additional_kwargs = { + "hide_from_ui": True, + "human_input_response": response_metadata, + } + if include_current_scope: + additional_kwargs["knowledge_scope"] = { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-current"], + } + + body = RunCreateRequest( + assistant_id="researcher", + input={ + "messages": [ + { + "type": "human", + "content": "Use the current dataset", + "additional_kwargs": additional_kwargs, + } + ] + }, + ) + request = _make_start_run_request(RunManager(store=MemoryRunStore())) + captured: dict[str, object] = {} + recover_scope = AsyncMock(return_value={"version": 1, "mode": "disabled"}) + + async def fake_run_agent(*_args, **kwargs): + captured["graph_input"] = kwargs["graph_input"] + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=object()), + patch("app.gateway.services.run_agent", side_effect=fake_run_agent), + patch("app.gateway.services._recover_run_knowledge_scope", new=recover_scope), + patch( + "app.gateway.services._load_scope_agent_config", + new=AsyncMock(return_value=SimpleNamespace(tool_groups=["knowledge"])), + ), + ): + record = await start_run(body, "thread-clarification-scope", request) + await record.task + + graph_input = captured["graph_input"] + assert isinstance(graph_input, dict) + message = graph_input["messages"][0] + assert message.additional_kwargs["knowledge_scope"] == expected_scope + assert recover_scope.await_count == recovery_calls + + +@pytest.mark.parametrize( + ("include_current_scope", "expected_scope", "recovery_calls"), + [ + ( + True, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-current"], + }, + 0, + ), + ( + False, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-source"], + }, + 1, + ), + ], + ids=["current-selection-wins", "omitted-selection-recovers"], +) +@pytest.mark.asyncio +async def test_edit_replay_scope_uses_current_selection_or_recovers_when_omitted( + _stub_app_config, + include_current_scope, + expected_scope, + recovery_calls, +): + from unittest.mock import AsyncMock, patch + + from app.gateway.routers.thread_runs import RunCreateRequest + from app.gateway.services import start_run + from deerflow.runtime import RunManager + from deerflow.runtime.runs.store.memory import MemoryRunStore + + set_app_config( + AppConfig.model_validate( + { + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "knowledge_base": {"enabled": True, "scope_selection_enabled": True}, + "tools": [ + { + "name": "knowledge_search", + "group": "knowledge", + "use": "deerflow.community.ragflow.tools:knowledge_search_tool", + } + ], + } + ) + ) + additional_kwargs = {} + if include_current_scope: + additional_kwargs["knowledge_scope"] = { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-current"], + } + body = RunCreateRequest( + assistant_id="researcher", + input={ + "messages": [ + { + "type": "human", + "content": "Edited question", + "additional_kwargs": additional_kwargs, + } + ] + }, + metadata={ + "replay_kind": "edit", + "regenerate_from_message_id": "assistant-source", + }, + ) + request = _make_start_run_request(RunManager(store=MemoryRunStore())) + captured: dict[str, object] = {} + recover_scope = AsyncMock( + return_value={ + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-source"], + } + ) + + async def fake_run_agent(*_args, **kwargs): + captured["graph_input"] = kwargs["graph_input"] + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=object()), + patch("app.gateway.services.run_agent", side_effect=fake_run_agent), + patch("app.gateway.services._recover_run_knowledge_scope", new=recover_scope), + patch( + "app.gateway.services._load_scope_agent_config", + new=AsyncMock(return_value=SimpleNamespace(tool_groups=["knowledge"])), + ), + ): + record = await start_run(body, "thread-edit-scope", request) + await record.task + + graph_input = captured["graph_input"] + assert isinstance(graph_input, dict) + message = graph_input["messages"][0] + assert message.additional_kwargs["knowledge_scope"] == expected_scope + assert recover_scope.await_count == recovery_calls + + def _make_start_run_persistence_context(): from types import SimpleNamespace diff --git a/backend/tests/test_knowledge_router.py b/backend/tests/test_knowledge_router.py new file mode 100644 index 000000000..6aa1f033d --- /dev/null +++ b/backend/tests/test_knowledge_router.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from _router_auth_helpers import make_authed_test_app +from fastapi.testclient import TestClient + +from app.gateway.auth.models import User +from app.gateway.deps import get_config +from app.gateway.routers import knowledge + + +def _config( + *, + enabled: bool = True, + api_key: str | None = "ragflow-secret", + scope_selection_enabled: bool = False, + datasets: list[str] | None = None, + provider: str = "deerflow.community.ragflow.tools:knowledge_search_tool", +) -> SimpleNamespace: + tool = SimpleNamespace( + use=provider, + model_extra={ + "base_url": "http://ragflow.test", + "api_key": api_key, + "timeout": 30, + **({"datasets": datasets} if datasets is not None else {}), + }, + ) + return SimpleNamespace( + knowledge_base=SimpleNamespace( + enabled=enabled, + scope_selection_enabled=scope_selection_enabled, + ), + get_tool_config=lambda name: tool if name == "knowledge_search" else None, + ) + + +def _user() -> User: + return User( + email="router-test@example.com", + password_hash="x", + system_role="user", + ) + + +def _app( + monkeypatch: pytest.MonkeyPatch, + client: object, + *, + config: SimpleNamespace | None = None, +): + app = make_authed_test_app(user_factory=_user) + app.include_router(knowledge.router) + app.dependency_overrides[get_config] = lambda: config or _config() + monkeypatch.setattr( + knowledge, + "_build_retrieval_client", + lambda settings: client, + ) + return app + + +def _enable_scope_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + knowledge, + "load_agent_config", + lambda name, *, user_id: SimpleNamespace( + name=name, + tool_groups=["knowledge"], + ), + ) + + +def test_retrieval_catalog_enforces_allowlist_and_normalizes_pages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _enable_scope_catalog(monkeypatch) + ragflow = SimpleNamespace( + list_datasets=AsyncMock( + side_effect=[ + [ + { + "id": "dataset-1", + "name": "Policies", + "embedding_model": "embed-a", + "chunk_count": 3, + } + ], + [ + { + "id": "dataset-2", + "name": "Empty", + "embedding_model": "", + "chunk_count": 0, + } + ], + ] + ) + ) + config = _config( + scope_selection_enabled=True, + datasets=["dataset-1", "dataset-2"], + ) + + with TestClient(_app(monkeypatch, ragflow, config=config)) as client: + response = client.get( + "/api/knowledge/retrieval-catalog/datasets", + params={"agent_name": "researcher", "page": 1, "page_size": 20}, + ) + + assert response.status_code == 200 + assert response.json() == { + "items": [ + {"id": "dataset-1", "name": "Policies", "selectable": True}, + {"id": "dataset-2", "name": "Empty", "selectable": False}, + ], + "page": 1, + "page_size": 20, + "total": 2, + } + assert [call.kwargs["dataset_id"] for call in ragflow.list_datasets.await_args_list] == ["dataset-1", "dataset-2"] + + +def test_retrieval_catalog_accepts_main_assistant( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _enable_scope_catalog(monkeypatch) + ragflow = SimpleNamespace( + list_datasets=AsyncMock( + return_value=[ + { + "id": "dataset-1", + "name": "Policies", + "embedding_model": "embed-a", + "chunk_count": 3, + } + ] + ) + ) + config = _config(scope_selection_enabled=True, datasets=["dataset-1"]) + + with TestClient(_app(monkeypatch, ragflow, config=config)) as client: + response = client.get( + "/api/knowledge/retrieval-catalog/datasets", + params={"agent_name": "lead_agent"}, + ) + + assert response.status_code == 200 + assert response.json()["items"] == [{"id": "dataset-1", "name": "Policies", "selectable": True}] + + +def test_retrieval_catalog_documents_reject_outside_allowlist_without_provider_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _enable_scope_catalog(monkeypatch) + ragflow = SimpleNamespace( + list_datasets=AsyncMock(), + list_documents=AsyncMock(), + ) + config = _config(scope_selection_enabled=True, datasets=["dataset-1"]) + + with TestClient(_app(monkeypatch, ragflow, config=config)) as client: + response = client.get( + "/api/knowledge/retrieval-catalog/datasets/dataset-2/documents", + params={"agent_name": "researcher"}, + ) + + assert response.status_code == 404 + ragflow.list_datasets.assert_not_awaited() + ragflow.list_documents.assert_not_awaited() + + +def test_retrieval_catalog_documents_marks_only_searchable_files_selectable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _enable_scope_catalog(monkeypatch) + ragflow = SimpleNamespace( + list_datasets=AsyncMock( + return_value=[ + { + "id": "dataset-1", + "name": "Policies", + "embedding_model": "embed-a", + "chunk_count": 3, + } + ] + ), + list_documents=AsyncMock( + return_value={ + "code": 0, + "data": { + "total": 3, + "docs": [ + { + "id": "doc-1", + "name": "Ready.pdf", + "run": "DONE", + "chunk_count": 2, + }, + { + "id": "doc-2", + "name": "Parsing.pdf", + "run": "RUNNING", + "chunk_count": 0, + }, + { + "id": "doc-3", + "name": "Empty.pdf", + "run": "DONE", + "chunk_count": 0, + }, + ], + }, + } + ), + ) + config = _config(scope_selection_enabled=True, datasets=["dataset-1"]) + + with TestClient(_app(monkeypatch, ragflow, config=config)) as client: + response = client.get( + "/api/knowledge/retrieval-catalog/datasets/dataset-1/documents", + params={ + "agent_name": "researcher", + "search": "ready", + "page": 2, + "page_size": 10, + }, + ) + + assert response.status_code == 200 + assert response.json()["items"] == [ + {"id": "doc-1", "name": "Ready.pdf", "selectable": True}, + {"id": "doc-2", "name": "Parsing.pdf", "selectable": False}, + {"id": "doc-3", "name": "Empty.pdf", "selectable": False}, + ] + ragflow.list_documents.assert_awaited_once_with( + "dataset-1", + params=[("page", "2"), ("page_size", "10"), ("keywords", "ready")], + ) + + +@pytest.mark.parametrize( + "config", + [ + _config(enabled=False, scope_selection_enabled=True), + _config(scope_selection_enabled=False), + _config( + scope_selection_enabled=True, + provider=("deerflow.community.lightrag.tools:knowledge_search_tool"), + ), + ], +) +def test_retrieval_catalog_fails_closed_when_capability_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + config: SimpleNamespace, +) -> None: + _enable_scope_catalog(monkeypatch) + ragflow = SimpleNamespace(list_datasets=AsyncMock()) + + with TestClient(_app(monkeypatch, ragflow, config=config)) as client: + response = client.get( + "/api/knowledge/retrieval-catalog/datasets", + params={"agent_name": "researcher"}, + ) + + assert response.status_code == 409 + ragflow.list_datasets.assert_not_awaited() + + +def test_management_routes_are_not_exposed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with TestClient(_app(monkeypatch, SimpleNamespace())) as client: + assert client.get("/api/knowledge/datasets").status_code == 404 + assert ( + client.post( + "/api/knowledge/datasets", + json={"name": "Deferred"}, + ).status_code + == 404 + ) + assert client.get("/api/knowledge/events").status_code == 404 diff --git a/backend/tests/test_knowledge_scope.py b/backend/tests/test_knowledge_scope.py new file mode 100644 index 000000000..5eaa06d32 --- /dev/null +++ b/backend/tests/test_knowledge_scope.py @@ -0,0 +1,239 @@ +import json + +import pytest +from pydantic import ValidationError + +from deerflow.knowledge_scope import ( + KNOWLEDGE_SCOPE_KEY, + KnowledgeScope, + canonicalize_knowledge_scope, + execution_scope, +) + + +def test_all_and_disabled_canonicalize_without_selection_fields() -> None: + assert canonicalize_knowledge_scope({"version": 1, "mode": "all"}) == { + "version": 1, + "mode": "all", + } + assert canonicalize_knowledge_scope({"version": 1, "mode": "disabled"}) == { + "version": 1, + "mode": "disabled", + } + + +@pytest.mark.parametrize( + "payload", + [ + {"version": 2, "mode": "all"}, + {"version": 1, "mode": "all", "future": True}, + {"version": 1, "mode": "all", "dataset_ids": ["dataset-a"]}, + {"version": 1, "mode": "disabled", "display": {"datasets": []}}, + {"version": 1, "mode": "selected", "dataset_ids": []}, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [{"dataset_id": "dataset-a", "document_ids": []}], + }, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [ + {"dataset_id": "dataset-a", "document_ids": ["doc-a"]}, + {"dataset_id": "dataset-a", "document_ids": ["doc-b"]}, + ], + }, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [{"dataset_id": "dataset-b", "document_ids": ["doc-a"]}], + }, + ], +) +def test_invalid_scope_shapes_are_rejected(payload: dict) -> None: + with pytest.raises(ValidationError): + canonicalize_knowledge_scope(payload) + + +def test_selected_scope_trims_and_stably_deduplicates_ids() -> None: + scope = canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": [" dataset-a ", "dataset-b", "dataset-a"], + "document_filters": [ + { + "dataset_id": " dataset-b ", + "document_ids": [" doc-1 ", "doc-2", "doc-1"], + } + ], + } + ) + + assert scope == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a", "dataset-b"], + "document_filters": [{"dataset_id": "dataset-b", "document_ids": ["doc-1", "doc-2"]}], + } + + +def test_execution_scope_drops_untrusted_display_snapshot() -> None: + scope = canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "display": {"datasets": [{"id": "dataset-a", "name": "Agriculture", "documents": []}]}, + } + ) + + assert execution_scope(scope) == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + } + assert KNOWLEDGE_SCOPE_KEY == "knowledge_scope" + + +@pytest.mark.parametrize( + "payload", + [ + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "display": {"datasets": [{"id": "dataset-b", "name": "Other"}]}, + }, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [{"dataset_id": "dataset-a", "document_ids": ["doc-a"]}], + "display": { + "datasets": [ + { + "id": "dataset-a", + "name": "Agriculture", + "documents": [{"id": "doc-b", "name": "Other.pdf"}], + } + ] + }, + }, + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "display": { + "datasets": [ + {"id": "dataset-a", "name": "First"}, + {"id": "dataset-a", "name": "Duplicate"}, + ] + }, + }, + ], +) +def test_display_must_be_unique_and_related_to_execution_selection(payload: dict) -> None: + with pytest.raises(ValidationError): + canonicalize_knowledge_scope(payload) + + +def test_scope_capacity_limits_are_rejected_without_truncation() -> None: + with pytest.raises(ValidationError): + canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": [f"dataset-{index}" for index in range(101)], + } + ) + + with pytest.raises(ValidationError): + canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [ + { + "dataset_id": "dataset-a", + "document_ids": [f"doc-{index}" for index in range(1001)], + } + ], + } + ) + + with pytest.raises(ValidationError): + canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": ["x" * 257], + } + ) + + +def test_display_capacity_and_canonical_json_byte_limit_are_enforced() -> None: + with pytest.raises(ValidationError): + canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": [f"dataset-{index}" for index in range(21)], + "display": {"datasets": [{"id": f"dataset-{index}", "name": f"Dataset {index}"} for index in range(21)]}, + } + ) + + oversized = { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [ + { + "dataset_id": "dataset-a", + "document_ids": [f"doc-{index}-{'x' * 240}" for index in range(300)], + } + ], + } + assert len(json.dumps(oversized, ensure_ascii=False).encode()) > 64 * 1024 + with pytest.raises(ValidationError): + canonicalize_knowledge_scope(oversized) + + +def test_model_rejects_more_than_fifty_display_documents_and_long_names() -> None: + document_ids = [f"doc-{index}" for index in range(51)] + with pytest.raises(ValidationError): + KnowledgeScope.model_validate( + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "document_filters": [{"dataset_id": "dataset-a", "document_ids": document_ids}], + "display": { + "datasets": [ + { + "id": "dataset-a", + "name": "Agriculture", + "documents": [{"id": document_id, "name": f"Document {index}"} for index, document_id in enumerate(document_ids)], + } + ] + }, + } + ) + + with pytest.raises(ValidationError): + canonicalize_knowledge_scope( + { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "display": { + "datasets": [ + {"id": "dataset-a", "name": "名" * 257}, + ] + }, + } + ) diff --git a/backend/tests/test_knowledge_scope_middleware.py b/backend/tests/test_knowledge_scope_middleware.py new file mode 100644 index 000000000..b776f93e5 --- /dev/null +++ b/backend/tests/test_knowledge_scope_middleware.py @@ -0,0 +1,181 @@ +from types import SimpleNamespace + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.tools import tool +from langgraph.runtime import Runtime + +from deerflow.agents.middlewares.knowledge_scope_middleware import ( + KnowledgeScopeMiddleware, +) +from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY +from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY + + +@tool +def knowledge_search(query: str) -> str: + """Search.""" + return query + + +@tool +def other_tool(query: str) -> str: + """Other.""" + return query + + +class _ModelRequest: + def __init__(self, messages, *, tools=(), runtime=None): + self.messages = list(messages) + self.tools = list(tools) + self.runtime = runtime + + def override(self, **kwargs): + return _ModelRequest( + kwargs.get("messages", self.messages), + tools=kwargs.get("tools", self.tools), + runtime=self.runtime, + ) + + +def _scope(mode: str = "selected") -> dict: + if mode != "selected": + return {"version": 1, "mode": mode} + return { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + "display": { + "datasets": [{"id": "dataset-a", "name": "Agriculture"}], + }, + } + + +def test_before_agent_projects_only_current_message_execution_scope() -> None: + historical = HumanMessage( + content="old", + id="old", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: _scope("all")}, + ) + current = HumanMessage( + content="new", + id="new", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: _scope()}, + ) + runtime = Runtime(context={CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: frozenset({"old"})}) + + KnowledgeScopeMiddleware().before_agent( + {"messages": [historical, current]}, + runtime, + ) + + assert runtime.context[KNOWLEDGE_SCOPE_RUNTIME_KEY] == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-a"], + } + + +def test_before_agent_uses_server_admitted_runtime_scope_for_recovery() -> None: + runtime = Runtime( + context={ + CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: frozenset({"source"}), + KNOWLEDGE_SCOPE_RUNTIME_KEY: _scope("disabled"), + } + ) + source = HumanMessage( + content="source", + id="source", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: _scope("selected")}, + ) + + KnowledgeScopeMiddleware().before_agent({"messages": [source]}, runtime) + + assert runtime.context[KNOWLEDGE_SCOPE_RUNTIME_KEY] == { + "version": 1, + "mode": "disabled", + } + + +@pytest.mark.parametrize("method_name", ["wrap_model_call", "awrap_model_call"]) +@pytest.mark.anyio +async def test_model_paths_strip_every_historical_scope_and_hide_disabled_tool( + method_name: str, +) -> None: + runtime = Runtime(context={KNOWLEDGE_SCOPE_RUNTIME_KEY: _scope("disabled")}) + messages = [ + HumanMessage( + content="old", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: _scope("all"), "keep": True}, + ), + AIMessage( + content="answer", + additional_kwargs={KNOWLEDGE_SCOPE_KEY: _scope(), "keep": True}, + ), + ] + request = _ModelRequest( + messages, + tools=[knowledge_search, other_tool], + runtime=runtime, + ) + captured = [] + middleware = KnowledgeScopeMiddleware() + + if method_name == "wrap_model_call": + middleware.wrap_model_call( + request, + lambda value: captured.append(value) or "ok", + ) + else: + + async def handler(value): + captured.append(value) + return "ok" + + await middleware.awrap_model_call(request, handler) + + assert [item.name for item in captured[0].tools] == ["other_tool"] + assert all(KNOWLEDGE_SCOPE_KEY not in message.additional_kwargs for message in captured[0].messages) + assert all(message.additional_kwargs["keep"] for message in captured[0].messages) + assert KNOWLEDGE_SCOPE_KEY in request.messages[0].additional_kwargs + + +def test_disabled_execution_guard_blocks_knowledge_tool() -> None: + runtime = Runtime(context={KNOWLEDGE_SCOPE_RUNTIME_KEY: _scope("disabled")}) + request = SimpleNamespace( + tool_call={"name": "knowledge_search", "id": "call-1"}, + runtime=runtime, + ) + + result = KnowledgeScopeMiddleware().wrap_tool_call( + request, + lambda _request: pytest.fail("disabled call must not execute"), + ) + + assert isinstance(result, ToolMessage) + assert result.status == "error" + assert result.tool_call_id == "call-1" + assert "disabled" in str(result.content).lower() + + +def test_legacy_message_without_scope_keeps_tools_and_runtime_unset() -> None: + runtime = Runtime(context={CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: frozenset()}) + middleware = KnowledgeScopeMiddleware() + middleware.before_agent( + {"messages": [HumanMessage(content="legacy", id="new")]}, + runtime, + ) + request = _ModelRequest( + [HumanMessage(content="legacy")], + tools=[knowledge_search], + runtime=runtime, + ) + captured = [] + + middleware.wrap_model_call( + request, + lambda value: captured.append(value) or "ok", + ) + + assert KNOWLEDGE_SCOPE_RUNTIME_KEY not in runtime.context + assert [item.name for item in captured[0].tools] == ["knowledge_search"] diff --git a/backend/tests/test_lightrag_tools.py b/backend/tests/test_lightrag_tools.py index ac3ceda2e..46ff7cb2d 100644 --- a/backend/tests/test_lightrag_tools.py +++ b/backend/tests/test_lightrag_tools.py @@ -362,6 +362,7 @@ def test_tool_assembly_hides_credentials_without_network_io(monkeypatch: pytest. ) config = SimpleNamespace( tools=[tool_config], + knowledge_base=SimpleNamespace(enabled=True), sandbox=SimpleNamespace(use="example.remote:Sandbox"), skill_evolution=SimpleNamespace(enabled=False), models=[], diff --git a/backend/tests/test_ragflow_client.py b/backend/tests/test_ragflow_client.py index 62b96360a..a942ec8ad 100644 --- a/backend/tests/test_ragflow_client.py +++ b/backend/tests/test_ragflow_client.py @@ -166,6 +166,39 @@ async def test_retrieve_rejects_empty_dataset_ids_before_request() -> None: assert called is False +@pytest.mark.anyio +async def test_retrieve_sends_nonempty_document_filter_and_rejects_empty_filter() -> None: + requests: list[dict] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(json.loads(request.content)) + return httpx.Response( + 200, + json={"code": 0, "data": {"chunks": [], "doc_aggs": [], "total": 0}}, + ) + + client = RAGFlowClient( + base_url="http://ragflow.test", + api_key="ragflow-secret", + transport=httpx.MockTransport(handler), + ) + + await client.retrieve( + "annual leave", + dataset_ids=["dataset-1"], + document_ids=["document-1"], + ) + assert requests[0]["document_ids"] == ["document-1"] + + with pytest.raises(ValueError, match="document_ids must be omitted or non-empty"): + await client.retrieve( + "annual leave", + dataset_ids=["dataset-1"], + document_ids=[], + ) + assert len(requests) == 1 + + @pytest.mark.anyio async def test_nonzero_api_code_is_normalized_and_redacts_api_key() -> None: async def handler(request: httpx.Request) -> httpx.Response: diff --git a/backend/tests/test_ragflow_tools.py b/backend/tests/test_ragflow_tools.py index 03e015e0e..a54777cb2 100644 --- a/backend/tests/test_ragflow_tools.py +++ b/backend/tests/test_ragflow_tools.py @@ -45,6 +45,7 @@ class FakeRAGFlowClient: retrieval_by_dataset_ids: Mapping[tuple[str, ...], dict] | None = None, retrieval_errors_by_dataset_ids: Mapping[tuple[str, ...], Exception] | None = None, error: Exception | None = None, + documents_by_dataset_id: Mapping[str, list[dict]] | None = None, ) -> None: self.datasets_by_id = dict(datasets_by_id or {}) self.dataset_errors_by_id = dict(dataset_errors_by_id or {}) @@ -53,7 +54,9 @@ class FakeRAGFlowClient: self.retrieval_by_dataset_ids = dict(retrieval_by_dataset_ids or {}) self.retrieval_errors_by_dataset_ids = dict(retrieval_errors_by_dataset_ids or {}) self.error = error + self.documents_by_dataset_id = dict(documents_by_dataset_id or {}) self.list_calls: list[str | None] = [] + self.document_list_calls: list[tuple[str, list[tuple[str, str]]]] = [] self.retrieve_calls: list[tuple[str, dict]] = [] async def list_datasets(self, *, dataset_id: str | None = None) -> list[dict]: @@ -66,6 +69,22 @@ class FakeRAGFlowClient: raise error return self.datasets_by_id.get(dataset_id, []) + async def list_documents( + self, + dataset_id: str, + *, + params: list[tuple[str, str]], + ) -> dict: + self.document_list_calls.append((dataset_id, params)) + requested_ids = [value for key, value in params if key == "ids"] + documents = self.documents_by_dataset_id.get(dataset_id, []) + if requested_ids: + documents = [document for document in documents if document.get("id") in requested_ids] + return { + "code": 0, + "data": {"docs": documents, "total": len(documents)}, + } + async def retrieve(self, query: str, **kwargs: object) -> dict: if self.error is not None: raise self.error @@ -409,6 +428,114 @@ async def test_grouped_retrieval_limits_concurrency_to_four(monkeypatch: pytest. assert fake.max_active_retrievals == 4 +@pytest.mark.anyio +async def test_selected_scope_dataset_validation_is_bounded_and_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + dataset_ids = [f"dataset-{index}" for index in range(5)] + + class ConcurrencyTrackingClient(FakeRAGFlowClient): + def __init__(self) -> None: + super().__init__(datasets_by_id={dataset_id: [_dataset(dataset_id, f"Dataset {index}")] for index, dataset_id in enumerate(dataset_ids)}) + self.active_dataset_lists = 0 + self.max_active_dataset_lists = 0 + + async def list_datasets(self, *, dataset_id: str | None = None) -> list[dict]: + self.list_calls.append(dataset_id) + self.active_dataset_lists += 1 + self.max_active_dataset_lists = max( + self.max_active_dataset_lists, + self.active_dataset_lists, + ) + try: + await asyncio.sleep(0.05) + return self.datasets_by_id.get(dataset_id or "", []) + finally: + self.active_dataset_lists -= 1 + + fake = ConcurrencyTrackingClient() + _install(monkeypatch, fake, config=_config(datasets=dataset_ids)) + + result = await ragflow_tools.knowledge_search( + "anything", + knowledge_scope={ + "version": 1, + "mode": "selected", + "dataset_ids": dataset_ids, + }, + ) + + assert result == "No relevant content found." + assert fake.max_active_dataset_lists == 4 + + +@pytest.mark.anyio +async def test_selected_scope_document_validation_is_bounded_and_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + dataset_ids = [f"dataset-{index}" for index in range(5)] + document_ids = [f"document-{index}" for index in range(5)] + + class ConcurrencyTrackingClient(FakeRAGFlowClient): + def __init__(self) -> None: + super().__init__( + datasets_by_id={dataset_id: [_dataset(dataset_id, f"Dataset {index}")] for index, dataset_id in enumerate(dataset_ids)}, + documents_by_dataset_id={ + dataset_id: [ + { + "id": document_ids[index], + "name": f"Document {index}", + "run": "DONE", + "chunk_count": 1, + } + ] + for index, dataset_id in enumerate(dataset_ids) + }, + ) + self.active_document_lists = 0 + self.max_active_document_lists = 0 + + async def list_documents( + self, + dataset_id: str, + *, + params: list[tuple[str, str]], + ) -> dict: + self.document_list_calls.append((dataset_id, params)) + self.active_document_lists += 1 + self.max_active_document_lists = max( + self.max_active_document_lists, + self.active_document_lists, + ) + try: + await asyncio.sleep(0.05) + documents = self.documents_by_dataset_id[dataset_id] + return { + "code": 0, + "data": {"docs": documents, "total": len(documents)}, + } + finally: + self.active_document_lists -= 1 + + fake = ConcurrencyTrackingClient() + _install(monkeypatch, fake, config=_config(datasets=dataset_ids)) + + result = await ragflow_tools.knowledge_search( + "anything", + knowledge_scope={ + "version": 1, + "mode": "selected", + "dataset_ids": dataset_ids, + "document_filters": [ + { + "dataset_id": dataset_id, + "document_ids": [document_ids[index]], + } + for index, dataset_id in enumerate(dataset_ids) + ], + }, + ) + + assert result == "No relevant content found." + assert fake.max_active_document_lists == 4 + + @pytest.mark.anyio async def test_dataset_without_embedding_metadata_returns_protocol_error(monkeypatch: pytest.MonkeyPatch) -> None: fake = FakeRAGFlowClient(all_datasets=[{"id": DATASET_ID_1, "name": "Broken", "chunk_count": 1}]) @@ -438,6 +565,151 @@ async def test_group_failure_remains_strict_and_redacts_secret_and_dataset_id(mo assert len(fake.retrieve_calls) == 2 +@pytest.mark.anyio +async def test_selected_scope_intersects_operator_allowlist_and_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake = FakeRAGFlowClient( + datasets_by_id={ + DATASET_ID_1: [_dataset(DATASET_ID_1, "Allowed")], + DATASET_ID_2: [_dataset(DATASET_ID_2, "Not allowed")], + }, + ) + _install(monkeypatch, fake, config=_config(datasets=[DATASET_ID_1])) + + result = await ragflow_tools.knowledge_search( + "leave", + knowledge_scope={ + "version": 1, + "mode": "selected", + "dataset_ids": [DATASET_ID_2], + }, + ) + + assert result == ("Error: The selected knowledge scope is no longer available; choose the knowledge bases again.") + assert fake.list_calls == [] + assert fake.retrieve_calls == [] + + +@pytest.mark.anyio +async def test_selected_document_scope_validates_membership_and_splits_groups( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake = FakeRAGFlowClient( + datasets_by_id={ + DATASET_ID_1: [_dataset(DATASET_ID_1, "All files")], + DATASET_ID_2: [_dataset(DATASET_ID_2, "One file")], + }, + documents_by_dataset_id={ + DATASET_ID_2: [ + { + "id": "doc-1", + "name": "guide.pdf", + "run": "DONE", + "chunk_count": 3, + } + ] + }, + ) + _install( + monkeypatch, + fake, + config=_config(datasets=[DATASET_ID_1, DATASET_ID_2]), + ) + + result = await ragflow_tools.knowledge_search( + "rice", + knowledge_scope={ + "version": 1, + "mode": "selected", + "dataset_ids": [DATASET_ID_1, DATASET_ID_2], + "document_filters": [{"dataset_id": DATASET_ID_2, "document_ids": ["doc-1"]}], + }, + ) + + assert result == "No relevant content found." + assert [call[1] for call in fake.retrieve_calls] == [ + { + "dataset_ids": [DATASET_ID_1], + "page_size": 8, + "similarity_threshold": 0.2, + "vector_similarity_weight": 0.3, + "top_k": 256, + }, + { + "dataset_ids": [DATASET_ID_2], + "document_ids": ["doc-1"], + "page_size": 8, + "similarity_threshold": 0.2, + "vector_similarity_weight": 0.3, + "top_k": 256, + }, + ] + assert fake.document_list_calls == [ + ( + DATASET_ID_2, + [ + ("page", "1"), + ("page_size", "1"), + ("ids", "doc-1"), + ], + ) + ] + + +@pytest.mark.anyio +@pytest.mark.parametrize( + "documents", + [ + [], + [{"id": "other", "run": "DONE", "chunk_count": 1}], + [{"id": "doc-1", "run": "RUNNING", "chunk_count": 0}], + [{"id": "doc-1", "run": "DONE", "chunk_count": 0}], + ], +) +async def test_invalid_or_unsearchable_document_selection_fails_closed( + monkeypatch: pytest.MonkeyPatch, + documents: list[dict], +) -> None: + fake = FakeRAGFlowClient( + datasets_by_id={ + DATASET_ID_1: [_dataset(DATASET_ID_1, "Selected")], + }, + documents_by_dataset_id={DATASET_ID_1: documents}, + ) + _install(monkeypatch, fake) + + result = await ragflow_tools.knowledge_search( + "rice", + knowledge_scope={ + "version": 1, + "mode": "selected", + "dataset_ids": [DATASET_ID_1], + "document_filters": [{"dataset_id": DATASET_ID_1, "document_ids": ["doc-1"]}], + }, + ) + + assert result == ("Error: The selected knowledge scope is no longer available; choose the knowledge bases or files again.") + assert fake.retrieve_calls == [] + + +@pytest.mark.anyio +async def test_disabled_scope_rejects_direct_tool_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake = FakeRAGFlowClient() + _install(monkeypatch, fake) + + result = await ragflow_tools.knowledge_search( + "rice", + knowledge_scope={"version": 1, "mode": "disabled"}, + ) + + assert result == "Error: Knowledge search is disabled for this turn." + assert fake.list_calls == [] + assert fake.retrieve_calls == [] + + @pytest.mark.anyio async def test_missing_dataset_binding_with_empty_catalog_returns_guidance(monkeypatch: pytest.MonkeyPatch) -> None: fake = FakeRAGFlowClient() @@ -691,6 +963,27 @@ def test_retrieval_settings_allow_omitting_dataset_ids(monkeypatch: pytest.Monke assert config.datasets is None +def test_retrieval_settings_do_not_fall_back_to_knowledge_capability_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + legacy_global = SimpleNamespace( + enabled=True, + base_url="http://legacy-ragflow.test", + api_key="legacy-secret", + timeout=60, + ) + app_config = SimpleNamespace( + knowledge_base=legacy_global, + get_tool_config=lambda _name: None, + ) + monkeypatch.setattr(ragflow_tools, "get_app_config", lambda: app_config) + + settings, error = ragflow_tools._settings_or_error() + + assert settings is None + assert error == "Error: knowledge_search is not configured; add its RAGFlow settings to the tools list in config.yaml." + + @pytest.mark.anyio async def test_explicitly_empty_dataset_allowlist_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: fake = FakeRAGFlowClient(all_datasets=[_dataset(DATASET_ID_1, "Must remain inaccessible")]) @@ -703,12 +996,15 @@ async def test_explicitly_empty_dataset_allowlist_fails_closed(monkeypatch: pyte assert fake.retrieve_calls == [] -def test_agent_exposes_only_query_on_single_search_tool() -> None: - assert not hasattr(ragflow_tools, "list_knowledge_bases_tool") - assert not hasattr(ragflow_tools, "list_knowledge_bases") +def test_agent_exposes_search_and_name_listing_tools() -> None: + assert hasattr(ragflow_tools, "list_knowledge_bases_tool") + assert hasattr(ragflow_tools, "list_knowledge_bases") assert ragflow_tools.knowledge_search_tool.name == "knowledge_search" assert ragflow_tools.knowledge_search_tool.coroutine is not None assert set(ragflow_tools.knowledge_search_tool.tool_call_schema.model_fields) == {"query"} + assert ragflow_tools.list_knowledge_bases_tool.name == "list_knowledge_bases" + assert ragflow_tools.list_knowledge_bases_tool.coroutine is not None + assert not ragflow_tools.list_knowledge_bases_tool.tool_call_schema.model_fields def test_tool_assembly_hides_bound_dataset_ids_without_network_io(monkeypatch: pytest.MonkeyPatch) -> None: @@ -723,6 +1019,7 @@ def test_tool_assembly_hides_bound_dataset_ids_without_network_io(monkeypatch: p ) config = SimpleNamespace( tools=[tool_config], + knowledge_base=SimpleNamespace(enabled=True), sandbox=SimpleNamespace(use="example.remote:Sandbox"), skill_evolution=SimpleNamespace(enabled=False), models=[], @@ -741,6 +1038,29 @@ def test_tool_assembly_hides_bound_dataset_ids_without_network_io(monkeypatch: p assert {tool.name for tool in tools}.isdisjoint({"list_knowledge_bases"}) +def test_tool_assembly_hides_configured_knowledge_provider_when_capability_is_disabled() -> None: + tool_config = ToolConfig( + name="knowledge_search", + group="knowledge", + use="deerflow.community.ragflow.tools:knowledge_search_tool", + base_url="http://ragflow.test", + api_key="ragflow-secret", + ) + config = SimpleNamespace( + tools=[tool_config], + knowledge_base=SimpleNamespace(enabled=False), + sandbox=SimpleNamespace(use="example.remote:Sandbox"), + skill_evolution=SimpleNamespace(enabled=False), + models=[], + acp_agents={}, + get_model_config=lambda _name: None, + ) + + tools = get_available_tools(include_mcp=False, app_config=config) + + assert {tool.name for tool in tools}.isdisjoint({"knowledge_search"}) + + def test_ragflow_package_has_explicit_init_file() -> None: package_dir = Path(ragflow_tools.__file__).resolve().parent diff --git a/backend/tests/test_subagent_batch_service.py b/backend/tests/test_subagent_batch_service.py index 2ac76bfd2..0568e194a 100644 --- a/backend/tests/test_subagent_batch_service.py +++ b/backend/tests/test_subagent_batch_service.py @@ -43,6 +43,11 @@ def _request(**overrides) -> BatchSubmitRequest: "system_prompt": "Work carefully.", }, "parent_model": "model-a", + "knowledge_scope": { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + }, }, } values.update(overrides) @@ -142,6 +147,11 @@ async def test_execute_item_marks_real_running_then_persists_terminal_result(mon assert repository.finalized["succeeded"] is True assert repository.finalized["result"] == "done" assert executor_kwargs["execution_capacity"] is execution_capacity + assert executor_kwargs["knowledge_scope"] == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + } @pytest.mark.asyncio diff --git a/backend/tests/test_task_tool_core_logic.py b/backend/tests/test_task_tool_core_logic.py index 16101b9f8..464747c57 100644 --- a/backend/tests/test_task_tool_core_logic.py +++ b/backend/tests/test_task_tool_core_logic.py @@ -3276,3 +3276,21 @@ def test_task_tool_forwards_no_criteria_by_default(monkeypatch): assert executor_kwargs["acceptance_criteria"] is None assert "" not in delegated_prompt + + +def test_task_tool_forwards_execution_only_knowledge_scope(monkeypatch): + runtime = _make_runtime() + runtime.context["__knowledge_scope_execution"] = { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + "display": {"datasets": [{"id": "dataset-1", "name": "Private label"}]}, + } + + executor_kwargs, _ = _capture_executor_call(monkeypatch, runtime=runtime) + + assert executor_kwargs["knowledge_scope"] == { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + } diff --git a/backend/tests/test_tool_error_handling_middleware.py b/backend/tests/test_tool_error_handling_middleware.py index 1c6fe0c74..e026e05f7 100644 --- a/backend/tests/test_tool_error_handling_middleware.py +++ b/backend/tests/test_tool_error_handling_middleware.py @@ -67,6 +67,21 @@ def _stub_runtime_middleware_imports(monkeypatch: pytest.MonkeyPatch) -> None: self.args = args self.kwargs = kwargs + class FakeInputSanitizationMiddleware(FakeMiddleware): + pass + + class FakeThreadDataMiddleware(FakeMiddleware): + pass + + class FakeSandboxMiddleware(FakeMiddleware): + pass + + class FakeDanglingToolCallMiddleware(FakeMiddleware): + pass + + class FakeSandboxAuditMiddleware(FakeMiddleware): + pass + class FakeLLMErrorHandlingMiddleware: def __init__(self, *, app_config): self.app_config = app_config @@ -82,22 +97,34 @@ def _stub_runtime_middleware_imports(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.thread_data_middleware", - _module("deerflow.agents.middlewares.thread_data_middleware", ThreadDataMiddleware=FakeMiddleware), + _module( + "deerflow.agents.middlewares.thread_data_middleware", + ThreadDataMiddleware=FakeThreadDataMiddleware, + ), ) monkeypatch.setitem( sys.modules, "deerflow.sandbox.middleware", - _module("deerflow.sandbox.middleware", SandboxMiddleware=FakeMiddleware), + _module( + "deerflow.sandbox.middleware", + SandboxMiddleware=FakeSandboxMiddleware, + ), ) monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.dangling_tool_call_middleware", - _module("deerflow.agents.middlewares.dangling_tool_call_middleware", DanglingToolCallMiddleware=FakeMiddleware), + _module( + "deerflow.agents.middlewares.dangling_tool_call_middleware", + DanglingToolCallMiddleware=FakeDanglingToolCallMiddleware, + ), ) monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.sandbox_audit_middleware", - _module("deerflow.agents.middlewares.sandbox_audit_middleware", SandboxAuditMiddleware=FakeMiddleware), + _module( + "deerflow.agents.middlewares.sandbox_audit_middleware", + SandboxAuditMiddleware=FakeSandboxAuditMiddleware, + ), ) @@ -109,6 +136,21 @@ def test_build_subagent_runtime_middlewares_threads_app_config_to_llm_middleware self.args = args self.kwargs = kwargs + class FakeInputSanitizationMiddleware(FakeMiddleware): + pass + + class FakeThreadDataMiddleware(FakeMiddleware): + pass + + class FakeSandboxMiddleware(FakeMiddleware): + pass + + class FakeDanglingToolCallMiddleware(FakeMiddleware): + pass + + class FakeSandboxAuditMiddleware(FakeMiddleware): + pass + class FakeLLMErrorHandlingMiddleware: def __init__(self, *, app_config): captured["app_config"] = app_config @@ -126,29 +168,38 @@ def test_build_subagent_runtime_middlewares_threads_app_config_to_llm_middleware monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.thread_data_middleware", - _module("deerflow.agents.middlewares.thread_data_middleware", ThreadDataMiddleware=FakeMiddleware), + _module( + "deerflow.agents.middlewares.thread_data_middleware", + ThreadDataMiddleware=FakeThreadDataMiddleware, + ), ) monkeypatch.setitem( sys.modules, "deerflow.sandbox.middleware", - _module("deerflow.sandbox.middleware", SandboxMiddleware=FakeMiddleware), + _module("deerflow.sandbox.middleware", SandboxMiddleware=FakeSandboxMiddleware), ) monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.dangling_tool_call_middleware", - _module("deerflow.agents.middlewares.dangling_tool_call_middleware", DanglingToolCallMiddleware=FakeMiddleware), + _module( + "deerflow.agents.middlewares.dangling_tool_call_middleware", + DanglingToolCallMiddleware=FakeDanglingToolCallMiddleware, + ), ) monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.sandbox_audit_middleware", - _module("deerflow.agents.middlewares.sandbox_audit_middleware", SandboxAuditMiddleware=FakeMiddleware), + _module( + "deerflow.agents.middlewares.sandbox_audit_middleware", + SandboxAuditMiddleware=FakeSandboxAuditMiddleware, + ), ) monkeypatch.setitem( sys.modules, "deerflow.agents.middlewares.input_sanitization_middleware", _module( "deerflow.agents.middlewares.input_sanitization_middleware", - InputSanitizationMiddleware=FakeMiddleware, + InputSanitizationMiddleware=FakeInputSanitizationMiddleware, neutralize_untrusted_tags=lambda value: value, ), ) @@ -156,7 +207,7 @@ def test_build_subagent_runtime_middlewares_threads_app_config_to_llm_middleware middlewares = build_subagent_runtime_middlewares(app_config=app_config, lazy_init=False) assert captured["app_config"] is app_config - # 9 baseline (InputSanitization, ToolOutputBudget, ToolResultSanitization, + # 10 baseline (InputSanitization, KnowledgeScope, ToolOutputBudget, ToolResultSanitization, # ThreadData, Sandbox, DanglingToolCall, LLMErrorHandling, SandboxAudit, # ToolErrorHandling) # + 1 ReadBeforeWriteMiddleware + 1 LoopDetectionMiddleware @@ -168,6 +219,7 @@ def test_build_subagent_runtime_middlewares_threads_app_config_to_llm_middleware # (all enabled by default). from deerflow.agents.middlewares.durable_context_middleware import DurableContextMiddleware from deerflow.agents.middlewares.dynamic_context_middleware import SubagentDateContextMiddleware + from deerflow.agents.middlewares.knowledge_scope_middleware import KnowledgeScopeMiddleware from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware from deerflow.agents.middlewares.skill_activation_middleware import SkillActivationMiddleware from deerflow.agents.middlewares.skill_tool_policy_middleware import SkillToolPolicyMiddleware @@ -176,9 +228,10 @@ def test_build_subagent_runtime_middlewares_threads_app_config_to_llm_middleware from deerflow.agents.middlewares.tool_output_budget_middleware import ToolOutputBudgetMiddleware from deerflow.agents.middlewares.tool_receipt_middleware import ToolReceiptMiddleware - assert len(middlewares) == 19 + assert len(middlewares) == 20 assert isinstance(middlewares[0], FakeMiddleware) # InputSanitizationMiddleware stub - assert isinstance(middlewares[1], ToolOutputBudgetMiddleware) + assert isinstance(middlewares[1], KnowledgeScopeMiddleware) + assert isinstance(middlewares[2], ToolOutputBudgetMiddleware) assert any(isinstance(m, ToolErrorHandlingMiddleware) for m in middlewares) # The receipt layer wraps ToolErrorHandlingMiddleware so receipts read the # deerflow_tool_meta status it stamps (guard-enforced, like ToolProgress). diff --git a/backend/tests/test_tool_output_budget_middleware.py b/backend/tests/test_tool_output_budget_middleware.py index b68291188..e5617fc80 100644 --- a/backend/tests/test_tool_output_budget_middleware.py +++ b/backend/tests/test_tool_output_budget_middleware.py @@ -1210,11 +1210,14 @@ class TestMiddlewareChainIntegration: middlewares = build_subagent_runtime_middlewares(app_config=app_config, lazy_init=False) # InputSanitizationMiddleware is the outermost wrap_model_call wrapper; - # ToolOutputBudgetMiddleware is the first wrap_tool_call handler. + # KnowledgeScopeMiddleware cleans model input immediately inside it; + # ToolOutputBudgetMiddleware remains immediately inside the scope guard. from deerflow.agents.middlewares.input_sanitization_middleware import InputSanitizationMiddleware + from deerflow.agents.middlewares.knowledge_scope_middleware import KnowledgeScopeMiddleware assert isinstance(middlewares[0], InputSanitizationMiddleware) - assert isinstance(middlewares[1], ToolOutputBudgetMiddleware) + assert isinstance(middlewares[1], KnowledgeScopeMiddleware) + assert isinstance(middlewares[2], ToolOutputBudgetMiddleware) def test_budget_middleware_in_lead_chain(self): from deerflow.agents.middlewares.tool_error_handling_middleware import build_lead_runtime_middlewares @@ -1223,9 +1226,11 @@ class TestMiddlewareChainIntegration: middlewares = build_lead_runtime_middlewares(app_config=app_config, lazy_init=False) from deerflow.agents.middlewares.input_sanitization_middleware import InputSanitizationMiddleware + from deerflow.agents.middlewares.knowledge_scope_middleware import KnowledgeScopeMiddleware assert isinstance(middlewares[0], InputSanitizationMiddleware) - assert isinstance(middlewares[1], ToolOutputBudgetMiddleware) + assert isinstance(middlewares[1], KnowledgeScopeMiddleware) + assert isinstance(middlewares[2], ToolOutputBudgetMiddleware) # =========================================================================== diff --git a/config.example.yaml b/config.example.yaml index 05e1ec087..ac449c7c0 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -20,7 +20,7 @@ # ============================================================================ # Bump this number when the config schema changes. # Run `make config-upgrade` to merge new fields into your local config.yaml. -config_version: 45 +config_version: 46 # Optional per-model request pacing (inside a models[] entry): # request_admission: @@ -727,6 +727,18 @@ models: # write_timeout: 60.0 # pool_timeout: 30.0 +# ============================================================================ +# Knowledge Capability +# ============================================================================ +# Optional knowledge capability. The feature is disabled by default. Provider +# connection and retrieval settings belong to the matching entry in `tools:` +# below (for example, the RAGFlow `knowledge_search` tool), not here. +knowledge_base: + enabled: false + # Custom-agent chat only. Keeps the selector hidden unless the effective + # knowledge_search provider is the bundled RAGFlow tool. + scope_selection_enabled: false + # ============================================================================ # Tool Groups Configuration # ============================================================================ @@ -775,6 +787,9 @@ tools: # top_k: 256 # max_chars_per_chunk: 800 # max_total_chars: 8000 + # - name: list_knowledge_bases + # group: knowledge + # use: deerflow.community.ragflow.tools:list_knowledge_bases_tool # LightRAG knowledge retrieval (read-only). Alternative provider for the same # knowledge_search tool; uncomment this entry instead of the RAGFlow one — diff --git a/deploy/helm/deer-flow/README.md b/deploy/helm/deer-flow/README.md index 692463b09..8eabdfa8d 100644 --- a/deploy/helm/deer-flow/README.md +++ b/deploy/helm/deer-flow/README.md @@ -135,7 +135,7 @@ they resolve from the `secrets` map): ```yaml config: | - config_version: 45 + config_version: 46 models: - name: gpt-4 use: langchain_openai:ChatOpenAI @@ -155,6 +155,10 @@ config: | connection_string: $DATABASE_URL stream_bridge: type: redis # cross-pod SSE; URL from DEER_FLOW_STREAM_BRIDGE_REDIS_URL + knowledge_base: + enabled: true + scope_selection_enabled: false + # Provider connection/retrieval settings belong on the knowledge_search tool. # Tools MUST be listed explicitly - the agent gets none otherwise # (BUILTIN_TOOLS only adds present_file + ask_clarification). The chart # default in values.yaml enables the sandbox tools + web tools (web_search, @@ -166,6 +170,7 @@ config: | - name: file:read - name: file:write - name: bash + - name: knowledge tools: - name: web_search group: web @@ -179,6 +184,14 @@ config: | group: web use: deerflow.community.image_search.tools:image_search_tool max_results: 5 + - name: knowledge_search + group: knowledge + use: deerflow.community.ragflow.tools:knowledge_search_tool + base_url: http://ragflow:9380 + api_key: $RAGFLOW_API_KEY + - name: list_knowledge_bases + group: knowledge + use: deerflow.community.ragflow.tools:list_knowledge_bases_tool - name: bash group: bash use: deerflow.sandbox.tools:bash_tool diff --git a/deploy/helm/deer-flow/values.yaml b/deploy/helm/deer-flow/values.yaml index e08c7934d..50472404a 100644 --- a/deploy/helm/deer-flow/values.yaml +++ b/deploy/helm/deer-flow/values.yaml @@ -249,11 +249,17 @@ ingress: # -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never # inline literal secret values here. The default enables provisioner sandbox. config: | - config_version: 45 + config_version: 46 log_level: info recursion_limit: 100 max_recursion_limit: 1000 + # Optional provider-agnostic knowledge capability and custom-agent scope selector. + # RAGFlow connection/retrieval settings belong on the knowledge_search tool. + knowledge_base: + enabled: false + scope_selection_enabled: false + models: [] # Example (uncomment & set the matching secret in `secrets`): # - name: gpt-4 @@ -339,6 +345,9 @@ config: | # top_k: 256 # max_chars_per_chunk: 800 # max_total_chars: 8000 + # - name: list_knowledge_bases + # group: knowledge + # use: deerflow.community.ragflow.tools:list_knowledge_bases_tool # Optional tenant-shared, read-only LightRAG retrieval; alternative # provider for the same knowledge_search tool (duplicate names keep the # first entry, so configure exactly one). Requires LightRAG v1.4.9+. Put diff --git a/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx b/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx index 50817656b..f55dea355 100644 --- a/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx +++ b/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx @@ -2,7 +2,7 @@ import { BotIcon, PlusSquare } from "lucide-react"; import { useParams, useRouter } from "next/navigation"; -import { useCallback, useEffect, useMemo, useState } from "react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import type { PromptInputMessage } from "@/components/ai-elements/prompt-input"; import { Button } from "@/components/ui/button"; @@ -18,6 +18,7 @@ import { InputBox, type InputBoxSubmitOptions, } from "@/components/workspace/input-box"; +import { KnowledgeScopeSelector } from "@/components/workspace/knowledge-scope-selector"; import { MessageList, MESSAGE_LIST_DEFAULT_PADDING_BOTTOM, @@ -38,8 +39,17 @@ import { useActiveGoal } from "@/components/workspace/use-active-goal"; import { useAgent } from "@/core/agents"; import { useAuth } from "@/core/auth/AuthProvider"; import { hasPermission, PERMISSIONS } from "@/core/auth/permissions"; -import { useBrowserControlEnabled } from "@/core/features"; +import { + useBrowserControlEnabled, + useKnowledgeBaseEnabled, +} from "@/core/features"; import { useI18n } from "@/core/i18n/hooks"; +import { + ALL_KNOWLEDGE_SCOPE, + buildKnowledgeScopeSnapshot, + KNOWLEDGE_SCOPE_KEY, + type KnowledgeScopeSelection, +} from "@/core/knowledge"; import { buildHumanInputResponseText, hasOpenHumanInputRequest, @@ -85,6 +95,7 @@ export default function AgentChatPage() { const [settings, setSettings] = useThreadSettings(threadId); const [localSettings, setLocalSettings] = useLocalSettings(); const { enabled: browserControlEnabled } = useBrowserControlEnabled(); + const { scopeSelectionEnabled } = useKnowledgeBaseEnabled(); const { tokenUsageEnabled } = useModels(); const threadTokenUsage = useThreadTokenUsage( isNewThread || isMock ? undefined : threadId, @@ -98,6 +109,51 @@ export default function AgentChatPage() { const contextUsage = selectContextUsage(threadTokenUsage.data); const { showNotification } = useNotification(); + const selectorVisible = + scopeSelectionEnabled && env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true"; + const agentKnowledgeEnabled = + agent !== null && + (agent.tool_groups == null || agent.tool_groups.includes("knowledge")); + const [knowledgeScope, setKnowledgeScope] = + useState(null); + const previousConversationRef = useRef({ + agentName: agent_name, + threadId, + isNewThread, + }); + + useEffect(() => { + setKnowledgeScope((current) => { + if (!selectorVisible) return null; + return current ?? ALL_KNOWLEDGE_SCOPE; + }); + }, [selectorVisible]); + + useEffect(() => { + const previous = previousConversationRef.current; + if (previous.agentName !== agent_name || previous.threadId !== threadId) { + const isNewThreadRouteReplacement = + previous.agentName === agent_name && + previous.isNewThread && + !isNewThread; + if (!isNewThreadRouteReplacement) { + setKnowledgeScope(selectorVisible ? ALL_KNOWLEDGE_SCOPE : null); + } + } + previousConversationRef.current = { + agentName: agent_name, + threadId, + isNewThread, + }; + }, [agent_name, isNewThread, selectorVisible, threadId]); + + const currentKnowledgeScopeSnapshot = useMemo( + () => + selectorVisible && agentKnowledgeEnabled && knowledgeScope + ? buildKnowledgeScopeSnapshot(knowledgeScope) + : null, + [agentKnowledgeEnabled, knowledgeScope, selectorVisible], + ); useEffect(() => { setIsWelcomeMode(isNewThread); @@ -116,6 +172,7 @@ export default function AgentChatPage() { } = useThreadStream({ threadId: isNewThread ? undefined : threadId, displayThreadId: threadId, + assistantId: agent_name, context: { ...settings.context, agent_name: agent_name }, isMock, onSend: () => { @@ -179,18 +236,27 @@ export default function AgentChatPage() { const handleSubmit = useCallback( (message: PromptInputMessage, options?: InputBoxSubmitOptions) => { + const scopedOptions = currentKnowledgeScopeSnapshot + ? { + ...options, + additionalKwargs: { + ...options?.additionalKwargs, + [KNOWLEDGE_SCOPE_KEY]: currentKnowledgeScopeSnapshot, + }, + } + : options; const sendPromise = sendMessage( threadId, message, { agent_name }, - options, + scopedOptions, ); if (message.files.length > 0) { return sendPromise; } void sendPromise; }, - [sendMessage, threadId, agent_name], + [currentKnowledgeScopeSnapshot, sendMessage, threadId, agent_name], ); const handleSubmitHumanInput = useCallback( @@ -207,6 +273,9 @@ export default function AgentChatPage() { additionalKwargs: { hide_from_ui: true, human_input_response: response, + ...(currentKnowledgeScopeSnapshot + ? { [KNOWLEDGE_SCOPE_KEY]: currentKnowledgeScopeSnapshot } + : {}), }, onSent: () => { sent = true; @@ -215,7 +284,7 @@ export default function AgentChatPage() { ); return sent; }, - [agent_name, sendMessage, threadId], + [agent_name, currentKnowledgeScopeSnapshot, sendMessage, threadId], ); const handleStop = useCallback(async () => { @@ -228,8 +297,15 @@ export default function AgentChatPage() { ); const handleEditAndRegenerate = useCallback( (messageId: string, replacementText: string) => - editAndRegenerateMessage(threadId, messageId, replacementText), - [editAndRegenerateMessage, threadId], + editAndRegenerateMessage( + threadId, + messageId, + replacementText, + currentKnowledgeScopeSnapshot + ? { [KNOWLEDGE_SCOPE_KEY]: currentKnowledgeScopeSnapshot } + : undefined, + ), + [currentKnowledgeScopeSnapshot, editAndRegenerateMessage, threadId], ); const tokenUsageInlineMode = tokenUsageEnabled @@ -439,6 +515,21 @@ export default function AgentChatPage() { agentSkillNames={agent?.skills} agentSkillsLoading={agentSkillsLoading} defaultModelName={agent?.model} + knowledgeScopeControl={ + selectorVisible && knowledgeScope ? ( + + ) : undefined + } autoFocus={isWelcomeMode} status={ thread.error @@ -458,6 +549,7 @@ export default function AgentChatPage() { disabled={ env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true" || isUploading || + (selectorVisible && agent === null) || (!isNewThread && isHistoryLoading) } onContextChange={(context, options) => { diff --git a/frontend/src/components/workspace/chats/chat-page.tsx b/frontend/src/components/workspace/chats/chat-page.tsx index da11636a8..9610f2ca8 100644 --- a/frontend/src/components/workspace/chats/chat-page.tsx +++ b/frontend/src/components/workspace/chats/chat-page.tsx @@ -18,6 +18,7 @@ import { InputBox, type InputBoxSubmitOptions, } from "@/components/workspace/input-box"; +import { KnowledgeScopeSelector } from "@/components/workspace/knowledge-scope-selector"; import { MessageList, MESSAGE_LIST_DEFAULT_PADDING_BOTTOM, @@ -38,8 +39,17 @@ import { useActiveGoal } from "@/components/workspace/use-active-goal"; import { Welcome } from "@/components/workspace/welcome"; import { useAuth } from "@/core/auth/AuthProvider"; import { hasPermission, PERMISSIONS } from "@/core/auth/permissions"; -import { useBrowserControlEnabled } from "@/core/features"; +import { + useBrowserControlEnabled, + useKnowledgeBaseEnabled, +} from "@/core/features"; import { useI18n } from "@/core/i18n/hooks"; +import { + ALL_KNOWLEDGE_SCOPE, + buildKnowledgeScopeSnapshot, + KNOWLEDGE_SCOPE_KEY, + type KnowledgeScopeSelection, +} from "@/core/knowledge"; import { buildHumanInputResponseText, hasOpenHumanInputRequest, @@ -127,6 +137,41 @@ export default function ChatPage() { }, [isNewThread]); const { showNotification } = useNotification(); + const { scopeSelectionEnabled } = useKnowledgeBaseEnabled(); + const selectorVisible = + scopeSelectionEnabled && env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true"; + const [knowledgeScope, setKnowledgeScope] = + useState(null); + const previousConversationRef = useRef({ threadId, isNewThread }); + + useEffect(() => { + setKnowledgeScope((current) => { + if (!selectorVisible) return null; + return current ?? ALL_KNOWLEDGE_SCOPE; + }); + }, [selectorVisible]); + + useEffect(() => { + const previous = previousConversationRef.current; + if ( + previous.threadId !== threadId || + previous.isNewThread !== isNewThread + ) { + const isNewThreadRouteReplacement = previous.isNewThread && !isNewThread; + if (!isNewThreadRouteReplacement) { + setKnowledgeScope(selectorVisible ? ALL_KNOWLEDGE_SCOPE : null); + } + } + previousConversationRef.current = { threadId, isNewThread }; + }, [isNewThread, selectorVisible, threadId]); + + const currentKnowledgeScopeSnapshot = useMemo( + () => + selectorVisible && knowledgeScope + ? buildKnowledgeScopeSnapshot(knowledgeScope) + : null, + [knowledgeScope, selectorVisible], + ); const { thread, @@ -257,13 +302,27 @@ export default function ChatPage() { if (submissionEpochRef.current !== submissionEpoch) { throw new Error("thread-submission-stale"); } - const sendPromise = sendMessage(threadId, message, undefined, options); + const scopedOptions = currentKnowledgeScopeSnapshot + ? { + ...options, + additionalKwargs: { + ...options?.additionalKwargs, + [KNOWLEDGE_SCOPE_KEY]: currentKnowledgeScopeSnapshot, + }, + } + : options; + const sendPromise = sendMessage( + threadId, + message, + undefined, + scopedOptions, + ); if (message.files.length > 0) { return sendPromise; } void sendPromise; }, - [sendMessage, threadId, ensureProjectThread], + [currentKnowledgeScopeSnapshot, sendMessage, threadId, ensureProjectThread], ); const handleSubmitHumanInput = useCallback( async (request: HumanInputRequest, response: HumanInputResponse) => { @@ -279,6 +338,9 @@ export default function ChatPage() { additionalKwargs: { hide_from_ui: true, human_input_response: response, + ...(currentKnowledgeScopeSnapshot + ? { [KNOWLEDGE_SCOPE_KEY]: currentKnowledgeScopeSnapshot } + : {}), }, onSent: () => { sent = true; @@ -287,7 +349,7 @@ export default function ChatPage() { ); return sent; }, - [sendMessage, threadId], + [currentKnowledgeScopeSnapshot, sendMessage, threadId], ); const handleStop = useCallback(async () => { await thread.stop(); @@ -299,8 +361,15 @@ export default function ChatPage() { ); const handleEditAndRegenerate = useCallback( (messageId: string, replacementText: string) => - editAndRegenerateMessage(threadId, messageId, replacementText), - [editAndRegenerateMessage, threadId], + editAndRegenerateMessage( + threadId, + messageId, + replacementText, + currentKnowledgeScopeSnapshot + ? { [KNOWLEDGE_SCOPE_KEY]: currentKnowledgeScopeSnapshot } + : undefined, + ), + [currentKnowledgeScopeSnapshot, editAndRegenerateMessage, threadId], ); const handleBranchTurn = useCallback( async (messageId: string, messageIds: string[]) => { @@ -527,6 +596,16 @@ export default function ChatPage() { isWelcomeMode={isWelcomeMode} threadId={threadId} draftThreadId={isNewThread ? "new" : threadId} + knowledgeScopeControl={ + selectorVisible && knowledgeScope ? ( + + ) : undefined + } autoFocus={isWelcomeMode} status={ thread.error diff --git a/frontend/src/components/workspace/input-box.tsx b/frontend/src/components/workspace/input-box.tsx index dbbd6174b..dd4330791 100644 --- a/frontend/src/components/workspace/input-box.tsx +++ b/frontend/src/components/workspace/input-box.tsx @@ -302,6 +302,7 @@ export function InputBox({ draftThreadId = threadId, draftAgentName, defaultModelName, + knowledgeScopeControl, initialValue, onContextChange, onFollowupsVisibilityChange, @@ -343,6 +344,8 @@ export function InputBox({ * (issue #4336). ``null`` / undefined = no agent default → use models[0]. */ defaultModelName?: string | null; + /** Optional knowledge-scope control rendered directly after mode. */ + knowledgeScopeControl?: React.ReactNode; initialValue?: string; onContextChange?: ( // Explicit selections contain only the fields changed by that action, @@ -2712,6 +2715,7 @@ export function InputBox({ + {knowledgeScopeControl} {supportReasoningEffort && context.mode !== "flash" && ( item.id === dataset.id, + ); + const datasets = [...selection.datasets]; + if (current >= 0) datasets[current] = dataset; + else datasets.push(dataset); + return { mode: "selected", datasets }; +} + +function DocumentSelector({ + agentName, + dataset, + disabled, + onChange, +}: { + agentName: string; + dataset: KnowledgeScopeDatasetSelection; + disabled: boolean; + onChange: (dataset: KnowledgeScopeDatasetSelection) => void; +}) { + const { t } = useI18n(); + const [page, setPage] = useState(1); + const [search, setSearch] = useState(""); + const query = useQuery({ + queryKey: [ + "knowledge", + "retrieval-catalog", + agentName, + dataset.id, + page, + search, + ], + queryFn: () => + listRetrievalCatalogDocuments({ + agentName, + datasetId: dataset.id, + page, + pageSize: PAGE_SIZE, + search, + }), + enabled: dataset.documents.mode === "selected", + retry: false, + }); + const selectedIds = new Set( + dataset.documents.mode === "selected" + ? dataset.documents.items.map((item) => item.id) + : [], + ); + + const toggleDocument = (document: RetrievalCatalogItem, checked: boolean) => { + if (!document.selectable) return; + const current = + dataset.documents.mode === "selected" ? dataset.documents.items : []; + onChange({ + ...dataset, + documents: { + mode: "selected", + items: checked + ? [...current.filter((item) => item.id !== document.id), document] + : current.filter((item) => item.id !== document.id), + }, + }); + }; + + return ( +
+
+ + +
+ {dataset.documents.mode === "selected" && ( + <> +
+ + { + setSearch(event.currentTarget.value); + setPage(1); + }} + /> +
+ {query.isPending ? ( + + ) : query.isError ? ( +

+ {t.knowledge.scope.loadFailed} +

+ ) : ( +
+ {query.data?.items.map((document) => ( + + ))} + +
+ )} + + )} +
+ ); +} + +function CatalogPagination({ + page, + pageSize, + total, + onPageChange, +}: { + page: number; + pageSize: number; + total: number; + onPageChange: (page: number) => void; +}) { + const { t } = useI18n(); + if (total <= pageSize) return null; + return ( +
+ + {page} + +
+ ); +} + +export function KnowledgeScopeSelector({ + agentName, + selection, + disabled = false, + unavailableReason, + onChange, +}: { + agentName: string; + selection: KnowledgeScopeSelection; + disabled?: boolean; + unavailableReason?: string; + onChange: (selection: KnowledgeScopeSelection) => void; +}) { + const { t } = useI18n(); + const [open, setOpen] = useState(false); + const [draft, setDraft] = useState(() => + cloneKnowledgeScopeSelection(selection), + ); + const [expandedDatasetId, setExpandedDatasetId] = useState( + null, + ); + const [page, setPage] = useState(1); + const [search, setSearch] = useState(""); + const query = useQuery({ + queryKey: ["knowledge", "retrieval-catalog", agentName, page, search], + queryFn: () => + listRetrievalCatalogDatasets({ + agentName, + page, + pageSize: PAGE_SIZE, + search, + }), + enabled: open && draft.mode === "selected" && !unavailableReason, + retry: false, + }); + + useEffect(() => { + if (open) setDraft(cloneKnowledgeScopeSelection(selection)); + }, [open, selection]); + + const counts = countKnowledgeScopeSelection(selection); + const label = + selection.mode === "all" + ? t.knowledge.scope.buttonAll + : selection.mode === "disabled" + ? t.knowledge.scope.buttonDisabled + : counts.documents > 0 + ? t.knowledge.scope.buttonDatasetsAndDocuments( + counts.datasets, + counts.documents, + ) + : t.knowledge.scope.buttonDatasets(counts.datasets); + const active = selection.mode !== "disabled"; + const draftInvalid = + draft.mode === "selected" && + (draft.datasets.length === 0 || + draft.datasets.some( + (dataset) => + dataset.documents.mode === "selected" && + dataset.documents.items.length === 0, + )); + const draftExceedsLimits = useMemo(() => { + if (draftInvalid) return false; + try { + buildKnowledgeScopeSnapshot(draft); + return false; + } catch { + return true; + } + }, [draft, draftInvalid]); + + const selectedDatasets = useMemo( + () => + draft.mode === "selected" + ? new Map(draft.datasets.map((item) => [item.id, item])) + : new Map(), + [draft], + ); + + const toggleDataset = (dataset: RetrievalCatalogItem, checked: boolean) => { + if (draft.mode !== "selected" || !dataset.selectable) return; + setDraft({ + mode: "selected", + datasets: checked + ? [ + ...draft.datasets.filter((item) => item.id !== dataset.id), + { ...dataset, documents: { mode: "all" } }, + ] + : draft.datasets.filter((item) => item.id !== dataset.id), + }); + if (!checked && expandedDatasetId === dataset.id) + setExpandedDatasetId(null); + }; + + const trigger = ( + + ); + + return ( + + {trigger} + + + {t.knowledge.scope.title} + {t.knowledge.scope.description} + +
+ {(["all", "selected", "disabled"] as const).map((mode) => ( + + ))} +
+ {draft.mode === "selected" && ( +
+
+ + { + setSearch(event.currentTarget.value); + setPage(1); + }} + /> +
+
+ {t.knowledge.scope.selectedCount(draft.datasets.length)} +
+ + {query.isPending ? ( + + ) : query.isError ? ( +

+ {t.knowledge.scope.loadFailed} +

+ ) : ( +
+ {query.data?.items.map((dataset) => { + const selected = selectedDatasets.get(dataset.id); + const expanded = expandedDatasetId === dataset.id; + return ( +
+
+ + toggleDataset( + dataset, + event.currentTarget.checked, + ) + } + /> + + {dataset.name} + + {!dataset.selectable && ( + + {t.knowledge.scope.notSearchable} + + )} + {selected && ( + + )} +
+ {selected && expanded && ( + + setDraft((current) => + replaceDataset(current, next), + ) + } + /> + )} +
+ ); + })} + +
+ )} + +
+ )} + {draftExceedsLimits && ( +

+ {t.knowledge.scope.selectionInvalid} +

+ )} + + + + + + + ); +} diff --git a/frontend/src/components/workspace/messages/knowledge-scope-summary.tsx b/frontend/src/components/workspace/messages/knowledge-scope-summary.tsx new file mode 100644 index 000000000..e9c928670 --- /dev/null +++ b/frontend/src/components/workspace/messages/knowledge-scope-summary.tsx @@ -0,0 +1,66 @@ +import { DatabaseIcon } from "lucide-react"; + +import { useI18n } from "@/core/i18n/hooks"; +import { + KNOWLEDGE_SCOPE_KEY, + readKnowledgeScopeSnapshot, +} from "@/core/knowledge"; + +export function KnowledgeScopeSummary({ + additionalKwargs, +}: { + additionalKwargs: Record | undefined; +}) { + const { t } = useI18n(); + const snapshot = readKnowledgeScopeSnapshot( + additionalKwargs?.[KNOWLEDGE_SCOPE_KEY], + ); + if (!snapshot) return null; + + if (snapshot.mode === "all") { + return ; + } + if (snapshot.mode === "disabled") { + return ; + } + + const datasetCount = snapshot.dataset_ids?.length ?? 0; + const documentCount = + snapshot.document_filters?.reduce( + (total, filter) => total + filter.document_ids.length, + 0, + ) ?? 0; + const datasetNames = snapshot.display?.datasets + .map((dataset) => dataset.name) + .filter(Boolean); + const documentNames = snapshot.display?.datasets + .flatMap( + (dataset) => dataset.documents?.map((document) => document.name) ?? [], + ) + .filter(Boolean); + const details = [ + datasetNames?.length ? datasetNames.join("、") : null, + documentNames?.length ? documentNames.join("、") : null, + ].filter((value): value is string => Boolean(value)); + return ( + + ); +} + +function SummaryText({ text, details }: { text: string; details?: string }) { + return ( +
+ + + {details ? `${text} · ${details}` : text} + +
+ ); +} diff --git a/frontend/src/components/workspace/messages/message-list-item.tsx b/frontend/src/components/workspace/messages/message-list-item.tsx index 418a4c333..9ec0526e0 100644 --- a/frontend/src/components/workspace/messages/message-list-item.tsx +++ b/frontend/src/components/workspace/messages/message-list-item.tsx @@ -69,6 +69,7 @@ import { ReferenceAttachmentSummary } from "../sidecar/reference-attachments"; import { SlashSkillChip } from "../slash-skill-chip"; import { Tooltip } from "../tooltip"; +import { KnowledgeScopeSummary } from "./knowledge-scope-summary"; import { MarkdownContent } from "./markdown-content"; import { createMarkdownLinkComponent } from "./markdown-link"; @@ -592,6 +593,11 @@ function MessageContent_({ ) : null} + | undefined + } + /> ); } diff --git a/frontend/src/core/features/api.ts b/frontend/src/core/features/api.ts index dc03c8e36..209bc0abc 100644 --- a/frontend/src/core/features/api.ts +++ b/frontend/src/core/features/api.ts @@ -15,6 +15,9 @@ export interface FeaturesResponse { enabled?: boolean; max_references?: number; }; + knowledge_base?: { + scope_selection_enabled?: boolean; + }; } export interface ConversationReferencesCapability { @@ -72,3 +75,12 @@ export async function fetchConversationReferencesCapability(): Promise { + const feature = (await fetchFeatures()).knowledge_base; + return { + scopeSelectionEnabled: feature?.scope_selection_enabled ?? false, + }; +} diff --git a/frontend/src/core/features/hooks.ts b/frontend/src/core/features/hooks.ts index e36ec0309..7c6a3d374 100644 --- a/frontend/src/core/features/hooks.ts +++ b/frontend/src/core/features/hooks.ts @@ -3,6 +3,7 @@ import { useQuery } from "@tanstack/react-query"; import { fetchBrowserControlEnabled, fetchConversationReferencesCapability, + fetchKnowledgeBaseFeature, fetchMcpTasksEnabled, fetchSubagentBatchesCapability, } from "./api"; @@ -67,3 +68,17 @@ export function useConversationReferencesCapability() { isLoading: isPending, }; } + +export function useKnowledgeBaseEnabled() { + const { data, isPending } = useQuery({ + queryKey: ["features", "knowledge_base"], + queryFn: fetchKnowledgeBaseFeature, + staleTime: 0, + refetchOnMount: true, + retry: false, + }); + return { + scopeSelectionEnabled: data?.scopeSelectionEnabled ?? false, + isLoading: isPending, + }; +} diff --git a/frontend/src/core/i18n/locales/en-US.ts b/frontend/src/core/i18n/locales/en-US.ts index f22cb7f91..d8cb73d89 100644 --- a/frontend/src/core/i18n/locales/en-US.ts +++ b/frontend/src/core/i18n/locales/en-US.ts @@ -403,6 +403,44 @@ export const enUS: Translations = { scheduledTasks: "Scheduled tasks", agentsDisabledTooltip: "Feature not enabled", }, + + // Knowledge scope for custom-agent chat + knowledge: { + scope: { + title: "Knowledge scope", + description: + "Choose which allowed knowledge bases and documents this agent may search.", + buttonAll: "Knowledge · All", + buttonDisabled: "Knowledge · Off", + buttonDatasets: (datasets) => + `Knowledge · ${datasets} ${datasets === 1 ? "base" : "bases"}`, + buttonDatasetsAndDocuments: (datasets, documents) => + `Knowledge · ${datasets} ${datasets === 1 ? "base" : "bases"} · ${documents} ${documents === 1 ? "file" : "files"}`, + allDatasets: "All allowed knowledge bases", + selectedDatasets: "Selected knowledge bases", + disabled: "Off", + allDocuments: "All searchable files", + selectedDocuments: "Selected files", + searchDatasets: "Search knowledge bases", + searchDocuments: "Search files", + selectedCount: (count) => `${count} selected`, + files: "Files", + notSearchable: "Not searchable", + loadFailed: + "The catalog could not be loaded. Your current selection is unchanged.", + selectionInvalid: "This selection exceeds the supported size limits.", + previous: "Previous", + next: "Next", + agentUnavailable: "This agent does not allow the knowledge tool group.", + apply: "Apply", + historyAll: "Knowledge: all allowed bases", + historyDisabled: "Knowledge: off", + historySelected: (datasets, documents) => + documents > 0 + ? `Knowledge: ${datasets} ${datasets === 1 ? "base" : "bases"}, ${documents} ${documents === 1 ? "file" : "files"}` + : `Knowledge: ${datasets} ${datasets === 1 ? "base" : "bases"}`, + }, + }, // Sidebar projects section projects: { title: "Projects", diff --git a/frontend/src/core/i18n/locales/types.ts b/frontend/src/core/i18n/locales/types.ts index 25138470b..7a36fe4c1 100644 --- a/frontend/src/core/i18n/locales/types.ts +++ b/frontend/src/core/i18n/locales/types.ts @@ -417,6 +417,40 @@ export interface Translations { emptyTrashFailed: string; }; + // Knowledge scope for custom-agent chat + knowledge: { + scope: { + title: string; + description: string; + buttonAll: string; + buttonDisabled: string; + buttonDatasets: (datasets: number) => string; + buttonDatasetsAndDocuments: ( + datasets: number, + documents: number, + ) => string; + allDatasets: string; + selectedDatasets: string; + disabled: string; + allDocuments: string; + selectedDocuments: string; + searchDatasets: string; + searchDocuments: string; + selectedCount: (count: number) => string; + files: string; + notSearchable: string; + loadFailed: string; + selectionInvalid: string; + previous: string; + next: string; + agentUnavailable: string; + apply: string; + historyAll: string; + historyDisabled: string; + historySelected: (datasets: number, documents: number) => string; + }; + }; + // Thread-scoped MCP background tasks backgroundTasks: { label: string; diff --git a/frontend/src/core/i18n/locales/zh-CN.ts b/frontend/src/core/i18n/locales/zh-CN.ts index d5bc5ba51..33e01f55f 100644 --- a/frontend/src/core/i18n/locales/zh-CN.ts +++ b/frontend/src/core/i18n/locales/zh-CN.ts @@ -377,6 +377,41 @@ export const zhCN: Translations = { scheduledTasks: "定时任务", agentsDisabledTooltip: "功能未启用", }, + + // 自定义智能体聊天中的知识库范围 + knowledge: { + scope: { + title: "知识库范围", + description: "选择该智能体本轮可以检索的知识库和文件。", + buttonAll: "知识库 · 全部", + buttonDisabled: "知识库 · 关闭", + buttonDatasets: (datasets) => `知识库 · ${datasets}库`, + buttonDatasetsAndDocuments: (datasets, documents) => + `知识库 · ${datasets}库 · ${documents}文件`, + allDatasets: "全部允许知识库", + selectedDatasets: "指定知识库", + disabled: "关闭", + allDocuments: "全部可检索文件", + selectedDocuments: "指定文件", + searchDatasets: "搜索知识库", + searchDocuments: "搜索文件", + selectedCount: (count) => `已选择 ${count} 个知识库`, + files: "文件", + notSearchable: "不可检索", + loadFailed: "目录加载失败，当前选择未改变。", + selectionInvalid: "当前选择超出支持的数量或大小限制。", + previous: "上一页", + next: "下一页", + agentUnavailable: "当前智能体未允许 knowledge 工具组。", + apply: "应用", + historyAll: "知识库：全部允许库", + historyDisabled: "知识库：关闭", + historySelected: (datasets, documents) => + documents > 0 + ? `知识库：${datasets}库 · ${documents}文件` + : `知识库：${datasets}库`, + }, + }, // Sidebar projects section projects: { title: "项目", diff --git a/frontend/src/core/knowledge/index.ts b/frontend/src/core/knowledge/index.ts new file mode 100644 index 000000000..f1021fe7b --- /dev/null +++ b/frontend/src/core/knowledge/index.ts @@ -0,0 +1,2 @@ +export * from "./scope"; +export * from "./scope-api"; diff --git a/frontend/src/core/knowledge/scope-api.ts b/frontend/src/core/knowledge/scope-api.ts new file mode 100644 index 000000000..0868a3ff0 --- /dev/null +++ b/frontend/src/core/knowledge/scope-api.ts @@ -0,0 +1,65 @@ +import { throwGatewayApiError } from "@/core/api/errors"; +import { fetch } from "@/core/api/fetcher"; +import { getBackendBaseURL } from "@/core/config"; + +export type RetrievalCatalogItem = { + id: string; + name: string; + selectable: boolean; +}; + +export type RetrievalCatalogPage = { + items: RetrievalCatalogItem[]; + page: number; + page_size: number; + total: number; +}; + +function catalogUrl(path: string): string { + return `${getBackendBaseURL()}/api/knowledge/retrieval-catalog${path}`; +} + +async function readCatalogPage( + response: Response, + fallback: string, +): Promise { + if (!response.ok) await throwGatewayApiError(response, fallback); + return (await response.json()) as RetrievalCatalogPage; +} + +export async function listRetrievalCatalogDatasets(options: { + agentName: string; + page: number; + pageSize?: number; + search?: string; +}): Promise { + const query = new URLSearchParams({ + agent_name: options.agentName, + page: String(options.page), + page_size: String(options.pageSize ?? 100), + }); + if (options.search?.trim()) query.set("search", options.search.trim()); + const response = await fetch(catalogUrl(`/datasets?${query}`)); + return readCatalogPage(response, "Failed to load the retrieval catalog."); +} + +export async function listRetrievalCatalogDocuments(options: { + agentName: string; + datasetId: string; + page: number; + pageSize?: number; + search?: string; +}): Promise { + const query = new URLSearchParams({ + agent_name: options.agentName, + page: String(options.page), + page_size: String(options.pageSize ?? 100), + }); + if (options.search?.trim()) query.set("search", options.search.trim()); + const response = await fetch( + catalogUrl( + `/datasets/${encodeURIComponent(options.datasetId)}/documents?${query}`, + ), + ); + return readCatalogPage(response, "Failed to load knowledge-base documents."); +} diff --git a/frontend/src/core/knowledge/scope.ts b/frontend/src/core/knowledge/scope.ts new file mode 100644 index 000000000..3b7ca5cd1 --- /dev/null +++ b/frontend/src/core/knowledge/scope.ts @@ -0,0 +1,316 @@ +export const KNOWLEDGE_SCOPE_KEY = "knowledge_scope"; + +const MAX_DATASETS = 100; +const MAX_DOCUMENTS = 1000; +const MAX_DISPLAY_DATASETS = 20; +const MAX_DISPLAY_DOCUMENTS = 50; +const MAX_ID_LENGTH = 256; +const MAX_NAME_LENGTH = 256; +const MAX_SCOPE_BYTES = 64 * 1024; + +export type KnowledgeScopeDocument = { id: string; name: string }; + +export type KnowledgeScopeDatasetSelection = { + id: string; + name: string; + documents: + | { mode: "all" } + | { mode: "selected"; items: KnowledgeScopeDocument[] }; +}; + +export type KnowledgeScopeSelection = + | { mode: "all" } + | { mode: "disabled" } + | { mode: "selected"; datasets: KnowledgeScopeDatasetSelection[] }; + +export type KnowledgeScopeSnapshot = { + version: 1; + mode: "all" | "selected" | "disabled"; + dataset_ids?: string[]; + document_filters?: Array<{ dataset_id: string; document_ids: string[] }>; + display?: { + datasets: Array<{ + id: string; + name: string; + documents?: KnowledgeScopeDocument[]; + }>; + }; +}; + +export const ALL_KNOWLEDGE_SCOPE: KnowledgeScopeSelection = { mode: "all" }; + +export function cloneKnowledgeScopeSelection( + selection: KnowledgeScopeSelection, +): KnowledgeScopeSelection { + if (selection.mode !== "selected") return { mode: selection.mode }; + return { + mode: "selected", + datasets: selection.datasets.map((dataset) => ({ + ...dataset, + documents: + dataset.documents.mode === "all" + ? { mode: "all" } + : { + mode: "selected", + items: dataset.documents.items.map((document) => ({ + ...document, + })), + }, + })), + }; +} + +function normalizeId(value: string): string { + const id = value.trim(); + if (!id || codePointLength(id) > MAX_ID_LENGTH) { + throw new Error("Knowledge scope contains an invalid identifier."); + } + return id; +} + +function codePointLength(value: string): number { + return [...value].length; +} + +function stableUnique(items: readonly T[]): T[] { + const result: T[] = []; + const seen = new Set(); + for (const item of items) { + const id = normalizeId(item.id); + if (!seen.has(id)) { + result.push({ ...item, id }); + seen.add(id); + } + } + return result; +} + +function byteLength(value: unknown): number { + return new TextEncoder().encode(JSON.stringify(value)).length; +} + +export function buildKnowledgeScopeSnapshot( + selection: KnowledgeScopeSelection, +): KnowledgeScopeSnapshot { + if (selection.mode !== "selected") { + return { version: 1, mode: selection.mode }; + } + + const datasets = stableUnique(selection.datasets); + if (datasets.length === 0 || datasets.length > MAX_DATASETS) { + throw new Error(`Select between 1 and ${MAX_DATASETS} knowledge bases.`); + } + + const documentFilters: NonNullable< + KnowledgeScopeSnapshot["document_filters"] + > = []; + let documentCount = 0; + for (const dataset of datasets) { + if (dataset.documents.mode !== "selected") continue; + const documents = stableUnique(dataset.documents.items); + if (documents.length === 0) { + throw new Error("Select at least one document or choose all documents."); + } + documentCount += documents.length; + documentFilters.push({ + dataset_id: dataset.id, + document_ids: documents.map((document) => document.id), + }); + } + if (documentCount > MAX_DOCUMENTS) { + throw new Error(`Select at most ${MAX_DOCUMENTS} documents.`); + } + + const snapshot: KnowledgeScopeSnapshot = { + version: 1, + mode: "selected", + dataset_ids: datasets.map((dataset) => dataset.id), + }; + if (documentFilters.length > 0) { + snapshot.document_filters = documentFilters; + } + + const filtersByDataset = new Map( + documentFilters.map((filter) => [ + filter.dataset_id, + new Set(filter.document_ids), + ]), + ); + let displayedDocuments = 0; + const displayDatasets: NonNullable< + NonNullable["datasets"] + > = []; + for (const dataset of datasets.slice(0, MAX_DISPLAY_DATASETS)) { + if (!dataset.name.trim() || codePointLength(dataset.name) > MAX_NAME_LENGTH) + continue; + const entry: (typeof displayDatasets)[number] = { + id: dataset.id, + name: dataset.name, + }; + const allowedDocuments = filtersByDataset.get(dataset.id); + if (allowedDocuments && dataset.documents.mode === "selected") { + const remaining = MAX_DISPLAY_DOCUMENTS - displayedDocuments; + const documents = stableUnique(dataset.documents.items) + .filter( + (document) => + allowedDocuments.has(document.id) && + Boolean(document.name.trim()) && + codePointLength(document.name) <= MAX_NAME_LENGTH, + ) + .slice(0, remaining) + // Catalog entries also carry provider metadata such as `selectable`. + // The message contract intentionally exposes only the stable display + // fields, so do not leak the catalog object into the wire snapshot. + .map(({ id, name }) => ({ id, name })); + if (documents.length > 0) { + entry.documents = documents; + displayedDocuments += documents.length; + } + } + displayDatasets.push(entry); + } + if (displayDatasets.length > 0) { + snapshot.display = { datasets: displayDatasets }; + } + + while (snapshot.display && byteLength(snapshot) > MAX_SCOPE_BYTES) { + let lastWithDocuments: + | NonNullable["datasets"][number] + | undefined; + for ( + let index = snapshot.display.datasets.length - 1; + index >= 0; + index -= 1 + ) { + const candidate = snapshot.display.datasets[index]; + if (candidate?.documents && candidate.documents.length > 0) { + lastWithDocuments = candidate; + break; + } + } + if (lastWithDocuments?.documents) { + lastWithDocuments.documents.pop(); + if (lastWithDocuments.documents.length === 0) { + delete lastWithDocuments.documents; + } + } else { + snapshot.display.datasets.pop(); + } + if (snapshot.display.datasets.length === 0) { + delete snapshot.display; + } + } + if (byteLength(snapshot) > MAX_SCOPE_BYTES) { + throw new Error("The selected knowledge scope is too large."); + } + return snapshot; +} + +export function countKnowledgeScopeSelection( + selection: KnowledgeScopeSelection, +) { + if (selection.mode !== "selected") { + return { datasets: 0, documents: 0 }; + } + return { + datasets: selection.datasets.length, + documents: selection.datasets.reduce( + (total, dataset) => + total + + (dataset.documents.mode === "selected" + ? dataset.documents.items.length + : 0), + 0, + ), + }; +} + +export function readKnowledgeScopeSnapshot( + value: unknown, +): KnowledgeScopeSnapshot | null { + if (!value || typeof value !== "object") return null; + const record = value as Record; + if ( + record.version !== 1 || + (record.mode !== "all" && + record.mode !== "selected" && + record.mode !== "disabled") + ) { + return null; + } + if (record.mode !== "selected") { + return { version: 1, mode: record.mode }; + } + if ( + !Array.isArray(record.dataset_ids) || + record.dataset_ids.length === 0 || + record.dataset_ids.some((id) => typeof id !== "string") + ) { + return null; + } + const rawFilters = record.document_filters; + if ( + rawFilters !== undefined && + (!Array.isArray(rawFilters) || + rawFilters.some( + (filter) => + !filter || + typeof filter !== "object" || + typeof (filter as Record).dataset_id !== "string" || + !Array.isArray((filter as Record).document_ids) || + ((filter as Record).document_ids as unknown[]).some( + (id) => typeof id !== "string", + ), + )) + ) { + return null; + } + + const snapshot: KnowledgeScopeSnapshot = { + version: 1, + mode: "selected", + dataset_ids: [...record.dataset_ids], + }; + if (Array.isArray(rawFilters)) { + snapshot.document_filters = rawFilters.map((filter) => { + const item = filter as { + dataset_id: string; + document_ids: string[]; + }; + return { + dataset_id: item.dataset_id, + document_ids: [...item.document_ids], + }; + }); + } + + const rawDisplay = record.display; + if (rawDisplay && typeof rawDisplay === "object") { + const rawDatasets = (rawDisplay as Record).datasets; + if (Array.isArray(rawDatasets)) { + const datasets = rawDatasets.flatMap((dataset) => { + if (!dataset || typeof dataset !== "object") return []; + const item = dataset as Record; + if (typeof item.id !== "string" || typeof item.name !== "string") { + return []; + } + const displayDataset: NonNullable< + KnowledgeScopeSnapshot["display"] + >["datasets"][number] = { id: item.id, name: item.name }; + if (Array.isArray(item.documents)) { + displayDataset.documents = item.documents.flatMap((document) => { + if (!document || typeof document !== "object") return []; + const displayDocument = document as Record; + return typeof displayDocument.id === "string" && + typeof displayDocument.name === "string" + ? [{ id: displayDocument.id, name: displayDocument.name }] + : []; + }); + } + return [displayDataset]; + }); + if (datasets.length > 0) snapshot.display = { datasets }; + } + } + return snapshot; +} diff --git a/frontend/src/core/threads/hooks.ts b/frontend/src/core/threads/hooks.ts index 3fffb7d30..63b30b674 100644 --- a/frontend/src/core/threads/hooks.ts +++ b/frontend/src/core/threads/hooks.ts @@ -81,6 +81,13 @@ import { export type ThreadStreamOptions = { threadId?: string | null | undefined; displayThreadId?: string | null | undefined; + /** + * Assistant identity sent to the Gateway for run admission and execution. + * Default-chat and sidecar callers use the lead agent; custom-agent pages + * pass their stable agent name so server-side capability checks see the same + * assistant that the runtime loads from the request context. + */ + assistantId?: string; context: LocalSettings["context"]; isMock?: boolean; onSend?: (threadId: string) => void; @@ -1727,6 +1734,7 @@ function isThreadMissingError(error: unknown): boolean { export function useThreadStream({ threadId, displayThreadId, + assistantId = "lead_agent", context, isMock, onSend, @@ -1849,7 +1857,7 @@ export function useThreadStream({ const thread = useStream({ client: getAPIClient(isMock), - assistantId: "lead_agent", + assistantId, threadId: onStreamThreadId, reconnectOnMount: true, fetchStateHistory: { limit: 1 }, @@ -2646,6 +2654,7 @@ export function useThreadStream({ threadId: string, humanMessageId: string, replacementText: string, + additionalKwargs?: Record, ) => { if (!humanMessageId) { return false; @@ -2672,7 +2681,25 @@ export function useThreadStream({ if (!response.ok) { throw new Error(await readResponseErrorMessage(response)); } - return (await response.json()) as EditRegeneratePrepareResponse; + const prepared = + (await response.json()) as EditRegeneratePrepareResponse; + if (!additionalKwargs || !Array.isArray(prepared.input.messages)) { + return prepared; + } + const messages = [...prepared.input.messages]; + for (let index = messages.length - 1; index >= 0; index -= 1) { + const message = messages[index]; + if (message?.type !== "human") continue; + messages[index] = { + ...message, + additional_kwargs: { + ...message.additional_kwargs, + ...additionalKwargs, + }, + }; + break; + } + return { ...prepared, input: { ...prepared.input, messages } }; }, getSupersededMessageIds: (prepared) => prepared.source_message_ids, getOptimisticMessages: (prepared) => prepared.input.messages ?? [], diff --git a/frontend/tests/e2e/agent-chat.spec.ts b/frontend/tests/e2e/agent-chat.spec.ts index cf0ba1203..3cd11af13 100644 --- a/frontend/tests/e2e/agent-chat.spec.ts +++ b/frontend/tests/e2e/agent-chat.spec.ts @@ -91,7 +91,10 @@ test.describe("Agent chat", () => { await textarea.fill("Review this code"); await textarea.press("Enter"); await expect.poll(() => streamBody).toBeDefined(); - expect(streamBody).toMatchObject({ context: { agent_name: "test-agent" } }); + expect(streamBody).toMatchObject({ + assistant_id: "test-agent", + context: { agent_name: "test-agent" }, + }); }); test("agent gallery page loads and shows agents", async ({ page }) => { diff --git a/frontend/tests/e2e/knowledge-scope.spec.ts b/frontend/tests/e2e/knowledge-scope.spec.ts new file mode 100644 index 000000000..6e176aa41 --- /dev/null +++ b/frontend/tests/e2e/knowledge-scope.spec.ts @@ -0,0 +1,216 @@ +import { expect, test } from "@playwright/test"; + +import { handleRunStream, mockLangGraphAPI } from "./utils/mock-api"; + +test.describe("custom-agent knowledge scope", () => { + test("selects one file and sends one immutable scope snapshot", async ({ + page, + }) => { + mockLangGraphAPI(page, { + agents: [ + { + name: "researcher", + description: "Research agent", + tool_groups: ["knowledge"], + }, + ], + features: { knowledgeScopeSelectionEnabled: true }, + }); + await page.route("**/api/knowledge/retrieval-catalog/datasets?*", (route) => + route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items: [{ id: "dataset-1", name: "Policies", selectable: true }], + page: 1, + page_size: 100, + total: 1, + }), + }), + ); + await page.route( + "**/api/knowledge/retrieval-catalog/datasets/dataset-1/documents?*", + (route) => + route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items: [ + { id: "doc-1", name: "Leave.pdf", selectable: true }, + { id: "doc-2", name: "Parsing.pdf", selectable: false }, + ], + page: 1, + page_size: 100, + total: 2, + }), + }), + ); + let streamBody: Record | undefined; + await page.route("**/api/langgraph/threads/*/runs/stream", (route) => { + streamBody = route.request().postDataJSON() as Record; + return handleRunStream(route); + }); + + await page.goto("/workspace/agents/researcher/chats/new"); + await expect(page.getByRole("link", { name: "Knowledge" })).toHaveCount(0); + const trigger = page.getByTestId("knowledge-scope-trigger"); + await expect(trigger).toHaveText(""); + await expect(trigger).toHaveAttribute("aria-pressed", "true"); + await expect(trigger).toHaveClass(/text-foreground/); + await expect(trigger).not.toHaveClass(/bg-primary\/10/); + + await trigger.click(); + await page.getByLabel("Off").check(); + await page.getByRole("button", { name: "Apply" }).click(); + await expect(trigger).toHaveAttribute("aria-pressed", "false"); + await expect(trigger).toHaveAttribute("aria-label", "Knowledge · Off"); + await expect(trigger).not.toHaveClass(/bg-primary\/10/); + + await trigger.click(); + await page.getByLabel("Selected knowledge bases").check(); + await page.getByLabel("Policies").check(); + await page.getByRole("button", { name: "Files" }).click(); + await page.getByLabel("Selected files").check(); + await page.getByLabel("Leave.pdf").check(); + await expect(page.getByLabel("Parsing.pdf")).toBeDisabled(); + await page.getByRole("button", { name: "Apply" }).click(); + await expect(trigger).toHaveAttribute("aria-pressed", "true"); + await expect(trigger).toHaveClass(/text-foreground/); + await expect(trigger).not.toHaveClass(/bg-primary\/10/); + + await page + .getByPlaceholder(/how can i assist you/i) + .fill("Find the policy"); + await page.getByRole("button", { name: "Submit" }).click(); + await expect.poll(() => streamBody).toBeDefined(); + + expect(streamBody).toMatchObject({ + assistant_id: "researcher", + input: { + messages: [ + { + type: "human", + additional_kwargs: { + knowledge_scope: { + version: 1, + mode: "selected", + dataset_ids: ["dataset-1"], + document_filters: [ + { dataset_id: "dataset-1", document_ids: ["doc-1"] }, + ], + display: { + datasets: [ + { + id: "dataset-1", + name: "Policies", + documents: [{ id: "doc-1", name: "Leave.pdf" }], + }, + ], + }, + }, + }, + }, + ], + }, + }); + }); + + test("main chat uses the selector and sends the lead-agent scope", async ({ + page, + }) => { + mockLangGraphAPI(page, { + features: { knowledgeScopeSelectionEnabled: true }, + }); + await page.route("**/api/knowledge/retrieval-catalog/datasets?*", (route) => + route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items: [{ id: "dataset-1", name: "Policies", selectable: true }], + page: 1, + page_size: 100, + total: 1, + }), + }), + ); + await page.route( + "**/api/knowledge/retrieval-catalog/datasets/dataset-1/documents?*", + (route) => + route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items: [{ id: "doc-1", name: "Leave.pdf", selectable: true }], + page: 1, + page_size: 100, + total: 1, + }), + }), + ); + let streamBody: Record | undefined; + await page.route("**/api/langgraph/threads/*/runs/stream", (route) => { + streamBody = route.request().postDataJSON() as Record; + return handleRunStream(route); + }); + + await page.goto("/workspace/chats/new"); + + await expect(page.getByPlaceholder(/how can i assist you/i)).toBeVisible(); + const trigger = page.getByTestId("knowledge-scope-trigger"); + await expect(trigger).toHaveText(""); + await expect(trigger).toHaveAttribute("aria-pressed", "true"); + await trigger.click(); + await page.getByLabel("Selected knowledge bases").check(); + await page.getByLabel("Policies").check(); + await page.getByRole("button", { name: "Files" }).click(); + await page.getByLabel("Selected files").check(); + await page.getByLabel("Leave.pdf").check(); + await page.getByRole("button", { name: "Apply" }).click(); + + await page + .getByPlaceholder(/how can i assist you/i) + .fill("Find the policy"); + await page.getByRole("button", { name: "Submit" }).click(); + await expect.poll(() => streamBody).toBeDefined(); + + expect(streamBody).toMatchObject({ + assistant_id: "lead_agent", + input: { + messages: [ + { + type: "human", + additional_kwargs: { + knowledge_scope: { + version: 1, + mode: "selected", + dataset_ids: ["dataset-1"], + document_filters: [ + { dataset_id: "dataset-1", document_ids: ["doc-1"] }, + ], + }, + }, + }, + ], + }, + }); + }); + + test("main chat hides the selector when configuration disables it", async ({ + page, + }) => { + let catalogRequests = 0; + mockLangGraphAPI(page, { + features: { knowledgeScopeSelectionEnabled: false }, + }); + await page.route("**/api/knowledge/retrieval-catalog/**", (route) => { + catalogRequests += 1; + return route.abort(); + }); + + await page.goto("/workspace/chats/new"); + + await expect(page.getByPlaceholder(/how can i assist you/i)).toBeVisible(); + await expect(page.getByTestId("knowledge-scope-trigger")).toHaveCount(0); + expect(catalogRequests).toBe(0); + }); +}); diff --git a/frontend/tests/e2e/utils/mock-api.ts b/frontend/tests/e2e/utils/mock-api.ts index 5b5154663..c45c13e12 100644 --- a/frontend/tests/e2e/utils/mock-api.ts +++ b/frontend/tests/e2e/utils/mock-api.ts @@ -155,6 +155,7 @@ export type MockAPIOptions = { agentsApiEnabled?: boolean; browserControlEnabled?: boolean; mcpTasksEnabled?: boolean; + knowledgeScopeSelectionEnabled?: boolean; }; runStreamHandler?: (route: Route) => Promise; }; @@ -387,6 +388,8 @@ export function mockLangGraphAPI(page: Page, options?: MockAPIOptions) { agentsApiEnabled: options?.features?.agentsApiEnabled ?? true, browserControlEnabled: options?.features?.browserControlEnabled ?? true, mcpTasksEnabled: options?.features?.mcpTasksEnabled ?? true, + knowledgeScopeSelectionEnabled: + options?.features?.knowledgeScopeSelectionEnabled ?? false, }; const upsertThread = (thread: MockThread) => { @@ -1842,6 +1845,10 @@ export function mockLangGraphAPI(page: Page, options?: MockAPIOptions) { agents_api: { enabled: featureFlags.agentsApiEnabled }, browser_control: { enabled: featureFlags.browserControlEnabled }, mcp_tasks: { enabled: featureFlags.mcpTasksEnabled }, + knowledge_base: { + scope_selection_enabled: + featureFlags.knowledgeScopeSelectionEnabled, + }, }), }); } diff --git a/frontend/tests/unit/components/workspace/knowledge-scope-selector.dom.test.tsx b/frontend/tests/unit/components/workspace/knowledge-scope-selector.dom.test.tsx new file mode 100644 index 000000000..3af98a470 --- /dev/null +++ b/frontend/tests/unit/components/workspace/knowledge-scope-selector.dom.test.tsx @@ -0,0 +1,121 @@ +import { afterEach, describe, expect, it, rs } from "@rstest/core"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { + act, + cleanup, + fireEvent, + render, + screen, + waitFor, +} from "@testing-library/react"; +import type { PropsWithChildren } from "react"; + +import { KnowledgeScopeSelector } from "@/components/workspace/knowledge-scope-selector"; +import { I18nProvider } from "@/core/i18n/context"; +import type { KnowledgeScopeSelection } from "@/core/knowledge"; + +afterEach(() => { + cleanup(); + rs.restoreAllMocks(); +}); + +function renderSelector(selection: KnowledgeScopeSelection) { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + + function Wrapper({ children }: PropsWithChildren) { + return ( + + {children} + + ); + } + + return render( + undefined} + />, + { wrapper: Wrapper }, + ); +} + +function requestUrl(input: RequestInfo | URL): string { + if (typeof input === "string") return input; + if (input instanceof URL) return input.href; + return input.url; +} + +describe("KnowledgeScopeSelector trigger", () => { + it("renders only the icon and stays highlighted while retrieval is active", () => { + renderSelector({ mode: "all" }); + + const trigger = screen.getByRole("button", { name: "Knowledge · All" }); + expect(trigger.textContent).toBe(""); + expect(trigger.getAttribute("aria-pressed")).toBe("true"); + expect(trigger.className).toContain("text-foreground"); + expect(trigger.className).not.toContain("bg-primary/10"); + expect(trigger.className).not.toContain("border-primary/20"); + expect(trigger.querySelector("svg")).not.toBeNull(); + }); + + it("returns to the neutral icon state when retrieval is off", () => { + renderSelector({ mode: "disabled" }); + + const trigger = screen.getByRole("button", { name: "Knowledge · Off" }); + expect(trigger.textContent).toBe(""); + expect(trigger.getAttribute("aria-pressed")).toBe("false"); + expect(trigger.className).not.toContain("bg-primary/10"); + expect(trigger.querySelector("svg")).not.toBeNull(); + }); + + it("does not load documents while an expanded dataset still uses all files", async () => { + const fetch = rs.spyOn(globalThis, "fetch").mockImplementation((input) => { + const url = requestUrl(input); + return Promise.resolve( + Response.json({ + items: url.includes("/documents?") + ? [{ id: "document-1", name: "Guide", selectable: true }] + : [{ id: "dataset-1", name: "Policies", selectable: true }], + page: 1, + page_size: 100, + total: 1, + }), + ); + }); + renderSelector({ + mode: "selected", + datasets: [ + { + id: "dataset-1", + name: "Policies", + documents: { mode: "all" }, + }, + ], + }); + + fireEvent.click(screen.getByRole("button", { name: "Knowledge · 1 base" })); + await screen.findByText("Policies"); + fireEvent.click(screen.getByRole("button", { name: "Files" })); + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + + expect( + fetch.mock.calls.some(([input]) => + requestUrl(input).includes("/documents?"), + ), + ).toBe(false); + + fireEvent.click(screen.getByRole("radio", { name: "Selected files" })); + await waitFor(() => { + expect( + fetch.mock.calls.some(([input]) => + requestUrl(input).includes("/documents?"), + ), + ).toBe(true); + }); + }); +}); diff --git a/frontend/tests/unit/core/agents/features.test.ts b/frontend/tests/unit/core/agents/features.test.ts index a85aad2a5..6ecdf5bf1 100644 --- a/frontend/tests/unit/core/agents/features.test.ts +++ b/frontend/tests/unit/core/agents/features.test.ts @@ -12,6 +12,7 @@ import { fetchAgentsApiEnabled } from "@/core/agents/api"; import { fetch as fetcher } from "@/core/api/fetcher"; import { fetchBrowserControlEnabled, + fetchKnowledgeBaseFeature, fetchMcpTasksEnabled, } from "@/core/features/api"; @@ -115,3 +116,28 @@ describe("fetchMcpTasksEnabled", () => { await expect(fetchMcpTasksEnabled()).rejects.toThrow(); }); }); + +describe("fetchKnowledgeBaseFeature", () => { + test("reads the knowledge-base retrieval-scope flags", async () => { + mockedFetch.mockResolvedValueOnce( + jsonResponse(200, { + agents_api: { enabled: true }, + knowledge_base: { + scope_selection_enabled: true, + }, + }), + ); + await expect(fetchKnowledgeBaseFeature()).resolves.toEqual({ + scopeSelectionEnabled: true, + }); + }); + + test("defaults to disabled when omitted", async () => { + mockedFetch.mockResolvedValueOnce( + jsonResponse(200, { agents_api: { enabled: true } }), + ); + await expect(fetchKnowledgeBaseFeature()).resolves.toEqual({ + scopeSelectionEnabled: false, + }); + }); +}); diff --git a/frontend/tests/unit/core/i18n/translations.test.ts b/frontend/tests/unit/core/i18n/translations.test.ts index 5a3ec532a..acbe27840 100644 --- a/frontend/tests/unit/core/i18n/translations.test.ts +++ b/frontend/tests/unit/core/i18n/translations.test.ts @@ -20,5 +20,6 @@ describe("core copy loading", () => { expect(chinese.channels.descriptions.buzz).toBe( "通过 DeerFlow 智能体接收 Buzz 频道消息和私聊。", ); + expect(chinese.knowledge.scope.title).toBe("知识库范围"); }); }); diff --git a/frontend/tests/unit/core/knowledge/scope.test.ts b/frontend/tests/unit/core/knowledge/scope.test.ts new file mode 100644 index 000000000..92b58bb84 --- /dev/null +++ b/frontend/tests/unit/core/knowledge/scope.test.ts @@ -0,0 +1,159 @@ +import { describe, expect, test } from "@rstest/core"; + +import { + buildKnowledgeScopeSnapshot, + cloneKnowledgeScopeSelection, + readKnowledgeScopeSnapshot, + type KnowledgeScopeSelection, +} from "@/core/knowledge/scope"; + +describe("knowledge scope snapshots", () => { + test("builds selected datasets and file filters with bounded display labels", () => { + const snapshot = buildKnowledgeScopeSnapshot({ + mode: "selected", + datasets: [ + { + id: " dataset-1 ", + name: "Policies", + documents: { + mode: "selected", + items: [ + { id: "doc-1", name: "Leave.pdf" }, + { id: "doc-1", name: "Duplicate.pdf" }, + ], + }, + }, + { + id: "dataset-2", + name: "Handbook", + documents: { mode: "all" }, + }, + ], + }); + + expect(snapshot).toEqual({ + version: 1, + mode: "selected", + dataset_ids: ["dataset-1", "dataset-2"], + document_filters: [{ dataset_id: "dataset-1", document_ids: ["doc-1"] }], + display: { + datasets: [ + { + id: "dataset-1", + name: "Policies", + documents: [{ id: "doc-1", name: "Leave.pdf" }], + }, + { id: "dataset-2", name: "Handbook" }, + ], + }, + }); + }); + + test("strips catalog metadata from displayed document snapshots", () => { + const catalogDocument = { + id: "doc-1", + name: "Leave.pdf", + selectable: true, + }; + const snapshot = buildKnowledgeScopeSnapshot({ + mode: "selected", + datasets: [ + { + id: "dataset-1", + name: "Policies", + documents: { mode: "selected", items: [catalogDocument] }, + }, + ], + }); + + expect(snapshot.display?.datasets[0]?.documents).toEqual([ + { id: "doc-1", name: "Leave.pdf" }, + ]); + }); + + test("rejects an empty selected document filter", () => { + expect(() => + buildKnowledgeScopeSnapshot({ + mode: "selected", + datasets: [ + { + id: "dataset-1", + name: "Policies", + documents: { mode: "selected", items: [] }, + }, + ], + }), + ).toThrow(/at least one document/i); + }); + + test("omits overlong display names without truncating execution ids", () => { + const snapshot = buildKnowledgeScopeSnapshot({ + mode: "selected", + datasets: [ + { + id: "dataset-1", + name: "x".repeat(257), + documents: { mode: "all" }, + }, + ], + }); + + expect(snapshot.dataset_ids).toEqual(["dataset-1"]); + expect(snapshot.display).toBeUndefined(); + }); + + test("counts astral Unicode characters as code points", () => { + const emojiName = "😀".repeat(256); + const snapshot = buildKnowledgeScopeSnapshot({ + mode: "selected", + datasets: [ + { + id: "dataset-1", + name: emojiName, + documents: { mode: "all" }, + }, + ], + }); + + expect(snapshot.display?.datasets[0]?.name).toBe(emojiName); + }); + + test("clones nested selection state for immutable message snapshots", () => { + const selection: KnowledgeScopeSelection = { + mode: "selected", + datasets: [ + { + id: "dataset-1", + name: "Policies", + documents: { + mode: "selected", + items: [{ id: "doc-1", name: "Leave.pdf" }], + }, + }, + ], + }; + const cloned = cloneKnowledgeScopeSelection(selection); + if (selection.mode === "selected") selection.datasets[0]!.name = "Renamed"; + + expect(cloned).not.toBe(selection); + expect(cloned.mode === "selected" && cloned.datasets[0]?.name).toBe( + "Policies", + ); + }); + + test("reads only recognized version-one modes", () => { + expect(readKnowledgeScopeSnapshot({ version: 1, mode: "all" })).toEqual({ + version: 1, + mode: "all", + }); + expect(readKnowledgeScopeSnapshot({ version: 2, mode: "all" })).toBeNull(); + expect( + readKnowledgeScopeSnapshot({ + version: 1, + mode: "selected", + dataset_ids: ["dataset-1"], + document_filters: [{ dataset_id: "dataset-1", document_ids: null }], + }), + ).toBeNull(); + }); +}); diff --git a/frontend/tests/unit/core/threads/stream-options.test.ts b/frontend/tests/unit/core/threads/stream-options.test.ts index cc089cb19..84e8fcaac 100644 --- a/frontend/tests/unit/core/threads/stream-options.test.ts +++ b/frontend/tests/unit/core/threads/stream-options.test.ts @@ -1,6 +1,6 @@ import { afterEach, expect, test, rs } from "@rstest/core"; -async function captureThreadStreamOptions() { +async function captureThreadStreamOptions(assistantId?: string) { let capturedOptions: Record | undefined; rs.resetModules(); @@ -70,6 +70,7 @@ async function captureThreadStreamOptions() { context: { mode: "flash", }, + assistantId, isMock: true, } as never); return null; @@ -97,3 +98,9 @@ test("does not subscribe to unsupported LangGraph events mode", async () => { expect(options).toHaveProperty("onUpdateEvent"); expect(options).toHaveProperty("onCustomEvent"); }); + +test("forwards the custom-agent assistant identity to the stream", async () => { + const options = await captureThreadStreamOptions("researcher"); + + expect(options).toHaveProperty("assistantId", "researcher"); +}); diff --git a/scripts/config-upgrade.sh b/scripts/config-upgrade.sh index 96b0ceaf9..bc036488a 100755 --- a/scripts/config-upgrade.sh +++ b/scripts/config-upgrade.sh @@ -74,6 +74,71 @@ print() # Each migration targets a specific version upgrade. # 'replacements': list of (old_string, new_string) applied to the raw YAML text. # This handles value changes that a dict merge cannot catch. +# 'data_transform': callable applied to the parsed config after text migrations. + +RAGFLOW_PROVIDER_KEYS = ( + 'base_url', + 'api_key', + 'timeout', + 'page_size', + 'similarity_threshold', + 'vector_similarity_weight', + 'top_k', + 'max_chars_per_chunk', + 'max_total_chars', +) + + +def migrate_knowledge_provider_settings(data): + # Move legacy RAGFlow settings to the provider tool and remove them from the generic block. + knowledge_base = data.get('knowledge_base') + tools = data.get('tools') + target = None + has_configured_knowledge_tool = False + if isinstance(tools, list): + has_configured_knowledge_tool = any( + isinstance(tool, dict) and tool.get('group') == 'knowledge' + for tool in tools + ) + target = next( + ( + tool + for tool in tools + if isinstance(tool, dict) + and tool.get('name') == 'knowledge_search' + and tool.get('use') == 'deerflow.community.ragflow.tools:knowledge_search_tool' + ), + None, + ) + + changes = [] + # Before the capability gate shipped, a tools-only knowledge configuration + # was valid and enabled by the presence of the provider tool itself. Preserve + # that provider-neutral behavior when the merge adds the example's + # ``enabled: false`` gate. Explicit operator values still win. + if not isinstance(knowledge_base, dict): + if not has_configured_knowledge_tool: + return changes + knowledge_base = data['knowledge_base'] = {'enabled': True} + changes.append('knowledge_base.enabled set to true (preserved configured knowledge tools)') + + if 'enabled' not in knowledge_base and has_configured_knowledge_tool: + knowledge_base['enabled'] = True + changes.append('knowledge_base.enabled set to true (preserved configured knowledge tools)') + + for key in RAGFLOW_PROVIDER_KEYS: + if key not in knowledge_base: + continue + if target is None: + changes.append(f'knowledge_base.{key} removed (no RAGFlow knowledge_search tool configured)') + elif key in target: + changes.append(f'knowledge_base.{key} removed (tools.knowledge_search.{key} preserved)') + else: + target[key] = knowledge_base[key] + changes.append(f'knowledge_base.{key} -> tools.knowledge_search.{key}') + del knowledge_base[key] + return changes + MIGRATIONS = { 1: { @@ -85,11 +150,10 @@ MIGRATIONS = { ('src.tools.', 'deerflow.tools.'), ], }, - # Future migrations go here: - # 2: { - # 'description': '...', - # 'replacements': [('old', 'new')], - # }, + 46: { + 'description': 'Preserve configured knowledge providers and move RAGFlow settings to the knowledge_search tool', + 'data_transform': migrate_knowledge_provider_settings, + }, } # Apply migrations in order for versions (user_version, example_version] @@ -107,6 +171,13 @@ for version in range(user_version + 1, example_version + 1): # Re-parse after text migrations user = yaml.safe_load(raw_text) or {} +# Apply structured migrations to the parsed config. +for version in range(user_version + 1, example_version + 1): + migration = MIGRATIONS.get(version) + transform = migration.get('data_transform') if migration else None + if transform: + migrated.extend(transform(user)) + if migrated: print(f'Applied {len(migrated)} migration(s):') for m in migrated: