mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-14 08:36:22 +08:00
feat(server): refactor copilot (#14892)
#### PR Dependency Tree * **PR #14892** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)
This commit is contained in:
DarkSky
@@ -0,0 +1,346 @@
|
||||
import { Injectable, Optional } from '@nestjs/common';
|
||||
|
||||
import { Models } from '../../../models';
|
||||
import type { AiActionRunStatus } from '../../../models/copilot-action-run';
|
||||
import {
|
||||
type NativeActionEvent,
|
||||
type NativeActionRuntimeInput,
|
||||
runNativeActionRecipePreparedStream,
|
||||
} from '../../../native';
|
||||
import type {
|
||||
CopilotImageOptions,
|
||||
CopilotProviderType,
|
||||
CopilotStructuredOptions,
|
||||
PromptMessage,
|
||||
} from '../providers/types';
|
||||
import type { ChatSession } from '../session';
|
||||
import {
|
||||
projectActionResultToAssistantTurn,
|
||||
summarizeActionResult,
|
||||
} from './action-output-projector';
|
||||
import {
|
||||
buildStructuredResponseFromSchemaJson,
|
||||
type RequiredStructuredOutputContract,
|
||||
} from './contracts';
|
||||
import { ExecutionPlanBuilder } from './execution-plan';
|
||||
import { TurnPersistence } from './hosts/turn-persistence';
|
||||
|
||||
type ActionRuntimeBridgeNativeInput = Omit<
|
||||
NativeActionRuntimeInput,
|
||||
'recipeId' | 'recipeVersion'
|
||||
>;
|
||||
|
||||
export type ActionRuntimeBridgeInput = {
|
||||
userId: string;
|
||||
workspaceId: string;
|
||||
docId?: string | null;
|
||||
session?: ChatSession;
|
||||
userMessageId?: string | null;
|
||||
compatSubmissionId?: string | null;
|
||||
actionId: string;
|
||||
actionVersion: string;
|
||||
attempt?: number;
|
||||
retryOf?: string | null;
|
||||
inputSnapshot?: unknown;
|
||||
nativeInput?: ActionRuntimeBridgeNativeInput;
|
||||
onRunCreated?: (
|
||||
context: ActionRuntimeBridgeRunContext
|
||||
) => Promise<void> | void;
|
||||
prepareStructuredRoutes?: {
|
||||
stepId?: string;
|
||||
modelId?: string;
|
||||
messages: PromptMessage[];
|
||||
options?: CopilotStructuredOptions;
|
||||
prefer?: CopilotProviderType;
|
||||
responseSchemaJson?: Record<string, unknown>;
|
||||
responseContract?: RequiredStructuredOutputContract;
|
||||
};
|
||||
prepareImageRoutes?: {
|
||||
stepId?: string;
|
||||
modelId?: string;
|
||||
messages: PromptMessage[];
|
||||
options?: CopilotImageOptions;
|
||||
prefer?: CopilotProviderType;
|
||||
};
|
||||
persistAttachment?: (attachment: unknown) => Promise<unknown> | unknown;
|
||||
signal?: AbortSignal;
|
||||
};
|
||||
|
||||
export type ActionRuntimeBridgeEvent = NativeActionEvent & {
|
||||
runId: string;
|
||||
};
|
||||
|
||||
export type ActionRuntimeBridgeRunContext = {
|
||||
runId: string;
|
||||
attempt: number;
|
||||
};
|
||||
|
||||
function extractResultArtifacts(result: unknown) {
|
||||
if (!result || typeof result !== 'object') {
|
||||
return [];
|
||||
}
|
||||
const value = result as { artifacts?: unknown; attachments?: unknown };
|
||||
if (Array.isArray(value.artifacts)) {
|
||||
return value.artifacts;
|
||||
}
|
||||
if (Array.isArray(value.attachments)) {
|
||||
return value.attachments;
|
||||
}
|
||||
return [];
|
||||
}
|
||||
|
||||
function resolveFinalStatus(
|
||||
event: NativeActionEvent | undefined,
|
||||
signal?: AbortSignal
|
||||
): Extract<AiActionRunStatus, 'succeeded' | 'failed' | 'aborted'> {
|
||||
if (signal?.aborted || event?.status === 'aborted') {
|
||||
return 'aborted';
|
||||
}
|
||||
if (event?.type === 'action_done' && event.status === 'succeeded') {
|
||||
return 'succeeded';
|
||||
}
|
||||
return 'failed';
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class ActionRuntimeBridge {
|
||||
constructor(
|
||||
private readonly models: Models,
|
||||
private readonly turnPersistence: TurnPersistence,
|
||||
@Optional() private readonly plans?: ExecutionPlanBuilder
|
||||
) {}
|
||||
|
||||
protected runNativeStream(
|
||||
input: NativeActionRuntimeInput,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
return runNativeActionRecipePreparedStream(input, signal);
|
||||
}
|
||||
|
||||
private async prepareNativeInput(
|
||||
input: ActionRuntimeBridgeInput
|
||||
): Promise<ActionRuntimeBridgeNativeInput & { input: unknown }> {
|
||||
const nativeInput = {
|
||||
...input.nativeInput,
|
||||
input: input.nativeInput?.input ?? {},
|
||||
};
|
||||
const structured = input.prepareStructuredRoutes;
|
||||
const image = input.prepareImageRoutes;
|
||||
if (!structured && !image) {
|
||||
return nativeInput;
|
||||
}
|
||||
if (!this.plans) {
|
||||
throw new Error('Action route preparation is not available');
|
||||
}
|
||||
const state =
|
||||
nativeInput.input && typeof nativeInput.input === 'object'
|
||||
? { ...(nativeInput.input as Record<string, unknown>) }
|
||||
: {};
|
||||
|
||||
if (structured) {
|
||||
const responseContract =
|
||||
structured.responseContract ??
|
||||
(buildStructuredResponseFromSchemaJson(
|
||||
structured.responseSchemaJson ?? { type: 'object' }
|
||||
) as RequiredStructuredOutputContract);
|
||||
const plan = await this.plans.buildStructuredPlan(
|
||||
{ modelId: structured.modelId },
|
||||
structured.messages,
|
||||
structured.options,
|
||||
structured.prefer ? { prefer: structured.prefer } : undefined,
|
||||
responseContract
|
||||
);
|
||||
const preparedRoutes = plan.nativeDispatch?.structured?.routes;
|
||||
if (!preparedRoutes?.length) {
|
||||
throw new Error('No native structured provider route prepared');
|
||||
}
|
||||
|
||||
const existingPreparedRoutes =
|
||||
state.preparedRoutes &&
|
||||
typeof state.preparedRoutes === 'object' &&
|
||||
!Array.isArray(state.preparedRoutes)
|
||||
? (state.preparedRoutes as Record<string, unknown>)
|
||||
: {};
|
||||
state.preparedRoutes = {
|
||||
...existingPreparedRoutes,
|
||||
[structured.stepId ?? 'generate']: preparedRoutes,
|
||||
};
|
||||
}
|
||||
|
||||
if (image) {
|
||||
const plan = await this.plans.buildImagePlan(
|
||||
{ modelId: image.modelId },
|
||||
image.messages,
|
||||
image.options,
|
||||
image.prefer ? { prefer: image.prefer } : undefined
|
||||
);
|
||||
const preparedRoutes = plan.nativeDispatch?.image?.routes;
|
||||
if (!preparedRoutes?.length) {
|
||||
throw new Error('No native image provider route prepared');
|
||||
}
|
||||
|
||||
const existingPreparedRoutes =
|
||||
state.preparedRoutes &&
|
||||
typeof state.preparedRoutes === 'object' &&
|
||||
!Array.isArray(state.preparedRoutes)
|
||||
? (state.preparedRoutes as Record<string, unknown>)
|
||||
: {};
|
||||
state.preparedRoutes = {
|
||||
...existingPreparedRoutes,
|
||||
[image.stepId ?? 'generate-image']: preparedRoutes,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
...nativeInput,
|
||||
input: state,
|
||||
};
|
||||
}
|
||||
|
||||
private async projectAssistantResult(
|
||||
input: ActionRuntimeBridgeInput,
|
||||
result: unknown,
|
||||
artifacts: unknown[],
|
||||
wasAborted: boolean
|
||||
) {
|
||||
if (!input.session) return null;
|
||||
const turn = projectActionResultToAssistantTurn({
|
||||
session: input.session,
|
||||
actionId: input.actionId,
|
||||
result,
|
||||
artifacts,
|
||||
wasAborted,
|
||||
});
|
||||
if (!turn) return null;
|
||||
return await this.turnPersistence.persistProjectedResult(
|
||||
input.session,
|
||||
turn,
|
||||
wasAborted
|
||||
);
|
||||
}
|
||||
|
||||
private async resolveAttempt(input: ActionRuntimeBridgeInput) {
|
||||
if (!input.retryOf) {
|
||||
return input.attempt ?? 1;
|
||||
}
|
||||
|
||||
const previous = await this.models.copilotActionRun.get(input.retryOf);
|
||||
if (!previous) {
|
||||
throw new Error('Retry source action run not found');
|
||||
}
|
||||
if (
|
||||
previous.userId !== input.userId ||
|
||||
previous.workspaceId !== input.workspaceId ||
|
||||
previous.actionId !== input.actionId ||
|
||||
previous.actionVersion !== input.actionVersion ||
|
||||
previous.sessionId !== (input.session?.config.sessionId ?? null)
|
||||
) {
|
||||
throw new Error('Retry source action run does not match current action');
|
||||
}
|
||||
if (input.attempt && input.attempt <= previous.attempt) {
|
||||
throw new Error('Retry attempt must be greater than source action run');
|
||||
}
|
||||
if (input.attempt) {
|
||||
return input.attempt;
|
||||
}
|
||||
return (previous?.attempt ?? 1) + 1;
|
||||
}
|
||||
|
||||
async *runStream(
|
||||
input: ActionRuntimeBridgeInput
|
||||
): AsyncIterableIterator<ActionRuntimeBridgeEvent> {
|
||||
const attempt = await this.resolveAttempt(input);
|
||||
const run = await this.models.copilotActionRun.create({
|
||||
userId: input.userId,
|
||||
workspaceId: input.workspaceId,
|
||||
docId: input.docId,
|
||||
sessionId: input.session?.config.sessionId,
|
||||
userMessageId: input.userMessageId,
|
||||
compatSubmissionId: input.compatSubmissionId,
|
||||
actionId: input.actionId,
|
||||
actionVersion: input.actionVersion,
|
||||
attempt,
|
||||
retryOf: input.retryOf,
|
||||
inputSnapshot: input.inputSnapshot,
|
||||
});
|
||||
await this.models.copilotActionRun.markRunning(run.id);
|
||||
await input.onRunCreated?.({
|
||||
runId: run.id,
|
||||
attempt,
|
||||
});
|
||||
|
||||
let finalEvent: NativeActionEvent | undefined;
|
||||
const attachments: unknown[] = [];
|
||||
try {
|
||||
const nativeInput = await this.prepareNativeInput({
|
||||
...input,
|
||||
});
|
||||
for await (const event of this.runNativeStream(
|
||||
{
|
||||
...nativeInput,
|
||||
recipeId: input.actionId,
|
||||
recipeVersion: input.actionVersion,
|
||||
},
|
||||
input.signal
|
||||
)) {
|
||||
finalEvent = event;
|
||||
let projectedEvent = event;
|
||||
if (event.type === 'attachment') {
|
||||
const attachment = input.persistAttachment
|
||||
? await input.persistAttachment(event.attachment)
|
||||
: event.attachment;
|
||||
attachments.push(attachment);
|
||||
projectedEvent = { ...event, attachment };
|
||||
}
|
||||
yield { ...projectedEvent, runId: run.id };
|
||||
}
|
||||
} catch (error) {
|
||||
finalEvent = {
|
||||
type: 'error',
|
||||
actionId: input.actionId,
|
||||
actionVersion: input.actionVersion,
|
||||
status: input.signal?.aborted ? 'aborted' : 'failed',
|
||||
errorCode: input.signal?.aborted
|
||||
? 'action_aborted'
|
||||
: 'action_bridge_stream_error',
|
||||
errorMessage:
|
||||
error instanceof Error ? error.message : 'action stream failed',
|
||||
};
|
||||
yield { ...finalEvent, runId: run.id };
|
||||
} finally {
|
||||
let status = resolveFinalStatus(finalEvent, input.signal);
|
||||
const result = finalEvent?.result;
|
||||
const artifacts =
|
||||
status === 'succeeded'
|
||||
? [...attachments, ...extractResultArtifacts(result)]
|
||||
: undefined;
|
||||
let assistantMessageId: string | null = null;
|
||||
let errorCode = status === 'succeeded' ? null : finalEvent?.errorCode;
|
||||
if (status === 'succeeded' || status === 'aborted') {
|
||||
try {
|
||||
assistantMessageId =
|
||||
(await this.projectAssistantResult(
|
||||
input,
|
||||
result,
|
||||
artifacts ?? [],
|
||||
status === 'aborted'
|
||||
)) ?? null;
|
||||
} catch {
|
||||
status = 'failed';
|
||||
errorCode = 'action_output_projection_failed';
|
||||
}
|
||||
}
|
||||
|
||||
await this.models.copilotActionRun.complete(run.id, {
|
||||
status,
|
||||
result: status === 'succeeded' ? result : undefined,
|
||||
artifacts: status === 'succeeded' ? artifacts : undefined,
|
||||
resultSummary:
|
||||
status === 'succeeded' ? summarizeActionResult(result) : null,
|
||||
errorCode,
|
||||
trace: finalEvent?.trace ?? undefined,
|
||||
assistantMessageId,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user