From 480deebf024ed30aa1f850b6b0db8394f18e2a90 Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:38:52 +0800 Subject: [PATCH] fix(knowledge-fs): harden research progress and scoring --- ...ogress-and-structured-scoring-hardening.md | 69 +++ knowledge-fs/apps/api/src/index.ts | 2 +- .../page-index-semantic-tree-search.test.ts | 21 + .../src/page-index-semantic-tree-search.ts | 32 ++ ...-task-progress-database-repository.test.ts | 17 + ...earch-task-progress-database-repository.ts | 4 +- .../api/src/research-task-runtime.test.ts | 112 ++++- .../packages/api/src/research-task-runtime.ts | 28 +- .../src/index.test.ts | 57 +++ .../dify-model-runtime-client/src/index.ts | 9 +- .../src/dify-model-runtime-llm.test.ts | 36 ++ knowledge-fs/packages/generation/src/index.ts | 27 +- .../__tests__/research-task-events.spec.ts | 114 +++++ .../__tests__/retrieval-test-page.spec.tsx | 226 ++++++++- web/features/new-rag/retrieval-test-page.tsx | 437 +++++++++++++----- .../services/knowledge-query-events.ts | 56 ++- .../new-rag/services/research-task-events.ts | 131 ++++++ web/i18n/ar-TN/dataset.json | 6 +- web/i18n/de-DE/dataset.json | 6 +- web/i18n/en-US/dataset.json | 14 +- web/i18n/es-ES/dataset.json | 6 +- web/i18n/fa-IR/dataset.json | 6 +- web/i18n/fr-FR/dataset.json | 6 +- web/i18n/hi-IN/dataset.json | 6 +- web/i18n/id-ID/dataset.json | 6 +- web/i18n/it-IT/dataset.json | 6 +- web/i18n/ja-JP/dataset.json | 6 +- web/i18n/ko-KR/dataset.json | 6 +- web/i18n/nl-NL/dataset.json | 6 +- web/i18n/pl-PL/dataset.json | 6 +- web/i18n/pt-BR/dataset.json | 6 +- web/i18n/ro-RO/dataset.json | 6 +- web/i18n/ru-RU/dataset.json | 6 +- web/i18n/sl-SI/dataset.json | 6 +- web/i18n/th-TH/dataset.json | 6 +- web/i18n/tr-TR/dataset.json | 6 +- web/i18n/uk-UA/dataset.json | 6 +- web/i18n/vi-VN/dataset.json | 6 +- web/i18n/zh-Hans/dataset.json | 6 +- web/i18n/zh-Hant/dataset.json | 6 +- .../images/new-rag/vibe-coding-star.svg | 8 + 41 files changed, 1361 insertions(+), 171 deletions(-) create mode 100644 knowledge-fs/.harness/changes/2026-07-31-research-progress-and-structured-scoring-hardening.md create mode 100644 web/features/new-rag/__tests__/research-task-events.spec.ts create mode 100644 web/features/new-rag/services/research-task-events.ts create mode 100644 web/public/images/new-rag/vibe-coding-star.svg diff --git a/knowledge-fs/.harness/changes/2026-07-31-research-progress-and-structured-scoring-hardening.md b/knowledge-fs/.harness/changes/2026-07-31-research-progress-and-structured-scoring-hardening.md new file mode 100644 index 00000000000..8067cc20bc0 --- /dev/null +++ b/knowledge-fs/.harness/changes/2026-07-31-research-progress-and-structured-scoring-hardening.md @@ -0,0 +1,69 @@ +# Research progress and structured semantic scoring hardening + +Date: 2026-07-31 + +## What changed + +- Routed PageIndex semantic candidate scoring through Dify's structured-output LLM endpoint. + Each bounded scoring batch now supplies a strict JSON schema for the exact candidate ids, while + the existing KnowledgeFS parser continues to reject missing, duplicated, unknown, or invalid + scores. The production batch size was reduced from ten candidates to five. +- Updated the Dify model-runtime client and generation provider to carry structured-output schemas, + consume Dify's validated `structured_output` stream field, and retain free-form response text + only as a compatibility fallback. +- Kept Research task stage transitions synchronized with heartbeat-updated row versions inside the + serialized lease lane, preventing a heartbeat from making the next durable transition use a + stale fence. +- Made the Research progress database repository accept PostgreSQL `BIGINT` timestamps returned as + bounded decimal strings while preserving non-negative safe-integer validation. +- Added Retrieval Test progress streaming through the console-issued Research capability, including + bounded SSE replay, cursor reconnects, terminal task/partial refreshes, actual per-stage timing, + and process-log presentation. +- Split active Research stage labels from completed milestone labels in every supported Web locale. + A streamed completion now refreshes the task first and leaves the existing active-to-completed + transition guard as the single owner of the final partial-result refresh. Historical terminal + replay does not trigger redundant task or partial requests. + +## Why + +- Free-form LLM output could spend tokens on prose or return malformed JSON before semantic scores, + even though Dify already owns tenant-scoped structured-output invocation and validation. +- Research heartbeats increment the durable row version while generation is running. Stage + transitions must read the latest fence from the same serialized lane to avoid false lease loss. +- PostgreSQL drivers commonly expose `BIGINT` values as strings, so valid persisted progress events + could fail when read back. +- Polling alone did not expose durable Research stage history or precise timing, while mixing SSE + terminal refreshes with the existing completion guard caused duplicate partial reads. +- Completed-milestone copy was misleading when reused for a stage that was still running. + +## Verification + +- Web Retrieval Test and Research event-stream focused suites pass. +- KnowledgeFS API semantic scoring, progress database, and Research runtime focused suites pass. +- Dify model-runtime client and generation-provider focused suites pass. +- TypeScript checks pass for `@knowledge/api`, `@knowledge/dify-model-runtime-client`, and + `@knowledge/generation`. +- Focused KnowledgeFS Biome checks, Web i18n synchronization, and `git diff --check` pass. +- The repository-wide Web TypeScript command remains blocked by pre-existing generated + `.next/types/validator.ts` route declarations unrelated to this change. +- Vite+ static `check` remains unavailable in the current local environment because its launcher + cannot resolve the Node binary; focused tests, i18n validation, KnowledgeFS Biome, and package + typechecks were used instead. + +## Performance and operational bounds + +- Semantic scoring still uses at most four concurrent LLM batches, with five candidates per batch, + bounded candidate text, output tokens, response characters, and per-batch timeouts. +- Research SSE responses remain server-bounded by page limit and connection duration. Clients + reconnect from the last durable sequence cursor rather than retaining an unbounded connection. +- Completed tasks perform one final partial refetch through the existing transition guard. Replayed + historical terminal events do not add redundant refresh requests. + +## Risks and follow-up + +- Smaller semantic-score batches increase the number of model calls for the same candidate window, + trading request count for more reliable structured responses within the configured output budget. +- Providers without native structured output still rely on Dify's prompt-based parsing before the + strict KnowledgeFS validation boundary. +- Progress history is held only for Research tasks selected during the current Retrieval Test page + session; switching or reloading reconstructs it from the durable cursor stream. diff --git a/knowledge-fs/apps/api/src/index.ts b/knowledge-fs/apps/api/src/index.ts index 8784a6ead0d..e9c541a182a 100644 --- a/knowledge-fs/apps/api/src/index.ts +++ b/knowledge-fs/apps/api/src/index.ts @@ -149,7 +149,7 @@ const rerankerOptions = createApiRerankerOptions(); const semanticEntityExtractionOptions = createApiSemanticEntityExtractionOptions(); const profileReasoningCapability = createApiProfileReasoningCapability(); const pageIndexSemanticTreeSearch = createPageIndexSemanticTreeSearch({ - batchSize: 10, + batchSize: 5, maxConcurrentBatches: 4, maxOutputTokens: profileReasoningCapability.maxOutputTokens, maxTextCharsPerCandidate: 1_500, diff --git a/knowledge-fs/packages/api/src/page-index-semantic-tree-search.test.ts b/knowledge-fs/packages/api/src/page-index-semantic-tree-search.test.ts index 104e9bb16f4..db96e070d3f 100644 --- a/knowledge-fs/packages/api/src/page-index-semantic-tree-search.test.ts +++ b/knowledge-fs/packages/api/src/page-index-semantic-tree-search.test.ts @@ -65,6 +65,27 @@ describe("PageIndex semantic tree search", () => { }), ); const firstPrompt = generate.mock.calls[0]?.[0].messages; + expect(generate.mock.calls[0]?.[0].structuredOutputSchema).toMatchObject({ + additionalProperties: false, + properties: { + scores: { + items: { + additionalProperties: false, + properties: { + candidateId: { enum: ["c1", "c2"], type: "string" }, + score: { maximum: 1, minimum: 0, type: "number" }, + }, + required: ["candidateId", "score", "reason"], + type: "object", + }, + maxItems: 2, + minItems: 2, + type: "array", + }, + }, + required: ["scores"], + type: "object", + }); expect(firstPrompt?.[0]?.content).toContain("do not score by keyword overlap alone"); expect(firstPrompt?.[1]?.content).toContain("How long must invoices be retained?"); expect(firstPrompt?.[1]?.content).toContain("Finance"); diff --git a/knowledge-fs/packages/api/src/page-index-semantic-tree-search.ts b/knowledge-fs/packages/api/src/page-index-semantic-tree-search.ts index ae81cbe5890..7562d4fa270 100644 --- a/knowledge-fs/packages/api/src/page-index-semantic-tree-search.ts +++ b/knowledge-fs/packages/api/src/page-index-semantic-tree-search.ts @@ -44,6 +44,7 @@ export interface GeneratePageIndexSemanticScoreInput { }[]; readonly model: string; readonly signal?: AbortSignal | undefined; + readonly structuredOutputSchema?: Readonly> | undefined; readonly temperature?: number | undefined; readonly tenantId?: string | undefined; } @@ -159,6 +160,7 @@ export function createPageIndexSemanticTreeSearch({ }), model: input.reasoningModel.model, signal: controller.signal, + structuredOutputSchema: semanticScoreOutputSchema(candidates), temperature: 0, tenantId, }), @@ -228,6 +230,36 @@ function semanticScoreMessages({ ]; } +function semanticScoreOutputSchema( + candidates: readonly PageIndexSemanticCandidate[], +): Readonly> { + return { + additionalProperties: false, + properties: { + scores: { + items: { + additionalProperties: false, + properties: { + candidateId: { + enum: candidates.map((candidate) => candidate.candidateId), + type: "string", + }, + reason: { maxLength: 500, minLength: 1, type: "string" }, + score: { maximum: 1, minimum: 0, type: "number" }, + }, + required: ["candidateId", "score", "reason"], + type: "object", + }, + maxItems: candidates.length, + minItems: candidates.length, + type: "array", + }, + }, + required: ["scores"], + type: "object", + }; +} + function parseSemanticScoreOutput( text: string, candidates: readonly PageIndexSemanticCandidate[], diff --git a/knowledge-fs/packages/api/src/research-task-progress-database-repository.test.ts b/knowledge-fs/packages/api/src/research-task-progress-database-repository.test.ts index e338633956a..134ddcf9e45 100644 --- a/knowledge-fs/packages/api/src/research-task-progress-database-repository.test.ts +++ b/knowledge-fs/packages/api/src/research-task-progress-database-repository.test.ts @@ -131,6 +131,23 @@ it("pages by tenant/task cursor and rejects unbounded reads", async () => { ).rejects.toThrow(/between 1 and 1/u); }); +it("reads PostgreSQL bigint timestamps returned as decimal strings", async () => { + const database = recordingDatabase("postgres", async () => ({ + rows: [progressRow({ created_at: String(Date.parse("2026-07-14T00:00:00.000Z")) })], + rowsAffected: 0, + })); + + await expect( + createRepository(database.adapter).list({ + limit: 1, + researchTaskJobId: JOB_ID, + tenantId: TENANT_ID, + }), + ).resolves.toMatchObject({ + items: [{ createdAt: "2026-07-14T00:00:00.000Z" }], + }); +}); + it("polls the durable ledger independently across replicas and releases subscriber bounds", async () => { const rows: DatabaseRow[] = []; const database = recordingDatabase("postgres", async (input) => { diff --git a/knowledge-fs/packages/api/src/research-task-progress-database-repository.ts b/knowledge-fs/packages/api/src/research-task-progress-database-repository.ts index 125831439e3..ea51dfea52d 100644 --- a/knowledge-fs/packages/api/src/research-task-progress-database-repository.ts +++ b/knowledge-fs/packages/api/src/research-task-progress-database-repository.ts @@ -9,7 +9,7 @@ import type { } from "@knowledge/core"; import { UuidSchema } from "@knowledge/core"; -import { numberColumn, stringColumn } from "./database-row-utils"; +import { nonnegativeSafeIntegerColumn, numberColumn, stringColumn } from "./database-row-utils"; import { databasePlaceholder, quoteDatabaseIdentifier } from "./database-sql-utils"; import { cloneJsonObject, jsonObjectColumn } from "./json-utils"; import type { ResearchTaskJobStage } from "./research-task-job"; @@ -397,7 +397,7 @@ function progressFromRow(row: DatabaseRow): ResearchTaskProgressEvent { if (!eventTypes.has(type) || !jobStages.has(stage)) { throw new Error("Research task progress row has an invalid event type or stage"); } - const createdAt = validTimestamp(numberColumn(row, "created_at")); + const createdAt = validTimestamp(nonnegativeSafeIntegerColumn(row, "created_at")); return { createdAt: new Date(createdAt).toISOString(), id: stringColumn(row, "id"), diff --git a/knowledge-fs/packages/api/src/research-task-runtime.test.ts b/knowledge-fs/packages/api/src/research-task-runtime.test.ts index ad33198623b..d566cc71699 100644 --- a/knowledge-fs/packages/api/src/research-task-runtime.test.ts +++ b/knowledge-fs/packages/api/src/research-task-runtime.test.ts @@ -344,6 +344,99 @@ describe("research task production runtime", () => { }); }); + it("advances with the latest row version after a heartbeat during generation", async () => { + const repository = new MemoryDurableRepository(baseJob()); + let heartbeatObserved: (() => void) | undefined; + const heartbeat = new Promise((resolve) => { + heartbeatObserved = resolve; + }); + const originalHeartbeat = repository.heartbeatExecution.bind(repository); + const heartbeatExecution = vi + .spyOn(repository, "heartbeatExecution") + .mockImplementation(async (input) => { + const updated = await originalHeartbeat(input); + heartbeatObserved?.(); + return updated; + }); + const runtime = createResearchTaskRuntime({ + ...runtimeOptions(repository), + generator: { + stream: async function* () { + await heartbeat; + yield traceStep("query.retrieve"); + yield traceStep("query.answer"); + }, + }, + heartbeatIntervalMs: 1, + leaseMs: 1_000, + now: () => 1_000, + }); + + await expect(runtime.tick()).resolves.toMatchObject({ + leased: 1, + succeeded: 1, + }); + expect(heartbeatExecution).toHaveBeenCalled(); + expect(repository.job.stage).toBe("completed"); + }); + + it("refreshes the stage-transition fence inside the serialized heartbeat lane", async () => { + const repository = new MemoryDurableRepository(baseJob()); + let releaseStageGuard: (() => void) | undefined; + const stageGuardRelease = new Promise((resolve) => { + releaseStageGuard = resolve; + }); + let stageGuardObserved: (() => void) | undefined; + const stageGuard = new Promise((resolve) => { + stageGuardObserved = resolve; + }); + let heartbeatObserved: (() => void) | undefined; + const heartbeat = new Promise((resolve) => { + heartbeatObserved = resolve; + }); + let assertionCount = 0; + const originalHeartbeat = repository.heartbeatExecution.bind(repository); + vi.spyOn(repository, "heartbeatExecution").mockImplementation(async (input) => { + const updated = await originalHeartbeat(input); + heartbeatObserved?.(); + return updated; + }); + const runtime = createResearchTaskRuntime({ + ...runtimeOptions(repository), + deletionFence: { + assertDeletionFenceUnchanged: async () => { + assertionCount += 1; + if (assertionCount === 2) { + stageGuardObserved?.(); + await stageGuardRelease; + } + }, + captureDeletionFence: async (scope) => ({ scope }) as never, + }, + generator: { + stream: async function* () { + yield traceStep("query.retrieve"); + yield traceStep("query.answer"); + }, + }, + heartbeatIntervalMs: 1, + leaseMs: 1_000, + now: () => 1_000, + }); + + const tick = runtime.tick(); + await stageGuard; + await heartbeat; + releaseStageGuard?.(); + + await expect(tick).resolves.toMatchObject({ + leased: 1, + retryScheduled: 0, + succeeded: 1, + }); + expect(repository.job.stage).toBe("completed"); + }); + it.each([ { error: `Research task partial result answer exceeds maxChars=${RESEARCH_TASK_PARTIAL_ANSWER_MAX_CHARS}`, @@ -1210,8 +1303,23 @@ class MemoryDurableRepository implements ResearchTaskDurableRepository { return [structuredClone(this.job)]; } - async heartbeatExecution(): Promise { - throw new Error("Unexpected heartbeat in a sub-interval test"); + async heartbeatExecution( + input: ResearchTaskExecutionFence & { + readonly leaseExpiresAt: number; + readonly workerId: string; + }, + ): Promise { + if (!this.matchesFence(input) || input.workerId !== this.job.workerId) { + return null; + } + this.job = { + ...this.job, + heartbeatAt: input.now, + leaseExpiresAt: input.leaseExpiresAt, + rowVersion: this.job.rowVersion + 1, + updatedAt: input.now, + }; + return structuredClone(this.job); } async advanceExecution( diff --git a/knowledge-fs/packages/api/src/research-task-runtime.ts b/knowledge-fs/packages/api/src/research-task-runtime.ts index ee7bca4347d..5a219b84040 100644 --- a/knowledge-fs/packages/api/src/research-task-runtime.ts +++ b/knowledge-fs/packages/api/src/research-task-runtime.ts @@ -266,10 +266,10 @@ export function createResearchTaskRuntime({ authorizationContext, abortSignal: abortController.signal, capabilityGrants, - current, deletionFence, deletionToken, generator, + getCurrent: () => current, manifests, now, partials, @@ -277,6 +277,9 @@ export function createResearchTaskRuntime({ publishProgress, repository, serialize, + updateCurrent: (updated) => { + current = updated; + }, }); await assertWritable(); const completed = await serialize(() => repository.completeExecution(fence(current, now()))); @@ -481,10 +484,10 @@ async function runResearchTask({ authorizationContext: initialAuthorizationContext, abortSignal, capabilityGrants, - current: initial, deletionFence, deletionToken, generator, + getCurrent, manifests, now, partials, @@ -492,6 +495,7 @@ async function runResearchTask({ publishProgress, repository, serialize, + updateCurrent, }: { readonly access: Pick; readonly allowLegacyProfileFallback: boolean; @@ -500,10 +504,10 @@ async function runResearchTask({ readonly capabilityGrants?: | Pick | undefined; - readonly current: ResearchTaskJob; readonly deletionFence?: DeletionLifecycleFenceGuard | undefined; readonly deletionToken?: DeletionLifecycleFenceToken | undefined; readonly generator: QueryGenerator; + readonly getCurrent: () => ResearchTaskJob; readonly manifests: KnowledgeSpaceManifestRepository; readonly now: () => number; readonly partials: ResearchTaskPartialResultRepository; @@ -515,8 +519,9 @@ async function runResearchTask({ ) => Promise; readonly repository: ResearchTaskDurableRepository; readonly serialize: (operation: () => Promise) => Promise; + readonly updateCurrent: (current: ResearchTaskJob) => void; }): Promise { - let current = initial; + let current = getCurrent(); let authorizationContext = initialAuthorizationContext; const assertWritable = async (): Promise => { if (deletionToken) { @@ -524,6 +529,7 @@ async function runResearchTask({ } }; const revalidate = async () => { + current = getCurrent(); if (abortSignal.aborted) { throw abortSignal.reason ?? new Error("Research task execution lease was lost"); } @@ -534,15 +540,19 @@ async function runResearchTask({ ); }; const advance = async (nextStage: ResearchTaskJobStage) => { - const previousStage = current.stage; await assertWritable(); - const updated = await serialize(() => - repository.advanceExecution({ ...fence(current, now()), nextStage }), - ); + const transition = await serialize(async () => { + current = getCurrent(); + const previousStage = current.stage; + const updated = await repository.advanceExecution({ ...fence(current, now()), nextStage }); + return { previousStage, updated }; + }); + const { previousStage, updated } = transition; if (!updated) { throw new Error("Research task stage transition lost its lease fence"); } current = updated; + updateCurrent(updated); await assertWritable(); await publishProgress(updated, "research_task.stage_changed", { previousStage }); }; @@ -684,7 +694,7 @@ async function runResearchTask({ tenantId: current.tenantId, }); } - return current; + return getCurrent(); } function durableRequestedMode( diff --git a/knowledge-fs/packages/dify-model-runtime-client/src/index.test.ts b/knowledge-fs/packages/dify-model-runtime-client/src/index.test.ts index b85ca284624..bd05cf362c5 100644 --- a/knowledge-fs/packages/dify-model-runtime-client/src/index.test.ts +++ b/knowledge-fs/packages/dify-model-runtime-client/src/index.test.ts @@ -140,6 +140,63 @@ describe("createDifyModelRuntimeClient", () => { expect(fetchImpl).toHaveBeenCalledOnce(); }); + it("routes structured LLM requests through Dify's structured-output endpoint", async () => { + let body: Record | undefined; + let url: string | undefined; + const responseBody = new ReadableStream({ + start(controller) { + controller.enqueue( + frame({ + data: { + delta: { finish_reason: "stop" }, + structured_output: { scores: [{ candidateId: "c1", score: 1 }] }, + }, + error: "", + }), + ); + controller.close(); + }, + }); + const client = createDifyModelRuntimeClient({ + apiKey: "inner-secret", + baseUrl: "http://api:5001", + fetch: vi.fn(async (input, init) => { + body = JSON.parse(String(init?.body)) as Record; + url = String(input); + return new Response(responseBody, { + headers: { "content-type": "text/event-stream" }, + status: 200, + }); + }), + }); + const structuredOutputSchema = { + properties: { scores: { type: "array" } }, + required: ["scores"], + type: "object", + }; + + const chunks: unknown[] = []; + for await (const chunk of client.invokeLlm({ + model: "openrouter/auto", + pluginId: "langgenius/openrouter", + promptMessages: [{ content: "score", role: "user" }], + provider: "openrouter", + structuredOutputSchema, + tenantId: "tenant-1", + })) { + chunks.push(chunk); + } + + expect(url).toBe("http://api:5001/inner/api/invoke/llm/structured-output"); + expect(body).toMatchObject({ structured_output_schema: structuredOutputSchema }); + expect(chunks).toEqual([ + { + delta: { finish_reason: "stop" }, + structured_output: { scores: [{ candidateId: "c1", score: 1 }] }, + }, + ]); + }); + it("lists tenant-active models through the Dify catalog endpoint", async () => { let body: unknown; const item = { diff --git a/knowledge-fs/packages/dify-model-runtime-client/src/index.ts b/knowledge-fs/packages/dify-model-runtime-client/src/index.ts index 93d81e5ba3e..7a6cdafc0d7 100644 --- a/knowledge-fs/packages/dify-model-runtime-client/src/index.ts +++ b/knowledge-fs/packages/dify-model-runtime-client/src/index.ts @@ -50,6 +50,7 @@ export interface DifyLlmInput extends DifyModelRequestContext { readonly completionParams?: Readonly> | undefined; readonly promptMessages: readonly unknown[]; readonly stop?: readonly string[] | undefined; + readonly structuredOutputSchema?: Readonly> | undefined; readonly tools?: readonly unknown[] | undefined; } @@ -209,9 +210,12 @@ export function createDifyModelRuntimeClient( async *invokeLlm(input) { const deadline = createDeadline(input.signal, requestTimeoutMs); try { + const structuredOutputSchema = input.structuredOutputSchema; const response = await request( fetchImpl, - `${baseUrl}/inner/api/invoke/llm`, + `${baseUrl}/inner/api/invoke/llm${ + structuredOutputSchema === undefined ? "" : "/structured-output" + }`, apiKey, { ...commonPayload(input), @@ -221,6 +225,9 @@ export function createDifyModelRuntimeClient( prompt_messages: [...input.promptMessages], stop: [...(input.stop ?? [])], stream: true, + ...(structuredOutputSchema === undefined + ? {} + : { structured_output_schema: { ...structuredOutputSchema } }), tools: [...(input.tools ?? [])], }, deadline.signal, diff --git a/knowledge-fs/packages/generation/src/dify-model-runtime-llm.test.ts b/knowledge-fs/packages/generation/src/dify-model-runtime-llm.test.ts index 644f55f62d4..8ab564e7716 100644 --- a/knowledge-fs/packages/generation/src/dify-model-runtime-llm.test.ts +++ b/knowledge-fs/packages/generation/src/dify-model-runtime-llm.test.ts @@ -119,6 +119,42 @@ describe("Dify model runtime LLM provider", () => { }); }); + it("prefers Dify's validated structured output over free-form response text", async () => { + let captured: DifyLlmInput | undefined; + const provider = createDifyModelRuntimeLlmProvider({ + ...BASE, + client: fakeClient( + () => [ + deltaChunk("I will explain before returning JSON."), + { + delta: { finish_reason: "stop" }, + structured_output: { + scores: [{ candidateId: "c1", reason: "direct", score: 1 }], + }, + }, + ], + (input) => { + captured = input; + }, + ), + }); + const structuredOutputSchema = { + properties: { scores: { type: "array" } }, + required: ["scores"], + type: "object", + }; + + const result = await provider.generate({ + messages: [{ content: "Score every candidate.", role: "user" }], + model: BASE.model, + structuredOutputSchema, + tenantId: "tenant-abc", + }); + + expect(captured?.structuredOutputSchema).toEqual(structuredOutputSchema); + expect(result.text).toBe('{"scores":[{"candidateId":"c1","reason":"direct","score":1}]}'); + }); + it("keeps content from Dify stream frames with null usage", async () => { const provider = createDifyModelRuntimeLlmProvider({ ...BASE, diff --git a/knowledge-fs/packages/generation/src/index.ts b/knowledge-fs/packages/generation/src/index.ts index 2a66c16df1a..03cfd6d6ad5 100644 --- a/knowledge-fs/packages/generation/src/index.ts +++ b/knowledge-fs/packages/generation/src/index.ts @@ -27,6 +27,7 @@ export interface GenerateTextInput { readonly messages: readonly LlmMessage[]; readonly model: string; readonly signal?: AbortSignal; + readonly structuredOutputSchema?: Readonly>; readonly temperature?: number; /** Tenant scope for Dify-managed model routing. */ readonly tenantId?: string; @@ -750,6 +751,7 @@ const DifyModelRuntimeLlmChunkSchema = z.object({ .partial() .optional(), model: z.string().optional(), + structured_output: z.unknown().optional(), }); /** @@ -792,6 +794,8 @@ export function createDifyModelRuntimeLlmProvider( const maxTokens = input.maxOutputTokens ?? options.maxOutputTokens; let finishReason = "stop"; + let structuredOutput: unknown; + let structuredOutputFallbackText = ""; let usage: LlmUsage | undefined; for await (const chunk of options.client.invokeLlm({ @@ -806,6 +810,9 @@ export function createDifyModelRuntimeLlmProvider( role: message.role, })), provider: options.provider, + ...(input.structuredOutputSchema === undefined + ? {} + : { structuredOutputSchema: input.structuredOutputSchema }), tenantId, ...(input.signal ? { signal: input.signal } : {}), })) { @@ -818,7 +825,15 @@ export function createDifyModelRuntimeLlmProvider( const content = parsed.data.delta?.message?.content; if (content) { - yield { delta: content, type: "delta" }; + if (input.structuredOutputSchema === undefined) { + yield { delta: content, type: "delta" }; + } else { + structuredOutputFallbackText += content; + } + } + + if (parsed.data.structured_output !== undefined) { + structuredOutput = parsed.data.structured_output; } const chunkUsage = parsed.data.delta?.usage; @@ -842,6 +857,16 @@ export function createDifyModelRuntimeLlmProvider( } } + if (input.structuredOutputSchema !== undefined) { + const text = + structuredOutput === undefined + ? structuredOutputFallbackText + : JSON.stringify(structuredOutput); + if (text) { + yield { delta: text, type: "delta" }; + } + } + yield { finishReason, // Dify invokes the exact tenant-scoped catalog route requested above. Individual diff --git a/web/features/new-rag/__tests__/research-task-events.spec.ts b/web/features/new-rag/__tests__/research-task-events.spec.ts new file mode 100644 index 00000000000..b5ddcc1a755 --- /dev/null +++ b/web/features/new-rag/__tests__/research-task-events.spec.ts @@ -0,0 +1,114 @@ +import type { KnowledgeFsStreamCapabilityResponse } from '@dify/contracts/api/console/knowledge-fs/types.gen' +import { streamResearchTaskEvents } from '../services/research-task-events' + +describe('streamResearchTaskEvents', () => { + it('gets the capability stream and validates progress events', async () => { + const encoder = new TextEncoder() + const responseBody = new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode( + 'id: 1\nevent: research_task.progress\ndata: {"createdAt":"2026-07-31T10:00:00.000Z","id":"event-1","payload":{},"researchTaskJobId":"task-1","sequence":1,"stage":"planning","type":"research_task.started"}\n\n', + ), + ) + controller.enqueue( + encoder.encode( + 'id: 2\nevent: completed\ndata: {"createdAt":"2026-07-31T10:00:12.000Z","id":"event-2","payload":{},"researchTaskJobId":"task-1","sequence":2,"stage":"completed","type":"research_task.stage_changed"}\n\n', + ), + ) + controller.close() + }, + }) + const fetchMock = vi + .spyOn(globalThis, 'fetch') + .mockResolvedValue(new Response(responseBody, { status: 200 })) + const events: unknown[] = [] + const capability: KnowledgeFsStreamCapabilityResponse = { + expires_at: '2026-07-31T10:05:00.000Z', + operation_id: 'streamResearchTask', + token: `header.${btoa(JSON.stringify({ trace_id: 'research-trace-1' }))}.signature`, + url: 'https://knowledge.example.test/research/events', + } + + const result = await streamResearchTaskEvents({ + capability, + onEvent: (event) => events.push(event), + }) + + expect(fetchMock).toHaveBeenCalledWith( + capability.url, + expect.objectContaining({ + credentials: 'omit', + headers: expect.objectContaining({ + Accept: 'text/event-stream', + Authorization: `Bearer ${capability.token}`, + 'X-Trace-ID': 'research-trace-1', + }), + method: 'GET', + }), + ) + expect(events).toEqual([ + { + createdAt: '2026-07-31T10:00:00.000Z', + id: 'event-1', + payload: {}, + researchTaskJobId: 'task-1', + sequence: 1, + stage: 'planning', + type: 'research_task.started', + }, + { + createdAt: '2026-07-31T10:00:12.000Z', + id: 'event-2', + payload: {}, + researchTaskJobId: 'task-1', + sequence: 2, + stage: 'completed', + type: 'research_task.stage_changed', + }, + ]) + expect(result).toEqual({ cursor: '2', reconnect: false, terminal: true }) + + fetchMock.mockRestore() + }) + + it('resumes from the latest cursor after a timeout', async () => { + const encoder = new TextEncoder() + const responseBody = new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode( + 'id: 3\nevent: research_task.progress\ndata: {"createdAt":"2026-07-31T10:00:20.000Z","id":"event-3","payload":{},"researchTaskJobId":"task-1","sequence":3,"stage":"retrieving","type":"research_task.stage_changed"}\n\n', + ), + ) + controller.enqueue( + encoder.encode('event: timeout\ndata: {"cursor":"3","researchTaskJobId":"task-1"}\n\n'), + ) + controller.close() + }, + }) + const fetchMock = vi + .spyOn(globalThis, 'fetch') + .mockResolvedValue(new Response(responseBody, { status: 200 })) + const capability: KnowledgeFsStreamCapabilityResponse = { + expires_at: '2026-07-31T10:05:00.000Z', + operation_id: 'streamResearchTask', + token: `header.${btoa(JSON.stringify({ trace_id: 'research-trace-1' }))}.signature`, + url: 'https://knowledge.example.test/research/events?knowledgeSpaceId=space-1', + } + + const result = await streamResearchTaskEvents({ + capability, + cursor: '2', + onEvent: vi.fn(), + }) + + expect(fetchMock).toHaveBeenCalledWith( + `${capability.url}&cursor=2`, + expect.objectContaining({ method: 'GET' }), + ) + expect(result).toEqual({ cursor: '3', reconnect: true, terminal: false }) + + fetchMock.mockRestore() + }) +}) diff --git a/web/features/new-rag/__tests__/retrieval-test-page.spec.tsx b/web/features/new-rag/__tests__/retrieval-test-page.spec.tsx index 203f78e1354..9a9a7d0fa34 100644 --- a/web/features/new-rag/__tests__/retrieval-test-page.spec.tsx +++ b/web/features/new-rag/__tests__/retrieval-test-page.spec.tsx @@ -15,7 +15,10 @@ const apiMock = vi.hoisted(() => ({ refetchPartials: vi.fn(), refetchTasks: vi.fn(), refetchTraces: vi.fn(), + researchTasks: [] as Array>, + streamCapability: vi.fn(), streamQuery: vi.fn(), + streamResearchEvents: vi.fn(), documentReferences: {} as Record, evidence: undefined as Record | undefined, traceDetail: undefined as Record | undefined, @@ -36,6 +39,10 @@ vi.mock('../services/knowledge-query-events', () => ({ streamKnowledgeQuery: apiMock.streamQuery, })) +vi.mock('../services/research-task-events', () => ({ + streamResearchTaskEvents: apiMock.streamResearchEvents, +})) + vi.mock('@tanstack/react-query', async (importOriginal) => { const original = await importOriginal() return { @@ -65,7 +72,7 @@ vi.mock('@tanstack/react-query', async (importOriginal) => { } if (resource === 'tasks') return { - data: { data: [] }, + data: { data: apiMock.researchTasks }, refetch: apiMock.refetchTasks, } if (resource === 'partials') @@ -81,6 +88,11 @@ vi.mock('@tanstack/react-query', async (importOriginal) => { vi.mock('@/service/client', () => ({ consoleClient: { knowledgeFs: { + tasks: { + byTaskId: { + streamCapability: { post: apiMock.streamCapability }, + }, + }, spaces: { byControlSpaceId: { queries: { admission: { post: apiMock.queryAdmission } }, @@ -178,7 +190,18 @@ describe('RetrievalTestPage', () => { }) apiMock.refetchTasks.mockResolvedValue(undefined) apiMock.refetchTraces.mockResolvedValue(undefined) + apiMock.researchTasks = [] + apiMock.streamCapability.mockResolvedValue({ + expires_at: '2026-07-31T10:30:00.000Z', + operation_id: 'streamResearchTask', + token: 'capability-token', + url: 'https://knowledge.example.test/research/events', + }) apiMock.streamQuery.mockResolvedValue(undefined) + apiMock.streamResearchEvents.mockResolvedValue({ + reconnect: false, + terminal: false, + }) apiMock.queryAdmission.mockResolvedValue({}) apiMock.createBadCase.mockResolvedValue({ id: 'bad-case-1' }) apiMock.createGolden.mockResolvedValue({ id: 'golden-1' }) @@ -228,6 +251,205 @@ describe('RetrievalTestPage', () => { }) }) + it('replays research progress events and shows actual stage durations', async () => { + apiMock.researchTasks = [ + { + completed_at: 1_800_000_025, + cost: {}, + created_at: 1_800_000_000, + id: 'research-completed', + knowledge_space_id: 'space-1', + metadata: {}, + mode: 'research', + query: 'Compare the refund policies', + stage: 'completed', + updated_at: 1_800_000_025, + }, + ] + apiMock.streamResearchEvents.mockImplementation( + async ({ + onEvent, + }: { + onEvent: (event: { + createdAt: string + id: string + payload: Record + researchTaskJobId: string + sequence: number + stage: string + type: string + }) => void + }) => { + const stages = [ + ['planning', 0], + ['retrieving', 2], + ['analyzing', 7], + ['generating', 14], + ['completed', 25], + ] as const + stages.forEach(([stage, seconds], index) => + onEvent({ + createdAt: new Date(1_800_000_000_000 + seconds * 1000).toISOString(), + id: `event-${index + 1}`, + payload: {}, + researchTaskJobId: 'research-completed', + sequence: index + 1, + stage, + type: index ? 'research_task.stage_changed' : 'research_task.started', + }), + ) + return { cursor: '5', reconnect: false, terminal: true } + }, + ) + const user = userEvent.setup() + renderPage() + + await user.click(screen.getByText('Compare the refund policies')) + expect( + screen.getByRole('button', { + name: 'dataset.newKnowledge.retrievalTest.processLog', + }), + ).toHaveAttribute('aria-pressed', 'true') + + await waitFor(() => + expect(apiMock.streamCapability).toHaveBeenCalledWith({ + body: { control_space_id: 'space-1' }, + params: { task_id: 'research-completed' }, + }), + ) + expect(apiMock.streamResearchEvents).toHaveBeenCalledWith( + expect.objectContaining({ + capability: expect.objectContaining({ operation_id: 'streamResearchTask' }), + }), + ) + expect(await screen.findByText('2s')).toBeInTheDocument() + expect(screen.getByText('5s')).toBeInTheDocument() + expect(screen.getByText('7s')).toBeInTheDocument() + expect(screen.getByText('11s')).toBeInTheDocument() + expect(apiMock.refetchTasks).not.toHaveBeenCalled() + expect(apiMock.refetchPartials).not.toHaveBeenCalled() + }) + + it('reconnects an active research event stream from its latest cursor', async () => { + apiMock.researchTasks = [ + { + cost: {}, + created_at: 1_800_000_000, + id: 'research-active', + knowledge_space_id: 'space-1', + metadata: {}, + mode: 'research', + query: 'Compare the refund policies', + stage: 'retrieving', + updated_at: 1_800_000_005, + }, + ] + apiMock.streamResearchEvents + .mockResolvedValueOnce({ cursor: '3', reconnect: true, terminal: false }) + .mockResolvedValueOnce({ cursor: '4', reconnect: false, terminal: true }) + const user = userEvent.setup() + renderPage() + + await user.click(screen.getByText('Compare the refund policies')) + + await waitFor(() => expect(apiMock.streamCapability).toHaveBeenCalledTimes(2)) + expect(apiMock.streamResearchEvents).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ cursor: '3' }), + ) + }) + + it('refreshes final partials only once when a streamed task completes', async () => { + const activeTask = { + cost: {}, + created_at: 1_800_000_000, + id: 'research-active', + knowledge_space_id: 'space-1', + metadata: {}, + mode: 'research', + query: 'Compare the refund policies', + stage: 'retrieving', + updated_at: 1_800_000_005, + } + apiMock.researchTasks = [activeTask] + apiMock.refetchTasks.mockImplementation(async () => { + apiMock.researchTasks = [ + { + ...activeTask, + completed_at: 1_800_000_025, + stage: 'completed', + updated_at: 1_800_000_025, + }, + ] + }) + apiMock.streamResearchEvents.mockImplementation( + async ({ + onEvent, + }: { + onEvent: (event: { + createdAt: string + id: string + payload: Record + researchTaskJobId: string + sequence: number + stage: string + type: string + }) => void + }) => { + onEvent({ + createdAt: '2027-01-15T08:00:25.000Z', + id: 'event-completed', + payload: {}, + researchTaskJobId: 'research-active', + sequence: 5, + stage: 'completed', + type: 'research_task.stage_changed', + }) + return { cursor: '5', reconnect: false, terminal: true } + }, + ) + const user = userEvent.setup() + renderPage() + + await user.click(screen.getByText('Compare the refund policies')) + + await waitFor(() => expect(apiMock.refetchTasks).toHaveBeenCalledOnce()) + await waitFor(() => expect(apiMock.refetchPartials).toHaveBeenCalledOnce()) + }) + + it('does not let the composer shortcut bypass an active research task', async () => { + apiMock.researchTasks = [ + { + cost: {}, + created_at: 1_800_000_000, + id: 'research-active', + knowledge_space_id: 'space-1', + metadata: {}, + mode: 'research', + query: 'Compare the refund policies', + stage: 'retrieving', + updated_at: 1_800_000_005, + }, + ] + const user = userEvent.setup() + renderPage() + + expect( + screen.getByText('dataset.newKnowledge.retrievalTest.retrievingActive · 2/4'), + ).toBeInTheDocument() + await user.click(screen.getByText('Compare the refund policies')) + const queryInput = screen.getByLabelText('dataset.newKnowledge.retrievalTest.queryPlaceholder') + expect( + screen.getByRole('button', { name: 'dataset.newKnowledge.retrievalTest.startResearch' }), + ).toBeDisabled() + + await user.click(queryInput) + await user.keyboard('{Control>}{Enter}{/Control}') + + expect(apiMock.planResearch).not.toHaveBeenCalled() + expect(apiMock.queryAdmission).not.toHaveBeenCalled() + }) + it('persists a selected trace as a production bad case', async () => { apiMock.traces = [ { @@ -333,7 +555,7 @@ describe('RetrievalTestPage', () => { expect( screen.getByRole('heading', { - name: 'dataset.newKnowledge.retrievalTest.result:{"mode":"dataset.newKnowledge.settings.retrievalMode.deep"}', + name: 'dataset.newKnowledge.retrievalTest.result', }), ).toBeInTheDocument() await user.click( diff --git a/web/features/new-rag/retrieval-test-page.tsx b/web/features/new-rag/retrieval-test-page.tsx index 9905bfe5f74..4057aaa4957 100644 --- a/web/features/new-rag/retrieval-test-page.tsx +++ b/web/features/new-rag/retrieval-test-page.tsx @@ -4,15 +4,18 @@ import type { KnowledgeFsResearchTaskPlanResponse, KnowledgeFsResearchTaskResponse, } from '@dify/contracts/api/console/knowledge-fs/types.gen' +import type { Hotkey } from '@tanstack/react-hotkeys' import type { RetrievalEvidence, RetrievalTestMode, RetrievalTestRecord, } from './retrieval-test-model' import type { KnowledgeQueryEvent } from './services/knowledge-query-events' +import type { ResearchTaskProgressEvent } from './services/research-task-events' import { Button } from '@langgenius/dify-ui/button' import { cn } from '@langgenius/dify-ui/cn' import { toast } from '@langgenius/dify-ui/toast' +import { matchesKeyboardEvent } from '@tanstack/react-hotkeys' import { useQuery, useQueryClient } from '@tanstack/react-query' import { useEffect, useMemo, useRef, useState } from 'react' import { useTranslation } from 'react-i18next' @@ -30,6 +33,7 @@ import { } from './retrieval-test-model' import { newKnowledgeDocumentDetailPath, newKnowledgeQualityPath } from './routes' import { streamKnowledgeQuery } from './services/knowledge-query-events' +import { streamResearchTaskEvents } from './services/research-task-events' type LocalQueryRun = { endedAt?: number @@ -51,11 +55,22 @@ type SelectedRun = { type QualityDecision = 'bad-case' | 'golden' const researchStageOrder = ['planning', 'retrieving', 'analyzing', 'generating'] as const +type ResearchStage = (typeof researchStageOrder)[number] +const runRetrievalHotkey = 'Mod+Enter' satisfies Hotkey function timeValue(value: number) { return value < 10_000_000_000 ? value * 1000 : value } +function formatRecordTime(value: number) { + return new Intl.DateTimeFormat(undefined, { + day: 'numeric', + hour: '2-digit', + minute: '2-digit', + month: 'short', + }).format(value) +} + function useClock(enabled: boolean) { const [now, setNow] = useState(() => Date.now()) useEffect(() => { @@ -66,6 +81,72 @@ function useClock(enabled: boolean) { return now } +function researchStageIndex(stage: KnowledgeFsResearchTaskResponse['stage']) { + if (stage === 'queued' || stage === 'paused') return 0 + if (stage === 'completed') return researchStageOrder.length + return researchStageOrder.findIndex((item) => item === stage) +} + +function estimatedStageDuration( + plan: KnowledgeFsResearchTaskPlanResponse | undefined, + stage: ResearchStage, +) { + if (!plan) return undefined + const stepNames: Record> = { + analyzing: new Set(['analyze']), + generating: new Set(['generate']), + planning: new Set(['plan']), + retrieving: new Set(['inspect', 'retrieve']), + } + const milliseconds = plan.steps.reduce((total, step) => { + if (!stepNames[stage].has(typeof step.name === 'string' ? step.name : '')) return total + return total + (typeof step.estimatedLatencyMs === 'number' ? step.estimatedLatencyMs : 0) + }, 0) + return milliseconds > 0 ? formatDuration(milliseconds) : undefined +} + +function researchProgressTime(event: ResearchTaskProgressEvent) { + return Date.parse(event.createdAt) +} + +function actualStageDuration( + events: ResearchTaskProgressEvent[], + stage: ResearchStage, + task: KnowledgeFsResearchTaskResponse, + now: number, +) { + const start = events.find((event) => event.stage === stage) + if (!start) return + const startedAt = researchProgressTime(start) + const next = events.find((event) => { + if (event.sequence <= start.sequence) return false + return ( + event.stage === 'canceled' || + event.stage === 'completed' || + event.stage === 'failed' || + (researchStageOrder.includes(event.stage as ResearchStage) && event.stage !== stage) + ) + }) + const endedAt = next + ? researchProgressTime(next) + : task.stage === stage + ? now + : task.completed_at + ? timeValue(task.completed_at) + : undefined + if (endedAt === undefined || endedAt < startedAt) return + return formatDuration(endedAt - startedAt) +} + +function mergeResearchProgressEvent( + events: ResearchTaskProgressEvent[], + event: ResearchTaskProgressEvent, +) { + const next = events.filter((candidate) => candidate.sequence !== event.sequence) + next.push(event) + return next.sort((left, right) => left.sequence - right.sequence) +} + function ScorePill({ score }: { score: number }) { const normalized = Math.max(0, Math.min(1, score)) return ( @@ -132,7 +213,7 @@ function EvidenceCard({ {evidence.score !== undefined && } -

+

{evidence.text}

{evidence.images.length > 0 && ( @@ -155,7 +236,7 @@ function EvidenceCard({ )} )} -