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
2026-05-04 00:36:47 +08:00
committed by GitHub
parent fa8f1a096c
commit d64f368623
239 changed files with 35859 additions and 16777 deletions
@@ -1,4 +1,4 @@
# Snapshot report for `src/__tests__/copilot.spec.ts`
# Snapshot report for `src/__tests__/copilot/copilot.spec.ts`
The actual snapshot is saved in `copilot.spec.ts.snap`.
@@ -52,12 +52,10 @@ Generated by [AVA](https://avajs.dev).
},
{
content: 'hello',
params: {},
role: 'user',
},
{
content: 'world',
params: {},
role: 'assistant',
},
]
@@ -74,12 +72,10 @@ Generated by [AVA](https://avajs.dev).
},
{
content: 'hello',
params: {},
role: 'user',
},
{
content: 'world',
params: {},
role: 'assistant',
},
]
@@ -96,22 +92,18 @@ Generated by [AVA](https://avajs.dev).
},
{
content: 'hello',
params: {},
role: 'user',
},
{
content: 'world',
params: {},
role: 'assistant',
},
{
content: 'aaa',
params: {},
role: 'user',
},
{
content: 'bbb',
params: {},
role: 'assistant',
},
]
@@ -128,22 +120,18 @@ Generated by [AVA](https://avajs.dev).
},
{
content: 'hello',
params: {},
role: 'user',
},
{
content: 'world',
params: {},
role: 'assistant',
},
{
content: 'aaa',
params: {},
role: 'user',
},
{
content: 'bbb',
params: {},
role: 'assistant',
},
]
@@ -445,6 +433,40 @@ Generated by [AVA](https://avajs.dev).
],
}
## capability policy host should gate pro model requests by subscription status
> should honor requested pro model
'gemini-2.5-pro'
> should fallback to default model
'gemini-2.5-flash'
> should fallback to default model when requesting pro model during trialing
'gemini-2.5-flash'
> should honor requested non-pro model during trialing
'gemini-2.5-flash'
> should pick default model when no requested model during trialing
'gemini-2.5-flash'
> should pick default model when no requested model during active
'gemini-2.5-flash'
> should honor requested pro model during active
'claude-sonnet-4-5@20250929'
> should fallback to default model when requesting non-optional model during active
'gemini-2.5-flash'
## should resolve model correctly based on subscription status and prompt config
> should honor requested pro model
@@ -0,0 +1,703 @@
# Snapshot report for `src/__tests__/copilot/native-provider.spec.ts`
The actual snapshot is saved in `native-provider.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## NativeProviderAdapter streamObject should map tool and text events
> Snapshot 1
[
{
args: {
doc_id: 'a1',
},
argumentParseError: undefined,
rawArgumentsText: undefined,
thought: undefined,
toolCallId: 'call_1',
toolName: 'doc_read',
type: 'tool-call',
},
{
args: {
doc_id: 'a1',
},
argumentParseError: undefined,
rawArgumentsText: undefined,
result: {
markdown: '# a1',
},
toolCallId: 'call_1',
toolName: 'doc_read',
type: 'tool-result',
},
{
textDelta: 'ok',
type: 'text-delta',
},
]
## buildCanonicalNativeRequest should only use explicit structured contract inputs
> Snapshot 1
{
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
}
## buildCanonicalNativeStructuredRequest should accept schema-only explicit structured response contracts
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
}
## buildCanonicalNativeStructuredRequest should honor explicit structured options contract before system responseFormat
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
ok: {
type: 'boolean',
},
},
required: [
'ok',
],
type: 'object',
},
strict: true,
}
## buildCanonicalNativeStructuredRequest should honor explicit responseSchema for array outputs
> Snapshot 1
{
items: {
additionalProperties: false,
properties: {
speaker: {
type: 'string',
},
text: {
type: 'string',
},
},
required: [
'speaker',
'text',
],
type: 'object',
},
type: 'array',
}
## buildCanonicalNativeStructuredRequest should consume explicit structured response contract without options.schema
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: false,
}
## buildCanonicalNativeStructuredRequest should accept explicit schema contracts without schemaHash
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
}
## buildNativeRequest should canonicalize Gemini attachments
> remote file url
[
{
text: 'summarize this attachment',
type: 'text',
},
{
source: {
media_type: 'application/pdf',
url: 'https://example.com/a.pdf',
},
type: 'file',
},
]
> remote image url
[
{
text: 'describe this image',
type: 'text',
},
{
source: {
media_type: 'image/png',
url: 'https://example.com/cat.png',
},
type: 'image',
},
]
> data url
[
{
text: 'read this note',
type: 'text',
},
{
source: {
data: 'aGVsbG8gd29ybGQ=',
media_type: 'text/plain',
},
type: 'file',
},
]
> remote audio url
[
{
text: 'transcribe this clip',
type: 'text',
},
{
source: {
media_type: 'audio/mpeg',
url: 'https://example.com/a.mp3',
},
type: 'audio',
},
]
> bytes and file handle
[
{
text: 'inspect these assets',
type: 'text',
},
{
source: {
data: 'aGVsbG8=',
file_name: 'hello.txt',
media_type: 'text/plain',
},
type: 'file',
},
{
source: {
file_handle: 'file_123',
file_name: 'report.pdf',
media_type: 'application/pdf',
},
type: 'file',
},
]
## buildNativeStructuredRequest should prefer explicit schema option
> Snapshot 1
{
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
}
## buildNativeStructuredRequest should ignore legacy params.schema fallback when explicit schema contract exists
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
}
## defineTool should precompute json schema at definition time
> Snapshot 1
{
additionalProperties: false,
properties: {
docId: {
type: 'string',
},
includeChildren: {
type: 'boolean',
},
},
required: [
'docId',
],
type: 'object',
}
## GeminiProvider should use native path for text-only requests
> Snapshot 1
{
include: [
'reasoning',
],
middleware: {
request: [
'normalize_messages',
'tool_schema_rewrite',
],
stream: [
'stream_event_normalize',
'citation_indexing',
],
},
reasoning: {
effort: 'medium',
},
remoteAttachmentRequests: [],
}
## GeminiProvider should use native path for structured requests
> Snapshot 1
{
request: {
messages: [
{
content: [
{
text: 'Return JSON only.',
type: 'text',
},
],
role: 'system',
},
{
content: [
{
text: 'Summarize AFFiNE in one short sentence.',
type: 'text',
},
],
role: 'user',
},
],
middleware: {
request: [
'normalize_messages',
'tool_schema_rewrite',
],
},
model: 'gemini-2.5-flash',
responseMimeType: 'application/json',
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
},
result: {
summary: 'AFFiNE native',
},
}
## GeminiProvider should use native structured path for audio attachments
> Snapshot 1
{
content: [
{
text: 'transcribe the audio',
type: 'text',
},
{
source: {
data: 'YXVkaW8tYnl0ZXM=',
media_type: 'audio/mpeg',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'https://example.com/a.mp3',
],
result: [
{
a: 'Speaker 1',
e: 1,
s: 0,
t: 'Hello',
},
],
}
## GeminiProvider should use native path for embeddings
> Snapshot 1
{
request: {
dimensions: 3,
inputs: [
'first',
'second',
],
model: 'gemini-embedding-001',
taskType: 'RETRIEVAL_DOCUMENT',
},
result: [
[
0.1,
0.2,
],
[
1.1,
1.2,
],
],
}
## GeminiProvider should canonicalize native text attachments
> remote file attachment
{
content: [
{
text: 'summarize this file',
type: 'text',
},
{
source: {
data: 'cGRmLWJ5dGVz',
media_type: 'application/pdf',
},
type: 'file',
},
],
remoteAttachmentRequests: [
'https://example.com/a.pdf',
],
}
> remote image attachment
{
content: [
{
text: 'describe this image',
type: 'text',
},
{
source: {
data: 'aW1hZ2UtYnl0ZXM=',
media_type: 'image/jpeg',
},
type: 'image',
},
],
remoteAttachmentRequests: [
'https://example.com/a.jpg',
],
}
> downloaded audio webm attachment
{
content: [
{
text: 'transcribe this clip',
type: 'text',
},
{
source: {
data: 'YXVkaW8tYnl0ZXM=',
media_type: 'audio/webm',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'https://example.com/a.webm',
],
}
> google file url attachment
{
content: [
{
text: 'summarize this file',
type: 'text',
},
{
source: {
media_type: 'application/pdf',
url: 'https://generativelanguage.googleapis.com/v1beta/files/file-123',
},
type: 'file',
},
],
remoteAttachmentRequests: [],
}
## PerplexityProvider should ignore attachments during text model matching
> Snapshot 1
[
{
text: 'summarize this',
type: 'text',
},
]
## GeminiVertexProvider should prefetch bearer token for native config
> Snapshot 1
{
auth_token: 'vertex-token',
base_url: 'https://vertex.example',
}
## GeminiVertexProvider should materialize remote attachments before native text path
> remote http url
{
content: [
{
text: 'transcribe the audio',
type: 'text',
},
{
source: {
data: 'YXVkaW8tYnl0ZXM=',
media_type: 'audio/mpeg',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'https://example.com/a.mp3',
],
}
> gs url
{
content: [
{
text: 'transcribe the audio',
type: 'text',
},
{
source: {
data: 'b3B1cy1ieXRlcw==',
media_type: 'audio/opus',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'gs://bucket/audio.opus',
],
}
## OpenAIProvider should use native structured dispatch
> Snapshot 1
{
request: {
messages: [
{
content: [
{
text: 'Return JSON only.',
type: 'text',
},
],
role: 'system',
},
{
content: [
{
text: 'Summarize AFFiNE in one sentence.',
type: 'text',
},
],
role: 'user',
},
],
middleware: {
request: [
'normalize_messages',
'tool_schema_rewrite',
],
},
model: 'gpt-4.1',
responseMimeType: 'application/json',
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
},
result: {
summary: 'AFFiNE structured',
},
}
## OpenAIProvider should prefer native output_json for structured dispatch
> Snapshot 1
{
summary: 'AFFiNE structured',
}
## OpenAIProvider should use native embedding dispatch
> Snapshot 1
{
request: {
dimensions: 8,
inputs: [
'alpha',
'beta',
],
model: 'text-embedding-3-small',
taskType: 'RETRIEVAL_DOCUMENT',
},
result: [
[
0.4,
0.5,
],
[
0.4,
0.5,
],
],
}
## OpenAIProvider should use native rerank dispatch
> Snapshot 1
{
request: {
candidates: [
{
id: 'react',
text: 'React is a UI library.',
},
{
id: 'weather',
text: 'The park is sunny today.',
},
],
model: 'gpt-4.1',
query: 'programming',
},
scores: [
0.8,
0.8,
],
}
@@ -0,0 +1,505 @@
# Snapshot report for `src/__tests__/copilot/provider-native.spec.ts`
The actual snapshot is saved in `provider-native.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## CopilotProviderFactory should return no prepared routes when native prepare returns null
> Snapshot 1
{
chat: [
length: 0,
prepared: undefined,
providerId: undefined,
],
embedding: [
length: 0,
prepared: undefined,
],
rerank: [
length: 0,
prepared: undefined,
],
structured: [
length: 0,
prepared: undefined,
],
}
## getActiveProviderMiddleware should merge defaults with profile override
> Snapshot 1
{
node: {
text: [
'citation_footnote',
'callout',
'thinking_format',
],
},
rust: {
request: [
'clamp_max_tokens',
],
stream: undefined,
},
}
## checkParams should infer remote image capability from url extension without host mime inference
> Snapshot 1
{
attachmentKinds: [
'image',
],
attachmentSourceKinds: [
'url',
],
inputTypes: [
'image',
'text',
],
}
## llmResolveRequestedModelMatch should preserve provider-prefixed optional matches
> prefixed optional hit
{
matchedOptionalModel: true,
selectedModel: 'openai-default/gemini-2.5-pro',
}
> prefixed optional miss
{
matchedOptionalModel: false,
selectedModel: 'gemini-2.5-flash',
}
## ExecutionPlan should serialize routed request state and reject host-only signal
> Snapshot 1
{
fallbackOrder: [
'openai-main',
],
transport: {
kind: 'chat',
request: {
messages: [
{
content: [
{
text: 'hello',
type: 'text',
},
],
role: 'user',
},
],
model: 'gpt-5-mini',
},
},
}
## NativeExecutionEngine should dispatch prepared text routes through native fallback
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello from primary',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello from fallback',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
]
## NativeExecutionEngine should prefer prepared native fallback dispatch for explicit routes
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
]
## ExecutionPlanBuilder should keep tool-loop chat routes on prepared dispatch path
> Snapshot 1
{
preparedTools: [
'answer',
],
transport: undefined,
}
## ExecutionPlanBuilder should keep single-route tool chat plans on prepared_routes path
> Snapshot 1
{
kind: 'chat',
request: {
messages: [
{
content: [
{
text: 'hello',
type: 'text',
},
],
role: 'user',
},
],
model: 'gpt-5-mini',
tools: [
{
description: 'Answer',
name: 'answer',
parameters: {
properties: {
value: {
type: 'string',
},
},
required: [
'value',
],
type: 'object',
},
},
],
},
}
## NativeExecutionEngine should route tool-loop chat prepared routes through native dispatch
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
'tools',
],
query: undefined,
schemaKeys: undefined,
toolNames: [
'answer',
],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello from fallback',
inputCount: 0,
keys: [
'messages',
'model',
'tools',
],
query: undefined,
schemaKeys: undefined,
toolNames: [
'answer',
],
},
},
]
## ExecutionPlanBuilder should build native prepared routes for structured, image, embedding and rerank
> Snapshot 1
{
embedding: {
routes: 2,
transport: undefined,
},
image: {
prepared: {
request: {
images: [],
model: 'gpt-image-1',
operation: 'generate',
prompt: 'draw a cat',
},
route: {
backendConfig: {
auth_token: 'image-key',
base_url: 'https://api.openai.com',
},
model: 'gpt-image-1',
protocol: 'openai_images',
providerId: 'openai-default',
},
},
routes: [
{
config: {
auth_token: 'image-key',
base_url: 'https://api.openai.com',
},
model: 'gpt-image-1',
protocol: 'openai_images',
provider_id: 'openai-default',
request: {
images: [],
model: 'gpt-image-1',
operation: 'generate',
prompt: 'draw a cat',
},
},
],
},
rerank: {
routes: 2,
transport: undefined,
},
structured: {
routes: 2,
transport: undefined,
},
}
## NativeExecutionEngine should dispatch structured prepared routes through native execution
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
'schema',
],
query: undefined,
schemaKeys: [
'ok',
],
toolNames: [],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello from fallback',
inputCount: 0,
keys: [
'messages',
'model',
'schema',
],
query: undefined,
schemaKeys: [
'ok',
],
toolNames: [],
},
},
]
## NativeExecutionEngine should dispatch embedding prepared routes through native execution
> Snapshot 1
{
called: true,
result: [
[
0.1,
0.2,
],
],
routes: [
{
model: 'text-embedding-3-small',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: null,
inputCount: 1,
keys: [
'inputs',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'text-embedding-3-small',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: null,
inputCount: 1,
keys: [
'inputs',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
],
}
## NativeExecutionEngine should dispatch rerank prepared routes through native execution
> Snapshot 1
{
called: true,
result: [
0.9,
0.1,
],
routes: [
{
model: 'gpt-4o-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 1,
firstContent: null,
inputCount: 0,
keys: [
'candidates',
'model',
'query',
],
query: 'programming',
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'gpt-4o-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 1,
firstContent: null,
inputCount: 0,
keys: [
'candidates',
'model',
'query',
],
query: 'programming fallback',
schemaKeys: undefined,
toolNames: [],
},
},
],
}
## NativeExecutionEngine should dispatch image plans through prepared native routes
> Snapshot 1
[
{
model: 'gpt-image-1',
providerId: 'openai-image',
requestShape: {
candidateCount: 0,
firstContent: null,
imageCount: 0,
inputCount: 0,
keys: [
'images',
'model',
'operation',
'prompt',
],
prompt: 'draw a cat',
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
]
@@ -1,50 +1,51 @@
import { randomUUID } from 'node:crypto';
import type { Prisma } from '@prisma/client';
import type { ExecutionContext, TestFn } from 'ava';
import ava from 'ava';
import Sinon from 'sinon';
import { z } from 'zod';
import { ServerFeature, ServerService } from '../../core';
import { AuthService } from '../../core/auth';
import { QuotaModule } from '../../core/quota';
import { Models } from '../../models';
import { llmImageDispatchPlan } from '../../native';
import { CopilotModule } from '../../plugins/copilot';
import { prompts, PromptService } from '../../plugins/copilot/prompt';
import { PromptService } from '../../plugins/copilot/prompt';
import {
CopilotProviderFactory,
CopilotProviderType,
StreamObject,
StreamObjectSchema,
} from '../../plugins/copilot/providers';
import { TranscriptionResponseSchema } from '../../plugins/copilot/transcript/schema';
import {
CopilotChatTextExecutor,
CopilotWorkflowService,
GraphExecutorState,
} from '../../plugins/copilot/workflow';
import {
CopilotChatImageExecutor,
CopilotCheckHtmlExecutor,
CopilotCheckJsonExecutor,
} from '../../plugins/copilot/workflow/executor';
import { ActionStreamHost } from '../../plugins/copilot/runtime/hosts/action-stream-host';
import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context';
import { ChatSession, ChatSessionService } from '../../plugins/copilot/session';
import { TranscriptPayloadSchema } from '../../plugins/copilot/transcript/schema';
import { CopilotTranscriptionService } from '../../plugins/copilot/transcript/service';
import { TestingPromptService } from '../mocks/prompt-service.mock';
import { createTestingModule, TestingModule } from '../utils';
import { TestAssets } from '../utils/copilot';
import {
assistantPrompt,
promptMessages,
singleUserPromptMessages,
userPrompt,
} from './prompt-test-helper';
type Tester = {
auth: AuthService;
module: TestingModule;
models: Models;
service: ServerService;
prompt: PromptService;
prompt: TestingPromptService;
factory: CopilotProviderFactory;
workflow: CopilotWorkflowService;
executors: {
image: CopilotChatImageExecutor;
text: CopilotChatTextExecutor;
html: CopilotCheckHtmlExecutor;
json: CopilotCheckJsonExecutor;
};
session: ChatSessionService;
actionStreams: ActionStreamHost;
transcript: CopilotTranscriptionService;
};
const test = ava as TestFn<Tester>;
let isCopilotConfigured = false;
@@ -65,6 +66,9 @@ const runIfCopilotConfigured = test.macro(
test.serial.before(async t => {
const module = await createTestingModule({
imports: [QuotaModule, CopilotModule],
tapModule: builder => {
builder.overrideProvider(PromptService).useClass(TestingPromptService);
},
});
const service = module.get(ServerService);
@@ -72,9 +76,11 @@ test.serial.before(async t => {
const auth = module.get(AuthService);
const models = module.get(Models);
const prompt = module.get(PromptService);
const prompt = module.get(PromptService) as TestingPromptService;
const factory = module.get(CopilotProviderFactory);
const workflow = module.get(CopilotWorkflowService);
const session = module.get(ChatSessionService);
const actionStreams = module.get(ActionStreamHost);
const transcript = module.get(CopilotTranscriptionService);
t.context.module = module;
t.context.auth = auth;
@@ -82,51 +88,15 @@ test.serial.before(async t => {
t.context.models = models;
t.context.prompt = prompt;
t.context.factory = factory;
t.context.workflow = workflow;
t.context.executors = {
image: module.get(CopilotChatImageExecutor),
text: module.get(CopilotChatTextExecutor),
html: module.get(CopilotCheckHtmlExecutor),
json: module.get(CopilotCheckJsonExecutor),
};
t.context.session = session;
t.context.actionStreams = actionStreams;
t.context.transcript = transcript;
});
test.serial.before(async t => {
const { prompt, executors, models, service } = t.context;
const { prompt } = t.context;
executors.image.register();
executors.text.register();
executors.html.register();
executors.json.register();
for (const name of await prompt.listNames()) {
await prompt.delete(name);
}
for (const p of prompts) {
await prompt.set(p.name, p.model, p.messages, p.config);
}
const user = await models.user.create({
email: `${randomUUID()}@affine.pro`,
});
await service.updateConfig(user.id, [
{
module: 'copilot',
key: 'scenarios',
value: {
enabled: true,
scenarios: {
image: 'flux-1/schnell',
complex_text_generation: 'gpt-5-mini',
coding: 'gpt-5-mini',
quick_decision_making: 'gpt-5-mini',
quick_text_generation: 'gpt-5-mini',
polish_and_summarize: 'gemini-2.5-flash',
},
},
},
]);
prompt.reset();
});
test.after(async t => {
@@ -332,10 +302,9 @@ const actions = [
{
name: 'Should chat with histories',
promptName: ['Chat With AFFiNE AI'],
messages: [
{
role: 'user' as const,
content: `
messages: promptMessages(
userPrompt(
`
Hi! Im going to send you a technical term related to real-time collaborative editing (e.g., CRDT, Operational Transformation, OT Composer, etc.). Whenever I send you a term:
1. Translate it into Chinese (send me the Chinese version).
2. Then translate that Chinese back into English (send me the retranslated English).
@@ -344,11 +313,10 @@ Hi! Im going to send you a technical term related to real-time collaborative
5. Finally, give the origin or “term history” (e.g., who introduced it, in which paper or year).
If you understand, please proceed by explaining the term “CRDT.”
`.trim(),
},
{
role: 'assistant' as const,
content: `
`.trim()
),
assistantPrompt(
`
1. **Chinese Translation:**
“CRDT” → **无冲突复制数据类型**
@@ -366,13 +334,12 @@ CRDTs enable **eventual consistency (最终一致性)** in real-time collaborati
4. **Origin / Term History:**
The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Carlos Baquero, and Marek Zawirski in their 2011 paper titled “Conflict-free Replicated Data Types” (published in the _Stabilization, Safety, and Security of Distributed Systems (SSS)_ conference). They formalized two families of CRDTs—state-based (“Convergent Replicated Data Types” or CvRDTs) and operation-based (“Commutative Replicated Data Types” or CmRDTs)—and proved their convergence properties under asynchronous, unreliable networks.
`.trim(),
},
{
role: 'user' as const,
content: `Thanks! Now please just tell me the **Chinese translation** and the **back-translated English term** that you provided previously for “CRDT.” Do not reprint the full introduction—only those two lines.`,
},
],
`.trim()
),
userPrompt(
'Thanks! Now please just tell me the **Chinese translation** and the **back-translated English term** that you provided previously for “CRDT.” Do not reprint the full introduction—only those two lines.'
)
),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
const lower = result.toLowerCase();
@@ -387,22 +354,18 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
{
name: 'Should not have citation',
promptName: ['Chat With AFFiNE AI'],
messages: [
{
role: 'user' as const,
content: 'what is AFFiNE AI?',
params: {
files: [
{
blobId: 'todo_md',
fileName: 'todo.md',
fileType: 'text/markdown',
fileContent: TestAssets.TODO,
},
],
},
messages: singleUserPromptMessages('what is AFFiNE AI?', {
params: {
files: [
{
blobId: 'todo_md',
fileName: 'todo.md',
fileType: 'text/markdown',
fileContent: TestAssets.TODO,
},
],
},
],
}),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
assertCitation(t, result, (t, c) => {
@@ -422,22 +385,18 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
{
name: 'Should have citation',
promptName: ['Chat With AFFiNE AI'],
messages: [
{
role: 'user' as const,
content: 'what is ssot',
params: {
docs: [
{
docId: 'SSOT',
docTitle: 'Single source of truth - Wikipedia',
fileType: 'text/markdown',
docContent: TestAssets.SSOT,
},
],
},
messages: singleUserPromptMessages('what is ssot', {
params: {
docs: [
{
docId: 'SSOT',
docTitle: 'Single source of truth - Wikipedia',
fileType: 'text/markdown',
docContent: TestAssets.SSOT,
},
],
},
],
}),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
assertCitation(t, result);
@@ -447,12 +406,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
{
name: 'stream objects',
promptName: ['Chat With AFFiNE AI'],
messages: [
{
role: 'user' as const,
content: 'what is AFFiNE AI',
},
],
messages: singleUserPromptMessages('what is AFFiNE AI'),
verifier: (t: ExecutionContext<Tester>, result: string) => {
t.truthy(checkStreamObjects(result), 'should be valid stream objects');
},
@@ -461,13 +415,9 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
{
name: 'Gemini native text',
promptName: ['Chat With AFFiNE AI'],
messages: [
{
role: 'user' as const,
content:
'In one short sentence, explain what AFFiNE AI is and mention AFFiNE by name.',
},
],
messages: singleUserPromptMessages(
'In one short sentence, explain what AFFiNE AI is and mention AFFiNE by name.'
),
config: { model: 'gemini-2.5-flash' },
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
@@ -482,13 +432,9 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
{
name: 'Gemini native stream objects',
promptName: ['Chat With AFFiNE AI'],
messages: [
{
role: 'user' as const,
content:
'Respond with one short sentence about AFFiNE AI and mention AFFiNE by name.',
},
],
messages: singleUserPromptMessages(
'Respond with one short sentence about AFFiNE AI and mention AFFiNE by name.'
),
config: { model: 'gemini-2.5-flash' },
verifier: (t: ExecutionContext<Tester>, result: string) => {
t.truthy(checkStreamObjects(result), 'should be valid stream objects');
@@ -501,92 +447,18 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
prefer: CopilotProviderType.Gemini,
type: 'object' as const,
},
{
name: 'Should transcribe short audio',
promptName: ['Transcript audio'],
messages: [
{
role: 'user' as const,
content: 'transcript the audio',
attachments: [
'https://cdn.affine.pro/copilot-test/MP9qDGuYgnY+ILoEAmHpp3h9Npuw2403EAYMEA.mp3',
],
params: {
schema: TranscriptionResponseSchema,
},
},
],
verifier: (t: ExecutionContext<Tester>, result: string) => {
t.notThrows(() => {
TranscriptionResponseSchema.parse(JSON.parse(result));
});
},
type: 'structured' as const,
prefer: CopilotProviderType.Gemini,
},
{
name: 'Should transcribe middle audio',
promptName: ['Transcript audio'],
messages: [
{
role: 'user' as const,
content: 'transcript the audio',
attachments: [
'https://cdn.affine.pro/copilot-test/2ed05eo1KvZ2tWB_BAjFo67EAPZZY-w4LylUAw.m4a',
],
params: {
schema: TranscriptionResponseSchema,
},
},
],
verifier: (t: ExecutionContext<Tester>, result: string) => {
t.notThrows(() => {
TranscriptionResponseSchema.parse(JSON.parse(result));
});
},
type: 'structured' as const,
prefer: CopilotProviderType.Gemini,
},
{
name: 'Should transcribe long audio',
promptName: ['Transcript audio'],
messages: [
{
role: 'user' as const,
content: 'transcript the audio',
attachments: [
'https://cdn.affine.pro/copilot-test/nC9-e7P85PPI2rU29QWwf8slBNRMy92teLIIMw.opus',
],
params: {
schema: TranscriptionResponseSchema,
},
},
],
config: { model: 'gemini-2.5-pro' },
verifier: (t: ExecutionContext<Tester>, result: string) => {
t.notThrows(() => {
TranscriptionResponseSchema.parse(JSON.parse(result));
});
},
type: 'structured' as const,
prefer: CopilotProviderType.Gemini,
},
{
promptName: ['Conversation Summary'],
messages: [
{
role: 'user' as const,
content: '',
params: {
messages: [
{ role: 'user', content: 'what is single source of truth?' },
{ role: 'assistant', content: TestAssets.SSOT },
],
focus: 'technical decisions',
length: 'comprehensive',
},
messages: singleUserPromptMessages('', {
params: {
messages: [
userPrompt('what is single source of truth?'),
assistantPrompt(TestAssets.SSOT),
],
focus: 'technical decisions',
length: 'comprehensive',
},
],
}),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
const cleared = result.toLowerCase();
@@ -619,7 +491,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
'Section Edit',
'Chat With AFFiNE AI',
],
messages: [{ role: 'user' as const, content: TestAssets.SSOT }],
messages: singleUserPromptMessages(TestAssets.SSOT),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
const cleared = result.toLowerCase();
@@ -634,7 +506,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: ['Continue writing'],
messages: [{ role: 'user' as const, content: TestAssets.AFFiNE }],
messages: singleUserPromptMessages(TestAssets.AFFiNE),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(result.length > 0, 'should not be empty');
@@ -643,7 +515,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: ['Brainstorm ideas about this', 'Brainstorm mindmap'],
messages: [{ role: 'user' as const, content: TestAssets.AFFiNE }],
messages: singleUserPromptMessages(TestAssets.AFFiNE),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(checkMDList(result), 'should be a markdown list');
@@ -652,7 +524,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: 'Expand mind map',
messages: [{ role: 'user' as const, content: '- Single source of truth' }],
messages: singleUserPromptMessages('- Single source of truth'),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(checkMDList(result), 'should be a markdown list');
@@ -661,7 +533,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: 'Find action items from it',
messages: [{ role: 'user' as const, content: TestAssets.TODO }],
messages: singleUserPromptMessages(TestAssets.TODO),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(checkMDList(result), 'should be a markdown list');
@@ -670,7 +542,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: ['Explain this code', 'Check code error'],
messages: [{ role: 'user' as const, content: TestAssets.Code }],
messages: singleUserPromptMessages(TestAssets.Code),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(
@@ -683,13 +555,9 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: 'Translate to',
messages: [
{
role: 'user' as const,
content: TestAssets.SSOT,
params: { language: 'Simplified Chinese' },
},
],
messages: singleUserPromptMessages(TestAssets.SSOT, {
params: { language: 'Simplified Chinese' },
}),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
const cleared = result.toLowerCase();
@@ -702,15 +570,11 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: ['Generate a caption', 'Explain this image'],
messages: [
{
role: 'user' as const,
content: '',
attachments: [
'https://cdn.affine.pro/copilot-test/Qgqy9qZT3VGIEuMIotJYoCCH.jpg',
],
},
],
messages: singleUserPromptMessages('', {
attachments: [
'https://cdn.affine.pro/copilot-test/Qgqy9qZT3VGIEuMIotJYoCCH.jpg',
],
}),
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
const content = result.toLowerCase();
@@ -725,15 +589,11 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: ['Convert to sticker', 'Remove background', 'Upscale image'],
messages: [
{
role: 'user' as const,
content: '',
attachments: [
'https://cdn.affine.pro/copilot-test/Zkas098lkjdf-908231.jpg',
],
},
],
messages: singleUserPromptMessages('', {
attachments: [
'https://cdn.affine.pro/copilot-test/Zkas098lkjdf-908231.jpg',
],
}),
verifier: (t: ExecutionContext<Tester>, link: string) => {
t.truthy(checkUrl(link), 'should be a valid url');
},
@@ -741,12 +601,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
{
promptName: ['Generate image'],
messages: [
{
role: 'user' as const,
content: 'Panda',
},
],
messages: singleUserPromptMessages('Panda'),
config: { quality: 'low' },
verifier: (t: ExecutionContext<Tester>, link: string) => {
t.truthy(checkUrl(link), 'should be a valid url');
@@ -774,7 +629,9 @@ for (const {
const prompt = (await promptService.get(promptName))!;
t.truthy(prompt, 'should have prompt');
const finalConfig = Object.assign({}, prompt.config, config);
const modelId = finalConfig.model || prompt.model;
const modelId =
('model' in finalConfig ? finalConfig.model : undefined) ??
prompt.model;
const provider = (await factory.getProviderByModel(modelId, {
prefer,
}))!;
@@ -782,29 +639,11 @@ for (const {
await retry(`action: ${promptName}`, t, async t => {
switch (type) {
case 'text': {
const result = await provider.text(
const result = await getProviderRuntimeHost(provider).run.text(
{ modelId },
[
...prompt.finish(
messages.reduce(
// @ts-expect-error params not typed
(acc, m) => Object.assign(acc, m.params),
{}
)
),
...messages,
],
finalConfig
);
t.truthy(result, 'should return result');
verifier?.(t, result);
break;
}
case 'structured': {
const result = await provider.structure(
{ modelId },
[
...prompt.finish(
...promptService.finish(
prompt,
messages.reduce(
(acc, m) => Object.assign(acc, m.params),
{}
@@ -820,10 +659,13 @@ for (const {
}
case 'object': {
const streamObjects: StreamObject[] = [];
for await (const chunk of provider.streamObject(
for await (const chunk of getProviderRuntimeHost(
provider
).run.streamObject(
{ modelId },
[
...prompt.finish(
...promptService.finish(
prompt,
messages.reduce(
(acc, m) => Object.assign(acc, (m as any).params || {}),
{}
@@ -852,29 +694,39 @@ for (const {
: undefined,
});
}
const stream = provider.streamImages(
{ modelId },
[
...prompt.finish(
finalMessage.reduce(
// @ts-expect-error params not typed
(acc, m) => Object.assign(acc, m.params),
params
)
),
...finalMessage,
const imageMessages = [
...promptService.finish(
prompt,
finalMessage.reduce(
(acc, m) => Object.assign(acc, m.params),
params
)
),
...finalMessage,
];
const prepared = await getProviderRuntimeHost(
provider
).prepare.image({ modelId }, imageMessages, finalConfig);
t.truthy(prepared, 'should prepare image request');
const result = await llmImageDispatchPlan({
preparedRoutes: [
{
provider_id: prepared!.route.providerId,
protocol: prepared!.route.protocol,
model: prepared!.route.model,
config: prepared!.route.backendConfig,
request: prepared!.request,
},
],
finalConfig
);
});
const result = [];
for await (const attachment of stream) {
result.push(attachment);
}
t.truthy(result.length, 'should return result');
for (const r of result) {
verifier?.(t, r);
t.truthy(result.response.images.length, 'should return result');
for (const image of result.response.images) {
const link = image.data_base64
? `data:${image.media_type};base64,${image.data_base64}`
: image.url;
t.truthy(link);
verifier?.(t, link!);
}
break;
}
@@ -889,53 +741,278 @@ for (const {
}
}
// ==================== workflow ====================
// ==================== action recipes ====================
const workflows = [
function actionRunRecord(
input: Parameters<Models['copilotActionRun']['create']>[0]
) {
return {
id: `action-run-${randomUUID()}`,
userId: input.userId,
workspaceId: input.workspaceId,
docId: input.docId ?? null,
sessionId: input.sessionId ?? null,
userMessageId: input.userMessageId ?? null,
compatSubmissionId: input.compatSubmissionId ?? null,
assistantMessageId: null,
actionId: input.actionId,
actionVersion: input.actionVersion,
status: 'created' as const,
attempt: input.attempt ?? 1,
retryOf: input.retryOf ?? null,
inputSnapshot: (input.inputSnapshot ?? null) as Prisma.JsonValue,
result: null,
artifacts: null,
resultSummary: null,
errorCode: null,
trace: null,
createdAt: new Date(),
updatedAt: new Date(),
};
}
function installActionSessionMock(
t: ExecutionContext<Tester>,
{
name: 'brainstorm',
actionId,
actionPrompt,
content,
}: {
actionId: string;
actionPrompt: Awaited<ReturnType<TestingPromptService['get']>>;
content: string;
}
) {
const { models, session } = t.context;
const sandbox = Sinon.createSandbox();
const sessionId = `copilot-provider-action-${actionId}-${randomUUID()}`;
const userId = `copilot-provider-user-${randomUUID()}`;
const workspaceId = `copilot-provider-action-${actionId}`;
const docId = `copilot-provider-action-${actionId}-doc`;
const savedTurns: Array<{ role: string }> = [];
const userTurn = {
conversationId: sessionId,
role: 'user' as const,
content,
attachments: [],
renderTrace: [],
toolEvents: [],
metadata: { language: 'English' },
createdAt: new Date(),
};
const chatSession = new ChatSession(
{
userId,
sessionId,
workspaceId,
docId,
turns: [userTurn],
prompt: actionPrompt!,
},
(prompt, turns, params, maxTokenSize, sessionId) =>
t.context.prompt.renderSession(
prompt,
turns,
params,
maxTokenSize,
sessionId
),
async state => {
savedTurns.push(...state.turns);
}
);
sandbox
.stub(session, 'get')
.callsFake(async id => (id === sessionId ? chatSession : null));
sandbox.stub(session, 'appendTurn').callsFake(async input => {
savedTurns.push(input.turn);
return { ...input.turn, id: `assistant-${randomUUID()}` };
});
sandbox.stub(session, 'revertLatestMessage').resolves();
sandbox
.stub(models.copilotActionRun, 'create')
.callsFake(async input => actionRunRecord(input));
sandbox.stub(models.copilotActionRun, 'markRunning').callsFake(
async id =>
({
id,
status: 'running',
}) as never
);
sandbox.stub(models.copilotActionRun, 'complete').callsFake(
async (id, input) =>
({
id,
...input,
updatedAt: new Date(),
}) as never
);
return { sandbox, sessionId, userId, savedTurns };
}
const actionRecipeCases = [
{
actionId: 'mindmap.generate',
content: 'apple company',
verifier: (t: ExecutionContext, result: string) => {
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(checkMDList(result), 'should be a markdown list');
},
},
{
name: 'presentation',
actionId: 'slides.outline',
content: 'apple company',
verifier: (t: ExecutionContext, result: string) => {
for (const l of result.split('\n')) {
const line = l.trim();
if (!line) continue;
t.notThrows(() => {
JSON.parse(l.trim());
}, 'should be valid json');
}
verifier: (t: ExecutionContext<Tester>, result: string) => {
assertNotWrappedInCodeBlock(t, result);
t.assert(
result
.split('\n')
.filter(line => line.trim())
.every(line => /^( {2})*(-|\*|\+) .+$/.test(line)),
'should be a markdown list'
);
t.false(
result
.split('\n')
.filter(line => line.trim())
.every(line => {
try {
JSON.parse(line);
return true;
} catch {
return false;
}
}),
'should not expose raw NDJSON'
);
},
},
];
for (const { name, content, verifier } of workflows) {
test(
`should be able to run workflow: ${name}`,
for (const { actionId, content, verifier } of actionRecipeCases) {
test.serial(
`should be able to run action recipe: ${actionId}`,
runIfCopilotConfigured,
async t => {
const { workflow } = t.context;
await retry(`action recipe: ${actionId}`, t, async t => {
const { actionStreams, prompt } = t.context;
const actionPrompt = await prompt.get(actionId);
if (!actionPrompt) {
return t.fail(`prompt ${actionId} should exist`);
}
const { sandbox, sessionId, userId, savedTurns } =
installActionSessionMock(t, { actionId, actionPrompt, content });
await retry(`workflow: ${name}`, t, async t => {
let result = '';
for await (const ret of workflow.runGraph({ content }, name)) {
if (ret.status === GraphExecutorState.EnterNode) {
t.log('enter node:', ret.node.name);
} else if (ret.status === GraphExecutorState.ExitNode) {
t.log('exit node:', ret.node.name);
} else if (ret.status === GraphExecutorState.EmitAttachment) {
t.log('stream attachment:', ret);
} else {
result += ret.content;
try {
const prepared = await actionStreams.stream(userId, sessionId, {
actionId,
actionVersion: 'v1',
modelId: actionPrompt.model,
});
for await (const event of prepared.stream) {
if (event.type === 'action_done' && event.status === 'succeeded') {
if (typeof event.result === 'string') {
result += event.result;
} else if (event.result && typeof event.result === 'object') {
const value = event.result as {
content?: unknown;
text?: unknown;
result?: unknown;
};
result +=
typeof value.content === 'string'
? value.content
: typeof value.text === 'string'
? value.text
: typeof value.result === 'string'
? value.result
: '';
}
}
}
} finally {
sandbox.restore();
}
t.truthy(result, 'should return result');
verifier?.(t, result);
verifier(t, result);
t.true(
savedTurns.some(turn => turn.role === 'assistant'),
'should persist assistant turn through real conversation host'
);
});
}
);
}
const TRANSCRIPT_AUDIO_CASES = [
{
name: 'short audio',
url: 'https://cdn.affine.pro/copilot-test/MP9qDGuYgnY+ILoEAmHpp3h9Npuw2403EAYMEA.mp3',
mimeType: 'audio/mpeg',
modelId: 'gemini-2.5-flash',
},
{
name: 'middle audio',
url: 'https://cdn.affine.pro/copilot-test/2ed05eo1KvZ2tWB_BAjFo67EAPZZY-w4LylUAw.m4a',
mimeType: 'audio/m4a',
modelId: 'gemini-2.5-flash',
},
{
name: 'long audio',
url: 'https://cdn.affine.pro/copilot-test/nC9-e7P85PPI2rU29QWwf8slBNRMy92teLIIMw.opus',
mimeType: 'audio/opus',
modelId: 'gemini-2.5-pro',
},
];
for (const testCase of TRANSCRIPT_AUDIO_CASES) {
test(
`should run transcript task through native action bridge: ${testCase.name}`,
runIfCopilotConfigured,
async t => {
const { models, transcript } = t.context;
const userId = `copilot-provider-transcript-user-${randomUUID()}`;
const workspaceId = `copilot-provider-transcript-workspace-${randomUUID()}`;
const blobId = `copilot-provider-transcript-blob-${randomUUID()}`;
const payload = TranscriptPayloadSchema.parse({
sourceAudio: { blobId, mimeType: testCase.mimeType },
infos: [
{
url: testCase.url,
mimeType: testCase.mimeType,
index: 0,
},
],
});
const task = await models.copilotTranscriptTask.create({
userId,
workspaceId,
blobId,
strategy: 'gemini',
recipeId: 'transcript.audio.gemini',
recipeVersion: 'v1',
inputSnapshot: payload,
publicMeta: {
sourceAudio: payload.sourceAudio,
infos: payload.infos,
},
});
await retry('transcript native action recipe', t, async t => {
await transcript.transcriptTask({
taskId: task.id,
payload,
modelId: testCase.modelId,
});
const ready = await models.copilotTranscriptTask.get(task.id);
t.is(ready?.status, 'ready');
const parsed = TranscriptPayloadSchema.parse(ready?.protectedResult);
t.is(typeof parsed.normalizedTranscript, 'string');
});
}
);
@@ -967,7 +1044,7 @@ test(
const provider = (await factory.getProviderByModel('gpt-4o-mini'))!;
t.assert(provider, 'should have provider for rerank');
const scores = await provider.rerank(
const scores = await getProviderRuntimeHost(provider).run.rerank(
{ modelId: 'gpt-4o-mini' },
{
query,
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,42 @@
import test from 'ava';
import { summarizePreparedRoutes } from '../../plugins/copilot/runtime/execution-metrics';
test('summarizePreparedRoutes should report none when no route is prepared', t => {
t.deepEqual(
summarizePreparedRoutes([{ prepared: undefined }, { prepared: undefined }]),
{
routeCount: 2,
preparedCount: 0,
preparedMode: 'none',
}
);
});
test('summarizePreparedRoutes should report partial when only some routes are prepared', t => {
t.deepEqual(
summarizePreparedRoutes([
{ prepared: { route: {} } as never },
{ prepared: undefined },
]),
{
routeCount: 2,
preparedCount: 1,
preparedMode: 'partial',
}
);
});
test('summarizePreparedRoutes should report all when every route is prepared', t => {
t.deepEqual(
summarizePreparedRoutes([
{ prepared: { route: {} } as never },
{ prepared: { route: {} } as never },
]),
{
routeCount: 2,
preparedCount: 2,
preparedMode: 'all',
}
);
});
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,73 @@
import type { LlmRequest } from '../../native';
import type { PromptMessage } from '../../plugins/copilot/providers/types';
function createPromptMessage(
role: PromptMessage['role'],
content: string,
extra: Omit<PromptMessage, 'role' | 'content'> = {}
): PromptMessage {
return {
role,
content,
...extra,
};
}
export function userPrompt(
content: string,
extra: Omit<PromptMessage, 'role' | 'content'> = {}
): PromptMessage {
return createPromptMessage('user', content, extra);
}
export function assistantPrompt(
content: string,
extra: Omit<PromptMessage, 'role' | 'content'> = {}
): PromptMessage {
return createPromptMessage('assistant', content, extra);
}
export function systemPrompt(
content: string,
extra: Omit<PromptMessage, 'role' | 'content'> = {}
): PromptMessage {
return createPromptMessage('system', content, extra);
}
export function promptMessages(...messages: PromptMessage[]) {
return messages;
}
export function singleUserPromptMessages(
content: string,
extra: Omit<PromptMessage, 'role' | 'content'> = {}
) {
return promptMessages(userPrompt(content, extra));
}
export function jsonOnlyPromptMessages(userContent: string) {
return promptMessages(
systemPrompt('Return JSON only.'),
userPrompt(userContent)
);
}
type NativeTextMessage = LlmRequest['messages'][number];
export function nativeUserText(text: string): NativeTextMessage {
return {
role: 'user',
content: [{ type: 'text', text }],
};
}
export function nativeAssistantText(text: string): NativeTextMessage {
return {
role: 'assistant',
content: [{ type: 'text', text }],
};
}
export function nativeMessages(...messages: NativeTextMessage[]) {
return messages;
}
@@ -7,14 +7,7 @@ import { CopilotProviderType } from '../../plugins/copilot/providers/types';
test('resolveProviderMiddleware should include anthropic defaults', t => {
const middleware = resolveProviderMiddleware(CopilotProviderType.Anthropic);
t.deepEqual(middleware.rust?.request, [
'normalize_messages',
'tool_schema_rewrite',
]);
t.deepEqual(middleware.rust?.stream, [
'stream_event_normalize',
'citation_indexing',
]);
t.is(middleware.rust, undefined);
t.deepEqual(middleware.node?.text, ['citation_footnote', 'callout']);
});
@@ -24,10 +17,7 @@ test('resolveProviderMiddleware should merge defaults and overrides', t => {
node: { text: ['thinking_format'] },
});
t.deepEqual(middleware.rust?.request, [
'normalize_messages',
'clamp_max_tokens',
]);
t.deepEqual(middleware.rust?.request, ['clamp_max_tokens']);
t.deepEqual(middleware.node?.text, [
'citation_footnote',
'callout',
@@ -48,9 +38,6 @@ test('buildProviderRegistry should normalize profile middleware defaults', t =>
const profile = registry.profiles.get('openai-main');
t.truthy(profile);
t.deepEqual(profile?.middleware.rust?.stream, [
'stream_event_normalize',
'citation_indexing',
]);
t.is(profile?.middleware.rust, undefined);
t.deepEqual(profile?.middleware.node?.text, ['citation_footnote', 'callout']);
});
File diff suppressed because it is too large Load Diff
@@ -1,5 +1,7 @@
import test from 'ava';
import { OpenAIProvider } from '../../plugins/copilot/providers';
import { CopilotProviderLifecycleService } from '../../plugins/copilot/providers/lifecycle-service';
import {
buildProviderRegistry,
resolveModel,
@@ -142,6 +144,46 @@ test('resolveModel should follow defaults -> fallback -> order and apply filters
t.deepEqual(routed.candidateProviderIds, ['openai-main', 'fal-main']);
});
test('resolveModel should resolve bare model ids by provider priority order', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
priority: 10,
config: { apiKey: '1' },
},
{
id: 'anthropic-main',
type: CopilotProviderType.Anthropic,
priority: 5,
config: { apiKey: '2' },
},
{
id: 'fal-main',
type: CopilotProviderType.FAL,
priority: 1,
config: { apiKey: '3' },
},
],
defaults: {
[ModelOutputType.Text]: 'anthropic-main',
fallback: 'fal-main',
},
});
const routed = resolveModel({
registry,
modelId: 'shared-model',
});
t.deepEqual(routed.candidateProviderIds, [
'openai-main',
'anthropic-main',
'fal-main',
]);
});
test('stripProviderPrefix should only strip matched provider prefix', t => {
const registry = buildProviderRegistry({
profiles: [
@@ -166,3 +208,75 @@ test('stripProviderPrefix should only strip matched provider prefix', t => {
'gpt-5-mini'
);
});
test('CopilotProviderLifecycleService should register current profiles and unregister stale ones', async t => {
const calls: string[] = [];
let registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
{
id: 'openai-backup',
type: CopilotProviderType.OpenAI,
config: { apiKey: '2' },
},
],
});
const provider = {
type: CopilotProviderType.OpenAI,
configured(execution: { providerId?: string } | undefined) {
return execution?.providerId === 'openai-main';
},
};
const service = new CopilotProviderLifecycleService(
{
get(token: unknown) {
return token === OpenAIProvider ? provider : undefined;
},
} as any,
{
register(providerId: string) {
calls.push(`register:${providerId}`);
},
unregister(providerId: string) {
calls.push(`unregister:${providerId}`);
},
} as any,
{
getRegistry() {
return registry;
},
} as any
);
await service.syncProviders();
t.deepEqual(calls.slice().sort(), [
'register:openai-main',
'unregister:openai-backup',
]);
calls.length = 0;
registry = buildProviderRegistry({
profiles: [
{
id: 'openai-backup',
type: CopilotProviderType.OpenAI,
config: { apiKey: '2' },
},
],
});
provider.configured = (execution: { providerId?: string } | undefined) =>
execution?.providerId === 'openai-backup';
await service.syncProviders();
t.deepEqual(calls.slice().sort(), [
'register:openai-backup',
'unregister:openai-main',
]);
});
@@ -0,0 +1,201 @@
import serverNativeModule from '@affine/server-native';
import test from 'ava';
import { z } from 'zod';
import type {
LlmEmbeddingRequest,
LlmRerankRequest,
LlmStructuredRequest,
} from '../../native';
import { CopilotProvider } from '../../plugins/copilot/providers/provider';
import type { ProviderDriverSpec } from '../../plugins/copilot/providers/provider-runtime-contract';
import { CopilotProviderType } from '../../plugins/copilot/providers/types';
import {
buildStructuredResponseContract,
type RequiredStructuredOutputContract,
requireStructuredOutputContract,
} from '../../plugins/copilot/runtime/contracts';
import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context';
import { nativeUserText, singleUserPromptMessages } from './prompt-test-helper';
function structuredOptions(schema: z.ZodTypeAny) {
const { responseSchemaJson, schemaHash } =
buildStructuredResponseContract(schema);
return { responseSchemaJson, schemaHash };
}
function structuredContract(
schema: z.ZodTypeAny
): RequiredStructuredOutputContract {
const contract = buildStructuredResponseContract(schema);
const requiredContract = requireStructuredOutputContract(contract);
if (!requiredContract) {
throw new Error('structured response contract is required');
}
return requiredContract;
}
class TemplateOnlyProvider extends CopilotProvider<{ apiKey: string }> {
readonly type = CopilotProviderType.OpenAI;
protected resolveModelBackendKind() {
return 'openai_responses' as const;
}
readonly structuredRequests: LlmStructuredRequest[] = [];
readonly embeddingRequests: LlmEmbeddingRequest[] = [];
readonly rerankRequests: Array<{
model: string;
query: string;
candidates: Array<{ id?: string; text: string }>;
topN?: number;
}> = [];
configured() {
return true;
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: () => ({
base_url: 'https://api.openai.com',
auth_token: 'test-key',
}),
mapError: (error: unknown) => error,
structured: {},
embedding: {
defaultDimensions: 8,
},
rerank: {},
};
}
}
test('template-only provider should reuse base structured, embedding and rerank drivers', async t => {
const provider = new TemplateOnlyProvider();
const originalStructured = (serverNativeModule as any).llmStructuredDispatch;
const originalEmbedding = (serverNativeModule as any).llmEmbeddingDispatch;
const originalRerank = (serverNativeModule as any).llmRerankDispatch;
(serverNativeModule as any).llmStructuredDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
provider.structuredRequests.push(
JSON.parse(requestJson) as LlmStructuredRequest
);
return JSON.stringify({
id: 'structured_1',
model: 'gpt-5-mini',
output_text: '{"summary":"native"}',
output_json: { summary: 'native' },
usage: {
prompt_tokens: 3,
completion_tokens: 2,
total_tokens: 5,
},
finish_reason: 'stop',
});
};
(serverNativeModule as any).llmEmbeddingDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as LlmEmbeddingRequest;
provider.embeddingRequests.push(request);
return JSON.stringify({
model: request.model,
embeddings: request.inputs.map((_, index) => [index + 0.1, index + 0.2]),
});
};
(serverNativeModule as any).llmRerankDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as LlmRerankRequest;
provider.rerankRequests.push(request);
return JSON.stringify({
model: request.model,
scores: request.candidates.map((_candidate, index) =>
index === 0 ? 0.9 : 0.1
),
});
};
t.teardown(() => {
(serverNativeModule as any).llmStructuredDispatch = originalStructured;
(serverNativeModule as any).llmEmbeddingDispatch = originalEmbedding;
(serverNativeModule as any).llmRerankDispatch = originalRerank;
});
const structured = await getProviderRuntimeHost(provider).run.structured(
{ modelId: 'gpt-5-mini' },
singleUserPromptMessages('summarize this'),
structuredOptions(z.object({ summary: z.string() })),
structuredContract(z.object({ summary: z.string() }))
);
const embeddings = await getProviderRuntimeHost(provider).run.embedding(
{ modelId: 'text-embedding-3-small' },
['alpha', 'beta'],
{
dimensions: 8,
}
);
const scores = await getProviderRuntimeHost(provider).run.rerank(
{ modelId: 'gpt-4o-mini' },
{
query: 'alpha',
candidates: [
{ id: 'alpha', text: 'alpha result' },
{ id: 'beta', text: 'beta result' },
],
topK: 1,
}
);
t.is(structured, JSON.stringify({ summary: 'native' }));
t.deepEqual(embeddings, [
[0.1, 0.2],
[1.1, 1.2],
]);
t.deepEqual(scores, [0.9, 0.1]);
t.is(provider.structuredRequests.length, 1);
t.like(provider.structuredRequests[0], {
model: 'gpt-5-mini',
messages: [
{ role: 'user', content: nativeUserText('summarize this').content },
],
schema: {
type: 'object',
properties: {
summary: { type: 'string' },
},
required: ['summary'],
additionalProperties: false,
},
strict: true,
responseMimeType: 'application/json',
});
t.is(provider.structuredRequests[0]?.middleware, undefined);
t.deepEqual(provider.embeddingRequests, [
{
model: 'text-embedding-3-small',
inputs: ['alpha', 'beta'],
dimensions: 8,
taskType: 'RETRIEVAL_DOCUMENT',
},
]);
t.deepEqual(provider.rerankRequests, [
{
model: 'gpt-4o-mini',
query: 'alpha',
candidates: [
{ id: 'alpha', text: 'alpha result' },
{ id: 'beta', text: 'beta result' },
],
topN: 1,
},
]);
});
@@ -1,15 +1,26 @@
import serverNativeModule from '@affine/server-native';
import test from 'ava';
import { z } from 'zod';
import type { DocReader } from '../../core/doc';
import type { AccessController } from '../../core/permission';
import type { Models } from '../../models';
import { NativeLlmRequest, NativeLlmStreamEvent } from '../../native';
import {
ToolCallAccumulator,
ToolCallLoop,
ToolSchemaExtractor,
} from '../../plugins/copilot/providers/loop';
LlmRequest,
type LlmToolCallbackRequest,
type LlmToolCallbackResponse,
type LlmToolLoopStreamEvent,
llmValidateContract,
} from '../../native';
import {
buildToolContracts,
parseToolContract,
parseToolLoopStreamEvent,
} from '../../plugins/copilot/runtime/contracts';
import {
createToolExecutionCallback,
createToolLoopBridge,
} from '../../plugins/copilot/runtime/tool/bridge';
import {
buildBlobContentGetter,
createBlobReadTool,
@@ -30,100 +41,47 @@ import {
DOCUMENT_SYNC_PENDING_MESSAGE,
LOCAL_WORKSPACE_SYNC_REQUIRED_MESSAGE,
} from '../../plugins/copilot/tools/doc-sync';
import { defineTool } from '../../plugins/copilot/tools/tool';
import {
nativeMessages,
nativeUserText,
singleUserPromptMessages,
} from './prompt-test-helper';
test('ToolCallAccumulator should merge deltas and complete tool call', t => {
const accumulator = new ToolCallAccumulator();
accumulator.feedDelta({
type: 'tool_call_delta',
call_id: 'call_1',
name: 'doc_read',
arguments_delta: '{"doc_id":"',
});
accumulator.feedDelta({
type: 'tool_call_delta',
call_id: 'call_1',
arguments_delta: 'a1"}',
test('defineTool should freeze json schema at definition time', t => {
const tool = defineTool({
description: 'Read doc',
inputSchema: z.object({
doc_id: z.string(),
limit: z.number().optional(),
}),
execute: async () => ({}),
});
const completed = accumulator.complete({
type: 'tool_call',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
});
t.deepEqual(completed, {
id: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
thought: undefined,
t.deepEqual(tool.jsonSchema, {
type: 'object',
properties: {
doc_id: { type: 'string' },
limit: { type: 'number' },
},
additionalProperties: false,
required: ['doc_id'],
});
});
test('ToolCallAccumulator should preserve invalid JSON instead of swallowing it', t => {
const accumulator = new ToolCallAccumulator();
accumulator.feedDelta({
type: 'tool_call_delta',
call_id: 'call_1',
name: 'doc_read',
arguments_delta: '{"doc_id":',
});
const pending = accumulator.drainPending();
t.is(pending.length, 1);
t.deepEqual(pending[0]?.id, 'call_1');
t.deepEqual(pending[0]?.name, 'doc_read');
t.deepEqual(pending[0]?.args, {});
t.is(pending[0]?.rawArgumentsText, '{"doc_id":');
t.truthy(pending[0]?.argumentParseError);
});
test('ToolCallAccumulator should prefer native canonical tool arguments metadata', t => {
const accumulator = new ToolCallAccumulator();
accumulator.feedDelta({
type: 'tool_call_delta',
call_id: 'call_1',
name: 'doc_read',
arguments_delta: '{"stale":true}',
});
const completed = accumulator.complete({
type: 'tool_call',
call_id: 'call_1',
name: 'doc_read',
arguments: {},
arguments_text: '{"doc_id":"a1"}',
arguments_error: 'invalid json',
});
t.deepEqual(completed, {
id: 'call_1',
name: 'doc_read',
args: {},
rawArgumentsText: '{"doc_id":"a1"}',
argumentParseError: 'invalid json',
thought: undefined,
});
});
test('ToolSchemaExtractor should convert zod schema to json schema', t => {
test('buildToolContracts should project precomputed json schema', t => {
const toolSet = {
doc_read: {
doc_read: defineTool({
description: 'Read doc',
inputSchema: z.object({
doc_id: z.string(),
limit: z.number().optional(),
}),
execute: async () => ({}),
},
}),
};
const extracted = ToolSchemaExtractor.extract(toolSet);
const extracted = buildToolContracts(toolSet);
t.deepEqual(extracted, [
{
@@ -142,43 +100,224 @@ test('ToolSchemaExtractor should convert zod schema to json schema', t => {
]);
});
test('ToolCallLoop should execute tool call and continue to next round', async t => {
const dispatchRequests: NativeLlmRequest[] = [];
const originalMessages = [{ role: 'user', content: 'read doc' }] as const;
test('buildToolContracts should reject tool definitions without json schema', t => {
const error = t.throws(() =>
buildToolContracts({
doc_read: {
description: 'Read doc',
inputSchema: z.object({ doc_id: z.string() }),
execute: async () => ({}),
} as never,
})
);
t.regex(error.message, /missing precomputed jsonSchema/);
});
test('defineTool should prefer explicit json schema when provided', t => {
const extracted = buildToolContracts({
doc_read: defineTool({
description: 'Read doc',
jsonSchema: {
type: 'object',
properties: {
doc_id: { type: 'string' },
},
required: ['doc_id'],
},
inputSchema: z.object({
doc_id: z.string(),
ignored: z.number(),
}),
execute: async () => ({}),
}),
});
t.deepEqual(extracted, [
{
name: 'doc_read',
description: 'Read doc',
parameters: {
type: 'object',
properties: {
doc_id: { type: 'string' },
},
required: ['doc_id'],
},
},
]);
});
test('ToolContract should freeze stable tool schema and callback payloads', t => {
const tool = parseToolContract({
name: 'doc_read',
description: 'Read doc',
parameters: {
type: 'object',
properties: {
doc_id: { type: 'string' },
},
required: ['doc_id'],
},
});
const result = llmValidateContract<LlmToolCallbackResponse>(
'toolCallbackResponse',
{
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
output: { markdown: '# a1' },
}
);
const request = llmValidateContract<LlmToolCallbackRequest>(
'toolCallbackRequest',
{
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
}
);
t.is(tool.name, 'doc_read');
t.deepEqual(request.args, { doc_id: 'a1' });
t.deepEqual(result.args, { doc_id: 'a1' });
});
test('ToolLoopStreamEvent should reject malformed tool_result metadata at decode boundary', t => {
const event = parseToolLoopStreamEvent({
type: 'tool_result',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
output: { markdown: '# a1' },
});
t.is(event.type, 'tool_result');
const error = t.throws(() =>
parseToolLoopStreamEvent({
type: 'tool_result',
call_id: 'call_1',
output: { markdown: '# a1' },
})
);
t.truthy(error);
});
test('createNativeToolExecutionCallback should preserve tool execution ABI', async t => {
const callback = createToolExecutionCallback(
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
execute: async args => ({ markdown: `# ${String(args.doc_id)}` }),
},
},
{ messages: singleUserPromptMessages('read doc') }
);
const result = await callback({
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
});
t.deepEqual(result, {
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
argumentParseError: undefined,
output: { markdown: '# a1' },
});
});
test('createNativeToolLoopBridge should preserve native callback and stream ABI', async t => {
const capturedRequests: LlmRequest[] = [];
const originalMessages = singleUserPromptMessages('read doc');
const signal = new AbortController().signal;
let executedArgs: Record<string, unknown> | null = null;
let executedMessages: unknown;
let executedSignal: AbortSignal | undefined;
const dispatch = (request: NativeLlmRequest) => {
dispatchRequests.push(request);
const round = dispatchRequests.length;
const original = (serverNativeModule as any).llmDispatchToolLoopStream;
(serverNativeModule as any).llmDispatchToolLoopStream = (
_protocol: string,
_backendConfigJson: string,
requestJson: string,
maxSteps: number,
callback: (error: Error | null, eventJson: string) => void,
toolCallback: (error: Error | null, requestJson: string) => Promise<string>
) => {
capturedRequests.push(JSON.parse(requestJson) as LlmRequest);
t.is(maxSteps, 4);
return (async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
if (round === 1) {
yield {
type: 'tool_call_delta',
call_id: 'call_1',
name: 'doc_read',
arguments_delta: '{"doc_id":"a1"}',
};
yield {
void (async () => {
callback(
null,
JSON.stringify({
type: 'tool_call',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
};
yield { type: 'done', finish_reason: 'tool_calls' };
return;
}
})
);
yield { type: 'text_delta', text: 'done' };
yield { type: 'done', finish_reason: 'stop' };
const result = JSON.parse(
await toolCallback(
null,
JSON.stringify({
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
})
)
) as {
callId: string;
name: string;
args: Record<string, unknown>;
rawArgumentsText?: string;
argumentParseError?: string;
output: unknown;
isError?: boolean;
};
callback(
null,
JSON.stringify({
type: 'tool_result',
call_id: result.callId,
name: result.name,
arguments: result.args,
arguments_text: result.rawArgumentsText,
arguments_error: result.argumentParseError,
output: result.output,
is_error: result.isError,
})
);
callback(null, JSON.stringify({ type: 'text_delta', text: 'done' }));
callback(null, JSON.stringify({ type: 'done', finish_reason: 'stop' }));
callback(null, '__AFFINE_LLM_STREAM_END__');
})();
};
let executedArgs: Record<string, unknown> | null = null;
let executedMessages: unknown;
let executedSignal: AbortSignal | undefined;
const loop = new ToolCallLoop(
dispatch,
return {
abort() {},
};
};
t.teardown(() => {
(serverNativeModule as any).llmDispatchToolLoopStream = original;
});
const bridge = createToolLoopBridge(
{
protocol: 'openai_chat',
backendConfig: {
base_url: 'https://api.openai.com',
auth_token: 'test-key',
},
},
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
@@ -193,14 +332,12 @@ test('ToolCallLoop should execute tool call and continue to next round', async t
4
);
const events: NativeLlmStreamEvent[] = [];
for await (const event of loop.run(
const events: LlmToolLoopStreamEvent[] = [];
for await (const event of bridge(
{
model: 'gpt-5-mini',
stream: true,
messages: [
{ role: 'user', content: [{ type: 'text', text: 'read doc' }] },
],
stream: false,
messages: nativeMessages(nativeUserText('read doc')),
},
signal,
[...originalMessages]
@@ -211,105 +348,13 @@ test('ToolCallLoop should execute tool call and continue to next round', async t
t.deepEqual(executedArgs, { doc_id: 'a1' });
t.deepEqual(executedMessages, originalMessages);
t.is(executedSignal, signal);
t.true(
dispatchRequests[1]?.messages.some(message => message.role === 'tool')
);
t.deepEqual(dispatchRequests[1]?.messages[1]?.content, [
{
type: 'tool_call',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
arguments_text: '{"doc_id":"a1"}',
arguments_error: undefined,
thought: undefined,
},
]);
t.deepEqual(dispatchRequests[1]?.messages[2]?.content, [
{
type: 'tool_result',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
arguments_text: '{"doc_id":"a1"}',
arguments_error: undefined,
output: { markdown: '# doc' },
is_error: undefined,
},
]);
t.true(capturedRequests[0]?.stream);
t.deepEqual(
events.map(event => event.type),
['tool_call', 'tool_result', 'text_delta', 'done']
);
});
test('ToolCallLoop should surface invalid JSON as tool error without executing', async t => {
let executed = false;
let round = 0;
const loop = new ToolCallLoop(
request => {
round += 1;
const hasToolResult = request.messages.some(
message => message.role === 'tool'
);
return (async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
if (!hasToolResult && round === 1) {
yield {
type: 'tool_call_delta',
call_id: 'call_1',
name: 'doc_read',
arguments_delta: '{"doc_id":',
};
yield { type: 'done', finish_reason: 'tool_calls' };
return;
}
yield { type: 'done', finish_reason: 'stop' };
})();
},
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
execute: async () => {
executed = true;
return { markdown: '# doc' };
},
},
},
2
);
const events: NativeLlmStreamEvent[] = [];
for await (const event of loop.run({
model: 'gpt-5-mini',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'read doc' }] }],
})) {
events.push(event);
}
t.false(executed);
t.true(events[0]?.type === 'tool_result');
t.deepEqual(events[0], {
type: 'tool_result',
call_id: 'call_1',
name: 'doc_read',
arguments: {},
arguments_text: '{"doc_id":',
arguments_error:
events[0]?.type === 'tool_result' ? events[0].arguments_error : undefined,
output: {
message: 'Invalid tool arguments JSON',
rawArguments: '{"doc_id":',
error:
events[0]?.type === 'tool_result'
? events[0].arguments_error
: undefined,
},
is_error: true,
});
});
test('doc_read should return specific sync errors for unavailable docs', async t => {
const cases = [
{
@@ -434,7 +479,7 @@ test('document search tools should return sync error for local workspace', async
);
const semanticTool = createDocSemanticSearchTool(
buildDocSearchGetter(ac, contextService, null, models).bind(null, {
buildDocSearchGetter(ac, contextService, undefined, models).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
@@ -478,7 +523,7 @@ test('doc_semantic_search should return empty array when nothing matches', async
} as unknown as Parameters<typeof buildDocSearchGetter>[1];
const semanticTool = createDocSemanticSearchTool(
buildDocSearchGetter(ac, contextService, null, models).bind(null, {
buildDocSearchGetter(ac, contextService, undefined, models).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
File diff suppressed because it is too large Load Diff
@@ -1,12 +1,16 @@
import { randomBytes } from 'node:crypto';
import serverNativeModule from '@affine/server-native';
import type { ProviderMiddlewareConfig } from '../../plugins/copilot/config';
import {
CopilotChatOptions,
CopilotEmbeddingOptions,
CopilotImageOptions,
type CopilotProviderModel,
CopilotProviderType,
CopilotStructuredOptions,
ModelConditions,
ModelInputType,
ModelFullConditions,
ModelOutputType,
PromptMessage,
StreamObject,
@@ -15,130 +19,534 @@ import {
DEFAULT_DIMENSIONS,
OpenAIProvider,
} from '../../plugins/copilot/providers/openai';
import type { ProviderModelRuntimeContext } from '../../plugins/copilot/providers/provider-model-runtime';
import {
type CopilotProviderExecution,
createNativeExecutionDriverSpec,
type ProviderDriverSpec,
} from '../../plugins/copilot/providers/provider-runtime-contract';
import type { ProviderRuntimeContexts } from '../../plugins/copilot/runtime/provider-runtime-context';
import { sleep } from '../utils/utils';
export class MockCopilotProvider extends OpenAIProvider {
override readonly models = [
{
id: 'test',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text, ModelOutputType.Object],
defaultForOutputType: true,
},
],
},
{
id: 'test-image',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Image],
defaultForOutputType: true,
},
],
},
{
id: 'gpt-5',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
},
],
},
{
id: 'gpt-5-2025-08-07',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
},
],
},
{
id: 'gpt-5-mini',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
},
],
},
{
id: 'gpt-5-nano',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
},
],
},
{
id: 'gpt-image-1',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Image],
defaultForOutputType: true,
},
],
},
{
id: 'gemini-2.5-flash',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
},
],
},
{
id: 'gemini-2.5-pro',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
},
],
},
{
id: 'gemini-3.1-pro-preview',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
},
],
},
];
const LLM_STREAM_END_MARKER = '__AFFINE_LLM_STREAM_END__';
const MOCK_NATIVE_TEXT = 'generate text to text';
const MOCK_NATIVE_STREAM_TEXT = 'generate text to text stream';
override async text(
function mockUsage() {
return {
prompt_tokens: 1,
completion_tokens: 1,
total_tokens: 2,
};
}
function buildMockDispatchResponse(model: string, text: string) {
return {
id: 'mock-dispatch',
model,
message: {
role: 'assistant',
content: [{ type: 'text', text }],
},
usage: mockUsage(),
finish_reason: 'stop',
};
}
function buildMockStructuredValue(schema: any, key?: string): any {
if (!schema || typeof schema !== 'object') {
return key === 'title' ? 'Weekly Sync' : MOCK_NATIVE_TEXT;
}
if (Array.isArray(schema.anyOf) && schema.anyOf.length > 0) {
return buildMockStructuredValue(schema.anyOf[0], key);
}
if (Array.isArray(schema.oneOf) && schema.oneOf.length > 0) {
return buildMockStructuredValue(schema.oneOf[0], key);
}
if (Array.isArray(schema.enum) && schema.enum.length > 0) {
return schema.enum[0];
}
switch (schema.type) {
case 'object': {
const properties =
schema.properties && typeof schema.properties === 'object'
? schema.properties
: {};
return Object.fromEntries(
Object.entries(properties).map(([key, value]) => [
key,
buildMockStructuredValue(value, key),
])
);
}
case 'array':
return [buildMockStructuredValue(schema.items, key)];
case 'boolean':
return true;
case 'number':
case 'integer':
switch (key) {
case 'durationMinutes':
return 45;
case 's':
return 30;
case 'e':
return 53;
default:
return 1;
}
case 'null':
return null;
case 'string':
default:
switch (key) {
case 'title':
return 'Weekly Sync';
case 'description':
return 'Send recap';
case 'owner':
return 'A';
case 'deadline':
return 'Friday';
case 'speaker':
case 'a':
return 'A';
case 'attendees':
return 'A';
case 'start':
return '00:00:42';
case 'end':
return '00:01:05';
case 'text':
case 'transcription':
case 't':
return 'Hello, everyone.';
case 'keyPoints':
return 'Reviewed launch status';
case 'decisions':
return 'Ship on Monday';
case 'openQuestions':
return 'Need final QA sign-off';
case 'blockers':
return 'Waiting on analytics';
case 'summary':
return 'Reviewed launch status';
default:
return MOCK_NATIVE_TEXT;
}
}
}
function parseFirstRoute(routesJson: string) {
const routes = JSON.parse(routesJson) as Array<{
provider_id?: string;
model?: string;
request?: {
model?: string;
operation?: string;
prompt?: string;
schema?: unknown;
};
}>;
return routes[0];
}
function buildMockStructuredResponse(model: string, schema: unknown) {
const output_json = buildMockStructuredValue(schema);
return {
id: 'mock-structured-dispatch',
model,
output_text: JSON.stringify(output_json),
output_json,
usage: mockUsage(),
finish_reason: 'stop',
};
}
function emitMockTextStream(
model: string,
callback: (error: Error | null, eventJson: string) => void
) {
callback(null, JSON.stringify({ type: 'message_start', model }));
for (const text of MOCK_NATIVE_STREAM_TEXT) {
callback(null, JSON.stringify({ type: 'text_delta', text }));
}
callback(
null,
JSON.stringify({
type: 'done',
finish_reason: 'stop',
usage: mockUsage(),
})
);
callback(null, LLM_STREAM_END_MARKER);
}
export function installMockCopilotRuntime() {
const native = serverNativeModule as Record<string, any>;
const original = {
llmDispatchPrepared: native.llmDispatchPrepared,
llmDispatchPreparedStream: native.llmDispatchPreparedStream,
llmRenderBuiltInPrompt: native.llmRenderBuiltInPrompt,
llmRenderBuiltInSessionPrompt: native.llmRenderBuiltInSessionPrompt,
llmValidateJsonSchema: native.llmValidateJsonSchema,
llmStructuredDispatch: native.llmStructuredDispatch,
llmStructuredDispatchPrepared: native.llmStructuredDispatchPrepared,
llmEmbeddingDispatch: native.llmEmbeddingDispatch,
llmEmbeddingDispatchPrepared: native.llmEmbeddingDispatchPrepared,
llmRerankDispatch: native.llmRerankDispatch,
llmRerankDispatchPrepared: native.llmRerankDispatchPrepared,
llmImageDispatchPrepared: native.llmImageDispatchPrepared,
runNativeActionRecipePreparedStream:
native.runNativeActionRecipePreparedStream,
};
native.llmDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response: buildMockDispatchResponse(
route?.request?.model ?? route?.model ?? 'test',
MOCK_NATIVE_TEXT
),
});
};
native.llmDispatchPreparedStream = (
routesJson: string,
callback: (error: Error | null, eventJson: string) => void
) => {
const route = parseFirstRoute(routesJson);
emitMockTextStream(
route?.request?.model ?? route?.model ?? 'test',
callback
);
return { abort() {} };
};
native.llmStructuredDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as {
model?: string;
schema?: unknown;
};
return JSON.stringify(
buildMockStructuredResponse(request.model ?? 'test', request.schema)
);
};
native.llmStructuredDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response: buildMockStructuredResponse(
route?.request?.model ?? route?.model ?? 'test',
route?.request?.schema
),
});
};
native.llmValidateJsonSchema = (_schema: unknown, value: unknown) => value;
native.llmEmbeddingDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as {
model?: string;
dimensions?: number;
};
const length = request.dimensions ?? DEFAULT_DIMENSIONS;
return JSON.stringify({
model: request.model ?? 'test',
embeddings: [
Array.from({ length }, (_value, index) => (index % 128) + 1),
],
usage: { prompt_tokens: 1, total_tokens: 1 },
});
};
native.llmEmbeddingDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
const response = JSON.parse(
native.llmEmbeddingDispatch(
'',
'',
JSON.stringify(route?.request ?? { model: route?.model ?? 'test' })
)
) as Record<string, unknown>;
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response,
});
};
native.llmRerankDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as {
model?: string;
candidates?: unknown[];
};
const candidateCount = request.candidates?.length ?? 0;
return JSON.stringify({
model: request.model ?? 'test',
scores: Array.from(
{ length: candidateCount },
(_value, index) => candidateCount - index
),
});
};
native.llmRerankDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
const response = JSON.parse(
native.llmRerankDispatch(
'',
'',
JSON.stringify(route?.request ?? { model: route?.model ?? 'test' })
)
) as Record<string, unknown>;
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response,
});
};
native.llmImageDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
const model = route?.request?.model ?? route?.model ?? 'test-image';
const images = [
{
url: `https://example.com/${model}.jpg`,
media_type: 'image/jpeg',
},
];
if (route?.request?.operation === 'edit' && route.request.prompt) {
images.push({
url: `https://example.com/generated/${encodeURIComponent(route.request.prompt)}.jpg`,
media_type: 'image/jpeg',
});
}
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response: {
images,
},
});
};
native.runNativeActionRecipePreparedStream = (
input: {
recipeId: string;
recipeVersion?: string;
input?: Record<string, any>;
},
callback: (error: Error | null, eventJson: string) => void
) => {
const version = input.recipeVersion ?? 'v1';
const result = input.recipeId.startsWith('image.filter.')
? {
url: `https://example.com/${input.recipeId}.jpg`,
}
: MOCK_NATIVE_STREAM_TEXT;
const attachmentEvent = input.recipeId.startsWith('image.filter.')
? [
{
type: 'attachment',
actionId: input.recipeId,
actionVersion: version,
status: 'running',
attachment: result,
},
]
: [];
const events = [
{
type: 'action_start',
actionId: input.recipeId,
actionVersion: version,
status: 'running',
},
{
type: 'step_start',
actionId: input.recipeId,
actionVersion: version,
stepId: 'generate',
status: 'running',
},
...attachmentEvent,
{
type: 'step_end',
actionId: input.recipeId,
actionVersion: version,
stepId: 'generate',
status: 'running',
},
{
type: 'action_done',
actionId: input.recipeId,
actionVersion: version,
status: 'succeeded',
result,
trace: {
actionId: input.recipeId,
actionVersion: version,
status: 'succeeded',
lightweight: [
{ type: 'action_start', status: 'running' },
{ type: 'action_trace', status: 'succeeded' },
],
},
},
];
for (const event of events) {
callback(null, JSON.stringify(event));
}
callback(null, LLM_STREAM_END_MARKER);
return { abort() {} };
};
return () => {
Object.assign(native, original);
};
}
export class MockCopilotProvider extends OpenAIProvider {
private runtimeHostOverride?: ProviderRuntimeContexts;
protected override resolveModelRuntimeContext(): ProviderModelRuntimeContext {
const providerType = this.type as CopilotProviderType;
return {
type: providerType,
backendKind:
providerType === CopilotProviderType.Gemini
? 'gemini_api'
: 'openai_responses',
};
}
override getDriverSpec(): ProviderDriverSpec {
const spec = super.getDriverSpec();
return {
...spec,
image: {
prepareMessages: async messages => messages,
},
};
}
private resolveMockModelId(
cond: Pick<ModelFullConditions, 'modelId' | 'outputType'>
) {
if (cond.modelId === 'test') {
return 'gpt-5-mini';
}
if (cond.modelId === 'test-image') {
return 'gpt-image-1';
}
return cond.modelId;
}
private normalizeMockConditions(
cond: ModelFullConditions
): ModelFullConditions {
const modelId = this.resolveMockModelId(cond);
return modelId === cond.modelId ? cond : { ...cond, modelId };
}
protected override createDriverSpec(spec: ProviderDriverSpec) {
return createNativeExecutionDriverSpec(spec, {
createBackendConfig: spec.createBackendConfig,
mapError: spec.mapError,
checkParams: input => this.checkParams(input),
selectModel: (cond, execution) => this.selectModel(cond, execution),
getTools: this.getTools.bind(this),
getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this),
});
}
override async match(
cond: ModelFullConditions = {},
execution?: CopilotProviderExecution
) {
return await super.match(this.normalizeMockConditions(cond), execution);
}
override resolveModel(
modelId: string,
execution?: CopilotProviderExecution
): CopilotProviderModel | undefined {
const resolvedModelId = this.resolveMockModelId({ modelId });
return resolvedModelId
? super.resolveModel(resolvedModelId, execution)
: undefined;
}
override selectModel(
cond: ModelFullConditions,
execution?: CopilotProviderExecution
): CopilotProviderModel {
return super.selectModel(this.normalizeMockConditions(cond), execution);
}
override checkParams(input: Parameters<OpenAIProvider['checkParams']>[0]) {
return super.checkParams({
...input,
cond: this.normalizeMockConditions(input.cond),
});
}
override getActiveProviderMiddleware(): ProviderMiddlewareConfig {
return {};
}
overrideRuntimeHost(runtimeHost: ProviderRuntimeContexts) {
if (!this.runtimeHostOverride) {
const runtimeHostOverride: ProviderRuntimeContexts = {
...runtimeHost,
run: {
...runtimeHost.run,
text: this.text.bind(this),
streamText: this.streamTextRuntime.bind(this),
streamObject: this.streamObjectRuntime.bind(this),
structured: this.structure.bind(this),
embedding: this.embedding.bind(this),
},
};
this.runtimeHostOverride = runtimeHostOverride;
}
return this.runtimeHostOverride;
}
private async *streamTextRuntime(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions
): AsyncIterableIterator<string> {
yield* this.streamText(cond, messages, options);
}
private async *streamObjectRuntime(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions
): AsyncIterableIterator<StreamObject> {
yield* this.streamObject(cond, messages, options);
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
@@ -147,19 +555,27 @@ export class MockCopilotProvider extends OpenAIProvider {
...cond,
outputType: ModelOutputType.Text,
};
await this.checkParams({ messages, cond: fullCond, options });
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
return 'generate text to text';
}
override async *streamText(
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
await this.checkParams({ messages, cond: fullCond, options });
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
@@ -173,70 +589,58 @@ export class MockCopilotProvider extends OpenAIProvider {
}
}
override async structure(
async structure(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotStructuredOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Structured };
await this.checkParams({ messages, cond: fullCond, options });
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
return 'generate text to text';
}
override async *streamImages(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotImageOptions = {}
) {
const fullCond = { ...cond, outputType: ModelOutputType.Image };
await this.checkParams({ messages, cond: fullCond, options });
// make some time gap for history test case
await sleep(100);
const { content: prompt } = [...messages].pop() || {};
if (!prompt) throw new Error('Prompt is required');
const imageUrls = [
`https://example.com/${cond.modelId || 'test'}.jpg`,
prompt,
];
for (const imageUrl of imageUrls) {
yield imageUrl;
if (options.signal?.aborted) {
break;
}
}
return;
}
// ====== text to embedding ======
override async embedding(
async embedding(
cond: ModelConditions,
messages: string | string[],
options: CopilotEmbeddingOptions = { dimensions: DEFAULT_DIMENSIONS }
): Promise<number[][]> {
messages = Array.isArray(messages) ? messages : [messages];
const fullCond = { ...cond, outputType: ModelOutputType.Embedding };
await this.checkParams({ embeddings: messages, cond: fullCond, options });
await this.checkParams({
embeddings: messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
return [Array.from(randomBytes(options.dimensions)).map(v => v % 128)];
return [
Array.from(randomBytes(options.dimensions ?? DEFAULT_DIMENSIONS)).map(
v => v % 128
),
];
}
override async *streamObject(
async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<StreamObject> {
const fullCond = { ...cond, outputType: ModelOutputType.Object };
await this.checkParams({ messages, cond: fullCond, options });
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
@@ -1,11 +1,12 @@
export { createFactory } from './factory';
export * from './prompt-service.mock';
export * from './team-workspace.mock';
export * from './user.mock';
export * from './workspace.mock';
export * from './workspace-user.mock';
import { MockAccessToken } from './access-token.mock';
import { MockCopilotProvider } from './copilot.mock';
import { installMockCopilotRuntime, MockCopilotProvider } from './copilot.mock';
import { MockDocMeta } from './doc-meta.mock';
import { MockDocSnapshot } from './doc-snapshot.mock';
import { MockDocUser } from './doc-user.mock';
@@ -30,4 +31,10 @@ export const Mockers = {
AccessToken: MockAccessToken,
};
export { MockCopilotProvider, MockEventBus, MockJobQueue, MockMailer };
export {
installMockCopilotRuntime,
MockCopilotProvider,
MockEventBus,
MockJobQueue,
MockMailer,
};
@@ -0,0 +1,110 @@
import { Injectable } from '@nestjs/common';
import { CopilotPromptInvalid } from '../../base';
import { llmGetBuiltInPromptSpec, llmRenderBuiltInPrompt } from '../../native';
import { PromptService } from '../../plugins/copilot/prompt';
import type { Prompt } from '../../plugins/copilot/prompt/spec';
import type {
PromptConfig,
PromptMessage,
} from '../../plugins/copilot/providers/types';
@Injectable()
export class TestingPromptService extends PromptService {
private readonly customPrompts = new Map<string, Prompt>();
private readonly builtInPromptOverrides = new Map<string, Prompt>();
reset() {
this.customPrompts.clear();
this.builtInPromptOverrides.clear();
}
async set(
name: string,
model: string,
messages: PromptMessage[],
config?: PromptConfig | null,
extraConfig?: { optionalModels: string[] }
) {
this.assertCustomPromptName(name);
const existing = this.customPrompts.get(name);
this.customPrompts.set(name, {
name,
model,
action: existing?.action,
optionalModels: existing?.optionalModels?.length
? [...existing.optionalModels, ...(extraConfig?.optionalModels ?? [])]
: extraConfig?.optionalModels,
config: config ? structuredClone(config) : undefined,
messages: this.cloneMessages(messages),
});
}
async overrideBuiltIn(
name: string,
data: {
messages?: PromptMessage[];
model?: string;
config?: PromptConfig | null;
}
) {
const current = this.loadBuiltInPrompt(name);
if (!current) {
throw new CopilotPromptInvalid(
`Built-in prompt ${name} not found in native catalog`
);
}
const { config, messages, model } = data;
const next = this.clonePrompt(current);
if (model !== undefined) {
next.model = model;
}
if (config === null) {
next.config = undefined;
} else if (config !== undefined) {
next.config = structuredClone(config);
}
if (messages) {
next.messages = this.cloneMessages(messages);
}
this.builtInPromptOverrides.set(name, next);
}
protected override lookupCompatPrompt(name: string) {
return (
this.builtInPromptOverrides.get(name) ??
this.customPrompts.get(name) ??
null
);
}
private assertCustomPromptName(name: string) {
if (this.loadBuiltInPrompt(name)) {
throw new CopilotPromptInvalid(
`Built-in prompt ${name} is owned by native catalog`
);
}
}
private loadBuiltInPrompt(name: string): Prompt | null {
const spec = llmGetBuiltInPromptSpec(name);
if (!spec) return null;
const prompt = llmRenderBuiltInPrompt({ name, renderParams: {} });
return {
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels,
config: spec.config,
messages: prompt.messages.map(message => ({
role: message.role,
content: message.content,
...(message.params ? { params: message.params } : {}),
})),
};
}
}
@@ -48,7 +48,9 @@ let docId = 'doc1';
test.beforeEach(async t => {
await t.context.module.initTestingDB();
await t.context.copilotSession.createPrompt('prompt-name', 'gpt-5-mini');
await t.context.db.aiPrompt.create({
data: { name: 'prompt-name', model: 'gpt-5-mini', action: null },
});
user = await t.context.user.create({
email: 'test@affine.pro',
});
@@ -6,6 +6,7 @@ import ava, { ExecutionContext, TestFn } from 'ava';
import { CopilotPromptInvalid, CopilotSessionInvalidInput } from '../../base';
import {
CopilotSessionModel,
Models,
UpdateChatSessionOptions,
UserModel,
WorkspaceModel,
@@ -19,6 +20,7 @@ interface Context {
user: UserModel;
workspace: WorkspaceModel;
copilotSession: CopilotSessionModel;
models: Models;
}
const test = ava as TestFn<Context>;
@@ -28,6 +30,7 @@ test.before(async t => {
t.context.user = module.get(UserModel);
t.context.workspace = module.get(WorkspaceModel);
t.context.copilotSession = module.get(CopilotSessionModel);
t.context.models = module.get(Models);
t.context.db = module.get(PrismaClient);
t.context.module = module;
});
@@ -55,10 +58,12 @@ const TEST_PROMPTS = {
// Helper functions
const createTestPrompts = async (
copilotSession: CopilotSessionModel,
_copilotSession: CopilotSessionModel,
db: PrismaClient
) => {
await copilotSession.createPrompt(TEST_PROMPTS.NORMAL, 'gpt-5-mini');
await db.aiPrompt.create({
data: { name: TEST_PROMPTS.NORMAL, model: 'gpt-5-mini', action: null },
});
await db.aiPrompt.create({
data: { name: TEST_PROMPTS.ACTION, model: 'gpt-5-mini', action: 'edit' },
});
@@ -1000,6 +1005,146 @@ test('should cleanup empty sessions correctly', async t => {
);
});
test('should append durable message and account durable costs', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const { sessionId } = await createTestSession(t);
const appended = await copilotSession.appendMessage({
sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'hello durable world',
params: { foo: 'bar' },
createdAt: new Date(),
},
});
const afterAppend = await db.aiSession.findUniqueOrThrow({
where: { id: sessionId },
select: { messageCost: true, tokenCost: true },
});
t.truthy(appended.id);
t.is(afterAppend.messageCost, 1);
t.true(afterAppend.tokenCost > 0);
t.deepEqual(appended.params, { foo: 'bar' });
const appendedBare = await copilotSession.appendMessage({
sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'assistant',
content: 'assistant reply',
createdAt: new Date(),
},
});
const storedBare = await db.aiSessionMessage.findUniqueOrThrow({
where: { id: appendedBare.id },
select: { params: true },
});
t.deepEqual(appendedBare.params, {});
t.deepEqual(storedBare.params, {});
const oneDayAgo = new Date(Date.now() - 24 * 60 * 60 * 1000);
await db.aiSession.update({
where: { id: sessionId },
data: { updatedAt: oneDayAgo },
});
const cleanup = await copilotSession.cleanupEmptySessions(oneDayAgo);
const persisted = await db.aiSession.findUnique({
where: { id: sessionId },
select: { deletedAt: true, messageCost: true },
});
t.deepEqual(cleanup, { removed: 0, cleaned: 0 });
t.truthy(persisted);
t.is(persisted?.deletedAt, null);
t.is(persisted?.messageCost, 1);
});
test('should count action runs without double-counting legacy action sessions', async t => {
const { copilotSession, db, models } = t.context;
await createTestPrompts(copilotSession, db);
const regular = await createTestSession(t);
await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'regular message',
createdAt: new Date(),
},
});
const legacyAction = await createTestSession(t, {
promptName: TEST_PROMPTS.ACTION,
promptAction: 'edit',
});
const migratedAction = await createTestSession(t, {
promptName: TEST_PROMPTS.ACTION,
promptAction: 'edit',
});
const run = await models.copilotActionRun.create({
userId: user.id,
workspaceId: workspace.id,
sessionId: migratedAction.sessionId,
actionId: 'mindmap.generate',
actionVersion: 'v1',
});
await models.copilotActionRun.complete(run.id, {
status: 'succeeded',
result: { ok: true },
trace: [{ type: 'action_done', status: 'succeeded' }],
});
const retryRun = await models.copilotActionRun.create({
userId: user.id,
workspaceId: workspace.id,
sessionId: migratedAction.sessionId,
actionId: 'mindmap.generate',
actionVersion: 'v1',
attempt: 2,
retryOf: run.id,
});
await models.copilotActionRun.complete(retryRun.id, {
status: 'aborted',
errorCode: 'action_aborted',
trace: [{ type: 'error', status: 'aborted' }],
});
const persistedRetry = await models.copilotActionRun.get(retryRun.id);
const transcriptTask = await models.copilotTranscriptTask.create({
userId: user.id,
workspaceId: workspace.id,
blobId: 'audio-1',
strategy: 'gemini',
recipeId: 'transcript.audio.gemini',
recipeVersion: 'v1',
});
await models.copilotTranscriptTask.complete(transcriptTask.id, {
status: 'ready',
protectedResult: { normalizedTranscript: '00:00:01 A: Hello' },
});
await models.copilotTranscriptTask.settle(transcriptTask.id);
t.like(persistedRetry, {
status: 'aborted',
attempt: 2,
retryOf: run.id,
errorCode: 'action_aborted',
trace: [{ type: 'error', status: 'aborted' }],
});
t.is(await copilotSession.countUserMessages(user.id), 4);
t.truthy(legacyAction.sessionId);
});
test('should get sessions for title generation correctly', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
@@ -1,82 +0,0 @@
import test from 'ava';
import { NativeStreamAdapter } from '../native';
test('NativeStreamAdapter should support buffered and awaited consumption', async t => {
const adapter = new NativeStreamAdapter<number>(undefined);
adapter.push(1);
const first = await adapter.next();
t.deepEqual(first, { value: 1, done: false });
const pending = adapter.next();
adapter.push(2);
const second = await pending;
t.deepEqual(second, { value: 2, done: false });
adapter.push(null);
const done = await adapter.next();
t.true(done.done);
});
test('NativeStreamAdapter return should abort handle and end iteration', async t => {
let abortCount = 0;
const adapter = new NativeStreamAdapter<number>({
abort: () => {
abortCount += 1;
},
});
const ended = await adapter.return();
t.is(abortCount, 1);
t.true(ended.done);
const secondReturn = await adapter.return();
t.true(secondReturn.done);
t.is(abortCount, 1);
const next = await adapter.next();
t.true(next.done);
});
test('NativeStreamAdapter should abort when AbortSignal is triggered', async t => {
let abortCount = 0;
const controller = new AbortController();
const adapter = new NativeStreamAdapter<number>(
{
abort: () => {
abortCount += 1;
},
},
controller.signal
);
const pending = adapter.next();
controller.abort();
const done = await pending;
t.true(done.done);
t.is(abortCount, 1);
});
test('NativeStreamAdapter should end immediately for pre-aborted signal', async t => {
let abortCount = 0;
const controller = new AbortController();
controller.abort();
const adapter = new NativeStreamAdapter<number>(
{
abort: () => {
abortCount += 1;
},
},
controller.signal
);
const next = await adapter.next();
t.true(next.done);
t.is(abortCount, 1);
adapter.push(1);
const stillDone = await adapter.next();
t.true(stillDone.done);
});
File diff suppressed because one or more lines are too long
@@ -1,5 +1,11 @@
import { randomUUID } from 'node:crypto';
import type {
GraphQLQuery,
QueryOptions,
QueryResponse,
} from '@affine/graphql';
import { transformToForm } from '@affine/graphql';
import { INestApplication, ModuleMetadata } from '@nestjs/common';
import type { NestExpressApplication } from '@nestjs/platform-express';
import { TestingModuleBuilder } from '@nestjs/testing';
@@ -188,21 +194,59 @@ export class TestingApp extends ApplyType<INestApplication>() {
// TODO(@forehalo): directly make proxy for graphql queries defined in `@affine/graphql`
// by calling with `app.apis.createWorkspace({ ...variables })`
async gql<Data = any>(query: string, variables?: any): Promise<Data> {
const res = await this.POST('/graphql')
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query,
async gql<Data = any>(query: string, variables?: any): Promise<Data>;
async gql<Query extends GraphQLQuery>(
options: QueryOptions<Query>
): Promise<QueryResponse<Query>>;
async gql<Data = any, Query extends GraphQLQuery = GraphQLQuery>(
queryOrOptions: string | QueryOptions<Query>,
variables?: any
): Promise<Data | QueryResponse<Query>> {
const req = this.POST('/graphql').set({ 'x-request-id': 'test' });
let res: supertest.Response;
if (typeof queryOrOptions === 'string') {
res = await req.set('x-operation-name', 'test').send({
query: queryOrOptions,
variables,
});
} else {
const operationName = queryOrOptions.query.op || 'test';
req.set('x-operation-name', operationName);
if (queryOrOptions.query.file) {
const form = transformToForm({
query: queryOrOptions.query.query,
variables: queryOrOptions.variables,
operationName,
});
for (const [key, value] of form.entries()) {
if (value instanceof File) {
req.attach(key, Buffer.from(await value.arrayBuffer()), {
filename: value.name || key,
contentType: value.type || 'application/octet-stream',
});
} else {
req.field(key, value);
}
}
res = await req;
} else {
res = await req.send({
query: queryOrOptions.query.query,
variables: queryOrOptions.variables,
});
}
}
if (res.status !== 200) {
throw new Error(
`Failed to execute gql: ${query}, status: ${res.status}, body: ${JSON.stringify(
res.body,
null,
2
)}`
`Failed to execute gql: ${
typeof queryOrOptions === 'string'
? queryOrOptions
: queryOrOptions.query.query
}, status: ${res.status}, body: ${JSON.stringify(res.body, null, 2)}`
);
}
@@ -0,0 +1,118 @@
import { Injectable } from '@nestjs/common';
import type { Prisma } from '@prisma/client';
import { Prisma as PrismaClient } from '@prisma/client';
import { BaseModel } from './base';
export type AiActionRunStatus =
| 'created'
| 'running'
| 'succeeded'
| 'failed'
| 'aborted';
function nullableJson(
value: unknown
): Prisma.NullableJsonNullValueInput | Prisma.InputJsonValue {
return value === undefined
? PrismaClient.JsonNull
: (value as Prisma.InputJsonValue);
}
@Injectable()
export class CopilotActionRunModel extends BaseModel {
async create(
input: Pick<
Prisma.AiActionRunCreateArgs['data'],
'userId' | 'workspaceId' | 'actionId' | 'actionVersion'
> & { inputSnapshot?: unknown } & Omit<
Partial<Prisma.AiActionRunCreateArgs['data']>,
'inputSnapshot'
>
) {
return await this.db.aiActionRun.create({
data: {
userId: input.userId,
workspaceId: input.workspaceId,
docId: input.docId ?? null,
sessionId: input.sessionId ?? null,
userMessageId: input.userMessageId ?? null,
compatSubmissionId: input.compatSubmissionId ?? null,
actionId: input.actionId,
actionVersion: input.actionVersion,
status: 'created',
attempt: input.attempt ?? 1,
retryOf: input.retryOf ?? null,
inputSnapshot: nullableJson(input.inputSnapshot),
},
});
}
async markRunning(id: string) {
return await this.db.aiActionRun.update({
where: { id },
data: { status: 'running' },
});
}
async complete(
id: string,
input: Omit<
Prisma.AiActionRunUpdateArgs['data'],
'artifacts' | 'result' | 'trace'
> & {
result?: unknown;
artifacts?: unknown;
trace?: unknown;
}
) {
return await this.db.aiActionRun.update({
where: { id },
data: {
status: input.status,
result: nullableJson(input.result),
artifacts: nullableJson(input.artifacts),
resultSummary: input.resultSummary ?? null,
errorCode: input.errorCode ?? null,
trace: nullableJson(input.trace),
assistantMessageId: input.assistantMessageId ?? null,
},
});
}
async get(id: string) {
const row = await this.db.aiActionRun.findUnique({ where: { id } });
return row ?? null;
}
async countSucceededByUser(userId: string) {
return await this.db.aiActionRun.count({
where: {
userId,
status: 'succeeded',
NOT: {
actionId: {
startsWith: 'transcript.audio.',
},
},
},
});
}
async countLegacyPromptActionSessionsWithoutRun(userId: string) {
return await this.db.aiSession.count({
where: {
userId,
promptAction: {
not: null,
},
NOT: {
promptAction: '',
},
actionRuns: {
none: {},
},
},
});
}
}
@@ -11,6 +11,10 @@ import {
} from '../base';
import { getTokenEncoder } from '../native';
import type { PromptAttachment } from '../plugins/copilot/providers/types';
import {
type ChatMessage as CopilotChatMessage,
ChatMessageSchema,
} from '../plugins/copilot/types';
import { BaseModel } from './base';
export enum SessionType {
@@ -34,10 +38,14 @@ type ChatStreamObject = {
toolName?: string;
args?: Record<string, any>;
result?: any;
rawArgumentsText?: string;
argumentParseError?: string;
thought?: string;
};
type ChatMessage = {
id?: string | undefined;
compatSubmissionId?: string | null;
role: 'system' | 'assistant' | 'user';
content: string;
attachments?: ChatAttachment[] | null;
@@ -46,6 +54,19 @@ type ChatMessage = {
createdAt: Date;
};
type StoredChatMessage = Prisma.AiSessionMessageGetPayload<{
select: {
id: true;
compatSubmissionId: true;
role: true;
content: true;
attachments: true;
streamObjects: true;
params: true;
createdAt: true;
};
}>;
type PureChatSession = {
sessionId: string;
workspaceId: string;
@@ -84,7 +105,10 @@ type UpdateChatSessionMessage = ChatSessionBaseState & {
};
export type UpdateChatSessionOptions = ChatSessionBaseState &
Pick<Partial<ChatSession>, 'docId' | 'pinned' | 'promptName' | 'title'>;
Pick<
Partial<ChatSession>,
'docId' | 'pinned' | 'promptName' | 'promptAction' | 'title'
> & { promptModel?: string };
export type UpdateChatSession = ChatSessionBaseState & UpdateChatSessionOptions;
@@ -114,6 +138,26 @@ export type CleanupSessionOptions = Pick<
@Injectable()
export class CopilotSessionModel extends BaseModel {
private noActionPromptCondition(): Prisma.AiSessionWhereInput {
return {
OR: [{ promptAction: null }, { promptAction: '' }],
};
}
private async ensurePromptCompatRecord(prompt: ChatPrompt) {
await this.db.aiPrompt.upsert({
where: { name: prompt.name },
update: {},
create: {
name: prompt.name,
action: prompt.action,
model: prompt.model,
optionalModels: [],
config: {},
},
});
}
private sanitizeString<T extends string | null | undefined>(value: T): T {
if (typeof value !== 'string') {
return value;
@@ -154,6 +198,9 @@ export class CopilotSessionModel extends BaseModel {
toolCallId: this.sanitizeString(stream.toolCallId) ?? '',
toolName: this.sanitizeString(stream.toolName) ?? '',
args: this.sanitizeJsonValue(stream.args),
rawArgumentsText: this.sanitizeString(stream.rawArgumentsText),
argumentParseError: this.sanitizeString(stream.argumentParseError),
thought: this.sanitizeString(stream.thought),
};
case 'tool-result':
return {
@@ -162,6 +209,8 @@ export class CopilotSessionModel extends BaseModel {
toolName: this.sanitizeString(stream.toolName) ?? '',
args: this.sanitizeJsonValue(stream.args),
result: this.sanitizeJsonValue(stream.result),
rawArgumentsText: this.sanitizeString(stream.rawArgumentsText),
argumentParseError: this.sanitizeString(stream.argumentParseError),
};
}
}
@@ -279,6 +328,7 @@ export class CopilotSessionModel extends BaseModel {
private sanitizeMessage(message: ChatMessage): ChatMessage {
return {
...message,
compatSubmissionId: this.sanitizeString(message.compatSubmissionId),
content: this.sanitizeString(message.content) ?? '',
attachments: this.sanitizeAttachments(message.attachments),
params: this.sanitizeJsonValue(
@@ -290,6 +340,23 @@ export class CopilotSessionModel extends BaseModel {
};
}
private toPublicMessage(message: StoredChatMessage): CopilotChatMessage {
const { compatSubmissionId: _compatSubmissionId, ...publicMessage } =
message;
return ChatMessageSchema.parse({
...publicMessage,
attachments: publicMessage.attachments ?? undefined,
streamObjects: publicMessage.streamObjects ?? undefined,
params: publicMessage.params ?? undefined,
});
}
private isCountedUserMessage(
message: Pick<StoredChatMessage, 'role'>
): boolean {
return message.role === AiPromptRole.user;
}
getSessionType(session: Pick<ChatSession, 'docId' | 'pinned'>): SessionType {
if (session.pinned) return SessionType.Pinned;
if (!session.docId) return SessionType.Workspace;
@@ -316,13 +383,6 @@ export class CopilotSessionModel extends BaseModel {
return true;
}
// NOTE: just for test, remove it after copilot prompt model is ready
async createPrompt(name: string, model: string, action?: string) {
await this.db.aiPrompt.create({
data: { name, model, action: action ?? null },
});
}
@Transactional()
async create(state: ChatSession, reuseChat = false): Promise<string> {
// find and return existing session if session is chat session
@@ -358,6 +418,7 @@ export class CopilotSessionModel extends BaseModel {
reuseChat = false
): Promise<string> {
const { prompt, ...rest } = state;
await this.ensurePromptCompatRecord(prompt);
return await this.models.copilotSession.create(
{ ...rest, promptName: prompt.name, promptAction: prompt.action ?? null },
reuseChat
@@ -414,7 +475,7 @@ export class CopilotSessionModel extends BaseModel {
workspaceId: state.workspaceId,
docId: state.docId,
parentSessionId: null,
prompt: { action: { equals: null } },
...this.noActionPromptCondition(),
...extraCondition,
},
select: { id: true, deletedAt: true },
@@ -464,22 +525,28 @@ export class CopilotSessionModel extends BaseModel {
});
}
@Transactional()
async getMeta(sessionId: string) {
return await this.getExists(sessionId, {
id: true,
userId: true,
workspaceId: true,
docId: true,
parentSessionId: true,
pinned: true,
title: true,
promptName: true,
tokenCost: true,
createdAt: true,
updatedAt: true,
});
}
private getListConditions(
options: ListSessionOptions
): Prisma.AiSessionWhereInput {
const { userId, sessionId, workspaceId, docId, action, fork } = options;
function getNullCond<T>(
maybeBool: boolean | undefined,
wrap: (ret: { not: null } | null) => T = ret => ret as T
): T | undefined {
return maybeBool === true
? wrap({ not: null })
: maybeBool === false
? wrap(null)
: undefined;
}
function getEqCond<T>(maybeValue: T | undefined): T | undefined {
return maybeValue !== undefined ? maybeValue : undefined;
}
@@ -492,8 +559,13 @@ export class CopilotSessionModel extends BaseModel {
id: getEqCond(sessionId),
deletedAt: null,
pinned: getEqCond(options.pinned),
prompt: getNullCond(action, ret => ({ action: ret })),
parentSessionId: getNullCond(fork),
...(action === false ? this.noActionPromptCondition() : {}),
...(action === true ? { NOT: this.noActionPromptCondition() } : {}),
...(fork === true
? { parentSessionId: { not: null } }
: fork === false
? { parentSessionId: null }
: {}),
},
];
@@ -505,7 +577,7 @@ export class CopilotSessionModel extends BaseModel {
workspaceId: workspaceId,
docId: docId ?? null,
id: getEqCond(sessionId),
prompt: { action: null },
...this.noActionPromptCondition(),
// should only find forked session
parentSessionId: { not: null },
deletedAt: null,
@@ -587,7 +659,7 @@ export class CopilotSessionModel extends BaseModel {
docId: true,
parentSessionId: true,
pinned: true,
prompt: true,
promptAction: true,
},
{ userId }
);
@@ -597,7 +669,7 @@ export class CopilotSessionModel extends BaseModel {
// not allow to update action session
if (!internalCall) {
if (session.prompt.action) {
if (session.promptAction) {
throw new CopilotSessionInvalidInput(
`Cannot update action: ${session.id}`
);
@@ -608,12 +680,29 @@ export class CopilotSessionModel extends BaseModel {
}
}
let nextPromptAction: string | null | undefined;
if (promptName) {
const prompt = await this.db.aiPrompt.findFirst({
where: { name: promptName },
});
// always not allow to update to action prompt
if (!prompt || prompt.action) {
if (options.promptModel) {
await this.ensurePromptCompatRecord({
name: promptName,
action: options.promptAction,
model: options.promptModel,
});
}
nextPromptAction = options.promptAction;
if (nextPromptAction === undefined) {
const prompt = await this.db.aiPrompt.findFirst({
where: { name: promptName },
select: { action: true },
});
if (!prompt) {
throw new CopilotSessionInvalidInput(
`Prompt ${promptName} not found or not available for session ${sessionId}`
);
}
nextPromptAction = prompt.action ?? null;
}
if (nextPromptAction) {
throw new CopilotSessionInvalidInput(
`Prompt ${promptName} not found or not available for session ${sessionId}`
);
@@ -626,7 +715,13 @@ export class CopilotSessionModel extends BaseModel {
await this.db.aiSession.update({
where: { id: sessionId },
data: { docId, promptName, pinned, title: sanitizedTitle },
data: {
docId,
promptName,
promptAction: nextPromptAction,
pinned,
title: sanitizedTitle,
},
});
return sessionId;
@@ -672,6 +767,48 @@ export class CopilotSessionModel extends BaseModel {
});
}
@Transactional()
async getMessage(sessionId: string, messageId: string) {
const message = await this.db.aiSessionMessage.findFirst({
where: { id: messageId, sessionId },
select: {
id: true,
compatSubmissionId: true,
role: true,
content: true,
attachments: true,
streamObjects: true,
params: true,
createdAt: true,
},
});
return message ? this.toPublicMessage(message) : null;
}
@Transactional()
async findMessageByCompatSubmissionId(
sessionId: string,
compatSubmissionId: string
) {
const message = await this.db.aiSessionMessage.findFirst({
where: { sessionId, compatSubmissionId },
select: {
id: true,
compatSubmissionId: true,
role: true,
content: true,
attachments: true,
streamObjects: true,
params: true,
createdAt: true,
},
orderBy: { createdAt: 'asc' },
});
return message ? this.toPublicMessage(message) : null;
}
private calculateTokenSize(messages: any[], model: string): number {
const encoder = getTokenEncoder(model);
const content = messages.map(m => m.content).join('');
@@ -694,10 +831,13 @@ export class CopilotSessionModel extends BaseModel {
);
await this.db.aiSessionMessage.createMany({
data: sanitizedMessages.map(m => ({
...m,
compatSubmissionId: m.compatSubmissionId || undefined,
role: m.role,
content: m.content,
attachments: m.attachments || undefined,
params: m.params || undefined,
streamObjects: m.streamObjects || undefined,
createdAt: m.createdAt,
sessionId,
})),
});
@@ -714,18 +854,120 @@ export class CopilotSessionModel extends BaseModel {
}
}
@Transactional()
async appendMessage(state: {
sessionId: string;
userId: string;
prompt: { model: string };
message: ChatMessage;
}) {
const haveSession = await this.has(state.sessionId, state.userId);
if (!haveSession) {
throw new CopilotSessionNotFound();
}
const message = this.sanitizeMessage(state.message);
const tokenCost = this.calculateTokenSize([message], state.prompt.model);
const created = await this.db.aiSessionMessage.create({
data: {
sessionId: state.sessionId,
compatSubmissionId: message.compatSubmissionId || undefined,
role: message.role,
content: message.content,
attachments: message.attachments || undefined,
params: message.params || undefined,
streamObjects: message.streamObjects || undefined,
createdAt: message.createdAt,
},
select: {
id: true,
compatSubmissionId: true,
role: true,
content: true,
attachments: true,
streamObjects: true,
params: true,
createdAt: true,
},
});
await this.db.aiSession.update({
where: { id: state.sessionId },
data: {
messageCost:
message.role === AiPromptRole.user ? { increment: 1 } : undefined,
tokenCost: { increment: tokenCost },
},
});
return this.toPublicMessage(created);
}
@Transactional()
async trimAfterMessage(
sessionId: string,
messageId: string,
removeTargetMessage = false
) {
const session = await this.getExists(sessionId, {
id: true,
});
if (!session) {
throw new CopilotSessionNotFound();
}
const messages = await this.getMessages(
sessionId,
{ id: true, role: true, content: true, params: true },
{ createdAt: 'asc' }
);
const messageIndex = messages.findIndex(({ id }) => id === messageId);
if (messageIndex < 0) {
throw new CopilotSessionNotFound();
}
const ids = messages
.slice(messageIndex + (removeTargetMessage ? 0 : 1))
.map(({ id }) => id);
if (!ids.length) {
return;
}
await this.db.aiSessionMessage.deleteMany({ where: { id: { in: ids } } });
const remainingMessages = await this.getMessages(sessionId, {
role: true,
});
const userMessageCount = remainingMessages.filter(message =>
this.isCountedUserMessage(message)
).length;
if (userMessageCount <= 1) {
await this.db.aiSession.update({
where: { id: sessionId },
data: { title: null },
});
}
}
@Transactional()
async revertLatestMessage(
sessionId: string,
removeLatestUserMessage: boolean
) {
const id = await this.getExists(sessionId, { id: true }).then(
session => session?.id
);
if (!id) {
const session = await this.getExists(sessionId, {
id: true,
});
if (!session) {
throw new CopilotSessionNotFound();
}
const messages = await this.getMessages(id, { id: true, role: true });
const messages = await this.getMessages(session.id, {
id: true,
role: true,
content: true,
});
const ids = messages
.slice(
messages.findLastIndex(({ role }) => role === AiPromptRole.user) +
@@ -737,14 +979,16 @@ export class CopilotSessionModel extends BaseModel {
await this.db.aiSessionMessage.deleteMany({ where: { id: { in: ids } } });
// clear the title if there only one round of conversation left
const remainingMessages = await this.getMessages(id, { role: true });
const userMessageCount = remainingMessages.filter(
m => m.role === AiPromptRole.user
const remainingMessages = await this.getMessages(session.id, {
role: true,
});
const userMessageCount = remainingMessages.filter(message =>
this.isCountedUserMessage(message)
).length;
if (userMessageCount <= 1) {
await this.db.aiSession.update({
where: { id },
where: { id: session.id },
data: { title: null },
});
}
@@ -755,11 +999,26 @@ export class CopilotSessionModel extends BaseModel {
async countUserMessages(userId: string): Promise<number> {
const sessions = await this.db.aiSession.findMany({
where: { userId },
select: { messageCost: true, prompt: { select: { action: true } } },
select: { messageCost: true, promptAction: true },
});
return sessions
.map(({ messageCost, prompt: { action } }) => (action ? 1 : messageCost))
const regularMessageCost = sessions
.filter(({ promptAction }) => !promptAction)
.map(({ messageCost }) => messageCost)
.reduce((prev, cost) => prev + cost, 0);
const [actionRunCost, legacyActionSessionCost, transcriptSettlementCost] =
await Promise.all([
this.models.copilotActionRun.countSucceededByUser(userId),
this.models.copilotActionRun.countLegacyPromptActionSessionsWithoutRun(
userId
),
this.models.copilotTranscriptTask.countSettledByUser(userId),
]);
return (
regularMessageCost +
actionRunCost +
legacyActionSessionCost +
transcriptSettlementCost
);
}
async cleanupEmptySessions(earlyThen: Date) {
@@ -799,7 +1058,7 @@ export class CopilotSessionModel extends BaseModel {
deletedAt: null,
messages: { some: {} },
// only generate titles for non-actions sessions
prompt: { action: null },
...this.noActionPromptCondition(),
},
select: {
id: true,
@@ -0,0 +1,124 @@
import { Injectable } from '@nestjs/common';
import type { Prisma } from '@prisma/client';
import { Prisma as PrismaClient } from '@prisma/client';
import { BaseModel } from './base';
function nullableJson(
value: unknown
): Prisma.NullableJsonNullValueInput | Prisma.InputJsonValue {
return value === undefined
? PrismaClient.JsonNull
: (value as Prisma.InputJsonValue);
}
function isRecordNotFound(error: unknown) {
return (
error instanceof PrismaClient.PrismaClientKnownRequestError &&
error.code === 'P2025'
);
}
@Injectable()
export class CopilotTranscriptTaskModel extends BaseModel {
async create(
input: Pick<
Prisma.AiTranscriptTaskCreateArgs['data'],
| 'userId'
| 'workspaceId'
| 'blobId'
| 'strategy'
| 'recipeId'
| 'recipeVersion'
> &
Partial<Prisma.AiTranscriptTaskCreateArgs['data']>
) {
return await this.db.aiTranscriptTask.create({
data: {
userId: input.userId,
workspaceId: input.workspaceId,
blobId: input.blobId,
status: 'pending',
strategy: input.strategy,
recipeId: input.recipeId,
recipeVersion: input.recipeVersion,
inputSnapshot: nullableJson(input.inputSnapshot),
publicMeta: nullableJson(input.publicMeta),
},
});
}
async get(id: string) {
const row = await this.db.aiTranscriptTask.findUnique({ where: { id } });
return row ?? null;
}
async getWithUser(
userId: string,
workspaceId: string,
taskId?: string,
blobId?: string
) {
if (!taskId && !blobId) return null;
const row = await this.db.aiTranscriptTask.findFirst({
where: {
userId,
workspaceId,
...(taskId ? { id: taskId } : {}),
...(blobId ? { blobId } : {}),
},
orderBy: { createdAt: 'desc' },
});
return row ?? null;
}
async markRunning(id: string, actionRunId?: string | null) {
try {
return await this.db.aiTranscriptTask.update({
where: { id },
data: {
status: 'running',
...(actionRunId ? { actionRunId } : {}),
errorCode: null,
},
});
} catch (error) {
if (isRecordNotFound(error)) return null;
throw error;
}
}
async complete(id: string, input: Prisma.AiTranscriptTaskUpdateArgs['data']) {
try {
return await this.db.aiTranscriptTask.update({
where: { id },
data: {
status: input.status,
...(input.actionRunId ? { actionRunId: input.actionRunId } : {}),
publicMeta: nullableJson(input.publicMeta),
protectedResult: nullableJson(input.protectedResult),
errorCode: input.errorCode ?? null,
},
});
} catch (error) {
if (isRecordNotFound(error)) return null;
throw error;
}
}
async settle(id: string) {
const task = await this.get(id);
if (!task) return null;
return await this.db.aiTranscriptTask.update({
where: { id },
data: { status: 'settled', settledAt: task.settledAt ?? new Date() },
});
}
async countSettledByUser(userId: string) {
return await this.db.aiTranscriptTask.count({
where: { userId, status: 'settled' },
});
}
}
@@ -16,9 +16,11 @@ import { CalendarSubscriptionModel } from './calendar-subscription';
import { CommentModel } from './comment';
import { CommentAttachmentModel } from './comment-attachment';
import { AppConfigModel } from './config';
import { CopilotActionRunModel } from './copilot-action-run';
import { CopilotContextModel } from './copilot-context';
import { CopilotJobModel } from './copilot-job';
import { CopilotSessionModel } from './copilot-session';
import { CopilotTranscriptTaskModel } from './copilot-transcript-task';
import { CopilotWorkspaceConfigModel } from './copilot-workspace';
import { DocModel } from './doc';
import { DocUserModel } from './doc-user';
@@ -56,6 +58,8 @@ const MODELS = {
notification: NotificationModel,
userSettings: UserSettingsModel,
copilotSession: CopilotSessionModel,
copilotTranscriptTask: CopilotTranscriptTaskModel,
copilotActionRun: CopilotActionRunModel,
copilotContext: CopilotContextModel,
copilotWorkspace: CopilotWorkspaceConfigModel,
copilotJob: CopilotJobModel,
@@ -132,6 +136,7 @@ export * from './common';
export * from './copilot-context';
export * from './copilot-job';
export * from './copilot-session';
export * from './copilot-transcript-task';
export * from './copilot-workspace';
export * from './doc';
export * from './doc-user';
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,16 @@
import { Injectable } from '@nestjs/common';
import { promptAttachmentToUrl } from '../providers/utils';
import type { ChatMessage } from '../types';
@Injectable()
export class HistoryAttachmentUrlProjector {
projectMessages(messages: ChatMessage[]): ChatMessage[] {
return messages.map(message => ({
...message,
attachments: message.attachments
?.map(attachment => promptAttachmentToUrl(attachment))
.filter((attachment): attachment is string => !!attachment),
}));
}
}
@@ -0,0 +1,99 @@
import { Injectable } from '@nestjs/common';
import { AiPromptRole } from '@prisma/client';
import type { Conversation, Turn } from '../core';
import { chatMessageFromTurn } from '../core';
import type { ResolvedPrompt } from '../prompt';
import { type ChatHistory } from '../types';
import { HistoryAttachmentUrlProjector } from './history-attachment-url-projector';
import { HistoryPromptPreloadProjector } from './history-prompt-preload-projector';
import {
HistoryVisibilityPolicy,
type ProjectConversationOptions,
} from './history-visibility-policy';
export type CanonicalConversationHistory = {
conversation: Conversation;
turns: Turn[];
prompt: ResolvedPrompt;
tokenCost: number;
};
export type CanonicalConversationMeta = Omit<
CanonicalConversationHistory,
'turns'
>;
@Injectable()
export class CompatHistoryProjector {
constructor(
private readonly visibility: HistoryVisibilityPolicy,
private readonly preloadProjector: HistoryPromptPreloadProjector,
private readonly attachmentUrls: HistoryAttachmentUrlProjector
) {}
private projectSessionBase(
history: CanonicalConversationMeta
): Omit<ChatHistory, 'messages'> {
const { conversation, prompt, tokenCost } = history;
return {
userId: conversation.userId,
sessionId: conversation.id,
workspaceId: conversation.workspaceId,
docId: conversation.docId,
parentSessionId: conversation.parentId,
pinned: conversation.pinned,
title: conversation.title,
action: prompt.action || null,
model: prompt.model,
optionalModels: prompt.optionalModels || [],
promptName: prompt.name,
tokens: tokenCost,
createdAt: conversation.createdAt,
updatedAt: conversation.updatedAt,
};
}
projectSession(
history: CanonicalConversationMeta,
_options: ProjectConversationOptions
): Omit<ChatHistory, 'messages'> | undefined {
return this.projectSessionBase(history);
}
projectHistory(
history: CanonicalConversationHistory,
options: ProjectConversationOptions & {
withMessages: boolean;
withPrompt?: boolean;
}
): ChatHistory | undefined {
if (!this.visibility.shouldExposeHistory(history, options)) return;
const base = this.projectSessionBase(history);
const { turns } = history;
const messages = turns.map(turn => chatMessageFromTurn(turn));
const preload = this.preloadProjector.project(
history,
options.withMessages,
options.withPrompt
);
const projectedMessages = options.withMessages
? preload
.concat(messages)
.filter(
message =>
message.role !== AiPromptRole.user ||
!!message.content.trim() ||
!!message.attachments?.length
)
.map(message => ({ ...message }))
: [];
return {
...base,
messages: this.attachmentUrls.projectMessages(projectedMessages),
};
}
}
@@ -0,0 +1,37 @@
import { Injectable } from '@nestjs/common';
import { AiPromptRole } from '@prisma/client';
import { PromptService } from '../prompt/service';
import type { ChatMessage } from '../types';
import type { CanonicalConversationHistory } from './history-projector';
@Injectable()
export class HistoryPromptPreloadProjector {
constructor(private readonly prompts: PromptService) {}
project(
history: CanonicalConversationHistory,
withMessages: boolean,
withPrompt?: boolean
): ChatMessage[] {
if (!withMessages || !withPrompt) {
return [];
}
const preload = this.prompts
.finish(
history.prompt,
history.turns[0] ? history.turns[0].metadata : {},
history.conversation.id
)
.filter(({ role }) => role !== AiPromptRole.system) as ChatMessage[];
preload.forEach((message, index) => {
message.createdAt = new Date(
history.conversation.createdAt.getTime() - preload.length - index - 1
);
});
return preload;
}
}
@@ -0,0 +1,28 @@
import { Injectable } from '@nestjs/common';
import type { CanonicalConversationHistory } from './history-projector';
export type ProjectConversationOptions = {
requestUserId: string | undefined;
action?: boolean;
skipVisibilityFilter?: boolean;
};
@Injectable()
export class HistoryVisibilityPolicy {
shouldExposeHistory(
history: CanonicalConversationHistory,
options: ProjectConversationOptions
): boolean {
if (options.skipVisibilityFilter) {
return true;
}
return !(
(history.conversation.userId === options.requestUserId &&
!!options.action !== !!history.prompt.action) ||
(history.conversation.userId !== options.requestUserId &&
!!history.prompt.action)
);
}
}
@@ -0,0 +1,118 @@
import { randomUUID } from 'node:crypto';
import { Injectable } from '@nestjs/common';
import { Cache } from '../../../base';
import type { PromptMessage } from '../providers/types';
const SUBMISSION_TTL = 24 * 60 * 60 * 1000;
type StoredCompatSubmission = {
id: string;
sessionId: string;
content?: string;
attachments?: PromptMessage['attachments'];
params?: Record<string, any>;
createdAt: string;
};
type StoredAcceptedSubmission = {
sessionId: string;
turnId: string;
acceptedAt: string;
};
export type CompatSubmission = Omit<StoredCompatSubmission, 'createdAt'> & {
createdAt: Date;
};
export type AcceptedCompatSubmission = Omit<
StoredAcceptedSubmission,
'acceptedAt'
> & {
acceptedAt: Date;
};
@Injectable()
export class CompatSubmissionStore {
constructor(private readonly cache: Cache) {}
private submissionKey(token: string) {
return `copilot:submission:${token}`;
}
private acceptedKey(token: string) {
return `copilot:submission:${token}:accepted`;
}
private fromStoredSubmission(
submission?: StoredCompatSubmission
): CompatSubmission | undefined {
if (!submission) {
return;
}
return {
...submission,
createdAt: new Date(submission.createdAt),
};
}
private fromStoredAccepted(
accepted?: StoredAcceptedSubmission
): AcceptedCompatSubmission | undefined {
if (!accepted) {
return;
}
return {
...accepted,
acceptedAt: new Date(accepted.acceptedAt),
};
}
async create(
submission: Omit<CompatSubmission, 'id' | 'createdAt'>
): Promise<string> {
const token = randomUUID();
const stored: StoredCompatSubmission = {
...submission,
id: token,
createdAt: new Date().toISOString(),
};
await this.cache.set(this.submissionKey(token), stored, {
ttl: SUBMISSION_TTL,
});
return token;
}
async get(token: string): Promise<CompatSubmission | undefined> {
return this.fromStoredSubmission(
await this.cache.get<StoredCompatSubmission>(this.submissionKey(token))
);
}
async markAccepted(
token: string,
accepted: { sessionId: string; turnId: string }
) {
await this.cache.set<StoredAcceptedSubmission>(
this.acceptedKey(token),
{
...accepted,
acceptedAt: new Date().toISOString(),
},
{ ttl: SUBMISSION_TTL }
);
await this.cache.delete(this.submissionKey(token));
}
async getAccepted(
token: string
): Promise<AcceptedCompatSubmission | undefined> {
return this.fromStoredAccepted(
await this.cache.get<StoredAcceptedSubmission>(this.acceptedKey(token))
);
}
}
@@ -5,7 +5,6 @@ import {
StorageJSONSchema,
StorageProviderConfig,
} from '../../base';
import { CopilotPromptScenario } from './prompt/prompts';
import {
AnthropicOfficialConfig,
AnthropicVertexConfig,
@@ -41,6 +40,7 @@ export const RustRequestMiddlewareValues = [
'normalize_messages',
'clamp_max_tokens',
'tool_schema_rewrite',
'openai_request_compat',
] as const;
export type RustRequestMiddleware =
(typeof RustRequestMiddlewareValues)[number];
@@ -83,7 +83,7 @@ export type CopilotProviderProfile = CopilotProviderProfileCommon &
}[CopilotProviderType];
export type CopilotProviderDefaults = Partial<
Record<Exclude<ModelOutputType, ModelOutputType.Rerank>, string>
Record<Exclude<ModelOutputType, typeof ModelOutputType.Rerank>, string>
> & {
fallback?: string;
};
@@ -212,7 +212,6 @@ declare global {
key: string;
}>;
storage: ConfigItem<StorageProviderConfig>;
scenarios: ConfigItem<CopilotPromptScenario>;
providers: {
profiles: ConfigItem<CopilotProviderProfile[]>;
defaults: ConfigItem<CopilotProviderDefaults>;
@@ -235,23 +234,6 @@ defineModuleConfig('copilot', {
desc: 'Whether to enable the copilot plugin. <br> Document: <a href="https://docs.affine.pro/self-host-affine/administer/ai" target="_blank">https://docs.affine.pro/self-host-affine/administer/ai</a>',
default: false,
},
scenarios: {
desc: 'Use custom models in scenarios and override default settings.',
default: {
override_enabled: false,
scenarios: {
audio_transcribing: 'gemini-2.5-flash',
chat: 'gemini-2.5-flash',
embedding: 'gemini-embedding-001',
image: 'gpt-image-1',
coding: 'claude-sonnet-4-5@20250929',
complex_text_generation: 'gpt-5-mini',
quick_decision_making: 'gpt-5-mini',
quick_text_generation: 'gemini-2.5-flash',
polish_and_summarize: 'gemini-2.5-flash',
},
},
},
'providers.profiles': {
desc: 'The profile list for copilot providers.',
default: [],
@@ -750,14 +750,17 @@ export class CopilotContextResolver {
sniffMime(buffer, mimetype) || mimetype
);
await this.jobs.addFileEmbeddingQueue({
userId: user.id,
workspaceId: session.workspaceId,
contextId: session.id,
blobId: file.blobId,
fileId: file.id,
fileName: file.name,
});
await this.jobs.addFileEmbeddingQueue(
{
userId: user.id,
workspaceId: session.workspaceId,
contextId: session.id,
blobId: file.blobId,
fileId: file.id,
fileName: file.name,
},
{ priority: 0 }
);
return file;
} catch (e: any) {
@@ -1,5 +1,4 @@
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import {
Cache,
@@ -15,7 +14,7 @@ import {
ContextFile,
Models,
} from '../../../models';
import { getEmbeddingClient } from '../embedding/client';
import { CopilotEmbeddingClientService } from '../embedding/client';
import type { EmbeddingClient } from '../embedding/types';
import { ContextSession } from './session';
@@ -27,7 +26,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
private client: EmbeddingClient | undefined;
constructor(
private readonly moduleRef: ModuleRef,
private readonly embeddingClients: CopilotEmbeddingClientService,
private readonly cache: Cache,
private readonly models: Models
) {}
@@ -43,7 +42,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
}
private async setup() {
this.client = await getEmbeddingClient(this.moduleRef);
this.client = await this.embeddingClients.refresh();
}
async onApplicationBootstrap() {
@@ -59,8 +58,8 @@ export class CopilotContextService implements OnApplicationBootstrap {
}
// public this client to allow overriding in tests
get embeddingClient() {
return this.client as EmbeddingClient;
get embeddingClient(): EmbeddingClient | undefined {
return this.client ?? this.embeddingClients.getClient();
}
private async saveConfig(
@@ -175,8 +174,9 @@ export class CopilotContextService implements OnApplicationBootstrap {
signal?: AbortSignal,
threshold: number = 0.5
) {
if (!this.embeddingClient) return [];
const embedding = await this.embeddingClient.getEmbedding(content, signal);
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
if (!embedding) return [];
const blobChunks = await this.models.copilotWorkspace.matchBlobEmbedding(
@@ -187,7 +187,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
);
if (!blobChunks.length) return [];
return await this.embeddingClient.reRank(content, blobChunks, topK, signal);
return await client.reRank(content, blobChunks, topK, signal);
}
async matchWorkspaceFiles(
@@ -197,8 +197,9 @@ export class CopilotContextService implements OnApplicationBootstrap {
signal?: AbortSignal,
threshold: number = 0.5
) {
if (!this.embeddingClient) return [];
const embedding = await this.embeddingClient.getEmbedding(content, signal);
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
if (!embedding) return [];
const fileChunks = await this.models.copilotWorkspace.matchFileEmbedding(
@@ -209,7 +210,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
);
if (!fileChunks.length) return [];
return await this.embeddingClient.reRank(content, fileChunks, topK, signal);
return await client.reRank(content, fileChunks, topK, signal);
}
async matchWorkspaceDocs(
@@ -219,8 +220,9 @@ export class CopilotContextService implements OnApplicationBootstrap {
signal?: AbortSignal,
threshold: number = 0.5
) {
if (!this.embeddingClient) return [];
const embedding = await this.embeddingClient.getEmbedding(content, signal);
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
if (!embedding) return [];
const workspaceChunks =
@@ -232,12 +234,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
);
if (!workspaceChunks.length) return [];
return await this.embeddingClient.reRank(
content,
workspaceChunks,
topK,
signal
);
return await client.reRank(content, workspaceChunks, topK, signal);
}
async matchWorkspaceAll(
@@ -249,8 +246,9 @@ export class CopilotContextService implements OnApplicationBootstrap {
docIds?: string[],
scopedThreshold: number = 0.85
) {
if (!this.embeddingClient) return [];
const embedding = await this.embeddingClient.getEmbedding(content, signal);
const client = this.embeddingClient;
if (!client) return [];
const embedding = await client.getEmbedding(content, signal);
if (!embedding) return [];
const [fileChunks, blobChunks, workspaceChunks, scopedWorkspaceChunks] =
@@ -293,7 +291,7 @@ export class CopilotContextService implements OnApplicationBootstrap {
return [];
}
return await this.embeddingClient.reRank(
return await client.reRank(
content,
[
...fileChunks,
@@ -318,6 +316,18 @@ export class CopilotContextService implements OnApplicationBootstrap {
}));
}
@OnEvent('workspace.doc.embed.finished')
async onDocEmbedFinished({
contextId,
docId,
}: Events['workspace.doc.embed.finished']) {
const context = await this.get(contextId);
await context.saveDocRecord(docId, doc => ({
...(doc as ContextDoc),
status: ContextEmbedStatus.finished,
}));
}
@OnEvent('workspace.file.embed.finished')
async onFileEmbedFinish({
contextId,
@@ -13,22 +13,17 @@ import type { Request, Response } from 'express';
import {
BehaviorSubject,
catchError,
connect,
filter,
finalize,
from,
ignoreElements,
interval,
lastValueFrom,
map,
merge,
mergeMap,
Observable,
reduce,
Subject,
take,
takeUntil,
tap,
} from 'rxjs';
import {
@@ -36,28 +31,18 @@ import {
BlobNotFound,
CallMetric,
Config,
CopilotSessionNotFound,
mapSseError,
metrics,
NoCopilotProviderAvailable,
UnsplashIsNotConfigured,
} from '../../base';
import { ServerFeature, ServerService } from '../../core';
import { CurrentUser, Public } from '../../core/auth';
import { CopilotContextService } from './context/service';
import { CopilotProviderFactory } from './providers/factory';
import type { CopilotProvider } from './providers/provider';
import {
ModelInputType,
ModelOutputType,
type StreamObject,
} from './providers/types';
import { StreamObjectParser } from './providers/utils';
import { ChatSession, ChatSessionService } from './session';
ActionStreamHost,
projectActionEventToChatEvent,
} from './runtime/hosts/action-stream-host';
import { TurnOrchestrator } from './runtime/turn-orchestrator';
import { CopilotStorage } from './storage';
import { ChatMessage, ChatQuerySchema } from './types';
import { getSignal, getTools } from './utils';
import { CopilotWorkflowService, GraphExecutorState } from './workflow';
import { getSignal } from './utils';
export interface ChatEvent {
type: 'event' | 'attachment' | 'message' | 'error' | 'ping';
@@ -74,11 +59,8 @@ export class CopilotController implements BeforeApplicationShutdown {
constructor(
private readonly config: Config,
private readonly server: ServerService,
private readonly chatSession: ChatSessionService,
private readonly context: CopilotContextService,
private readonly provider: CopilotProviderFactory,
private readonly workflow: CopilotWorkflowService,
private readonly orchestrator: TurnOrchestrator,
private readonly actionStreams: ActionStreamHost,
private readonly storage: CopilotStorage
) {}
@@ -92,85 +74,6 @@ export class CopilotController implements BeforeApplicationShutdown {
this.ongoingStreamCount$.complete();
}
private async chooseProvider(
outputType: ModelOutputType,
userId: string,
sessionId: string,
messageId?: string,
modelId?: string
): Promise<{
provider: CopilotProvider;
model: string;
hasAttachment: boolean;
}> {
const [, session] = await Promise.all([
this.chatSession.checkQuota(userId),
this.chatSession.get(sessionId),
]);
if (!session || session.config.userId !== userId) {
throw new CopilotSessionNotFound();
}
const model = await session.resolveModel(
this.server.features.includes(ServerFeature.Payment),
modelId
);
const hasAttachment = messageId
? !!(await session.getMessageById(messageId)).attachments?.length
: false;
const provider = await this.provider.getProvider({
outputType,
modelId: model,
});
if (!provider) {
throw new NoCopilotProviderAvailable({ modelId: model });
}
return { provider, model, hasAttachment };
}
private async appendSessionMessage(
sessionId: string,
messageId?: string,
retry = false
): Promise<[ChatMessage | undefined, ChatSession]> {
const session = await this.chatSession.get(sessionId);
if (!session) {
throw new CopilotSessionNotFound();
}
let latestMessage = undefined;
if (!messageId || retry) {
// revert the latest message generated by the assistant
// if messageId is provided, we will also revert latest user message
await this.chatSession.revertLatestMessage(sessionId, !!messageId);
session.revertLatestMessage(!!messageId);
if (!messageId) {
latestMessage = session.latestUserMessage;
}
}
if (messageId) {
await session.pushByMessageId(messageId);
}
return [latestMessage, session];
}
private parseNumber(value: string | string[] | undefined) {
if (!value) {
return undefined;
}
const num = Number.parseInt(Array.isArray(value) ? value[0] : value, 10);
if (Number.isNaN(num)) {
return undefined;
}
return num;
}
private mergePingStream(
messageId: string,
source$: Observable<ChatEvent>
@@ -184,59 +87,12 @@ export class CopilotController implements BeforeApplicationShutdown {
return merge(source$.pipe(finalize(() => subject$.next(null))), ping$);
}
private async prepareChatSession(
user: CurrentUser,
sessionId: string,
query: Record<string, string | string[]>,
outputType: ModelOutputType
) {
let { messageId, retry, modelId, params } = ChatQuerySchema.parse(query);
private toMessageEvent(messageId: string | undefined, data: string | object) {
return { type: 'message' as const, id: messageId, data };
}
const { provider, model } = await this.chooseProvider(
outputType,
user.id,
sessionId,
messageId,
modelId
);
const [latestMessage, session] = await this.appendSessionMessage(
sessionId,
messageId,
retry
);
const context = await this.context.getBySessionId(sessionId);
const contextParams =
(Array.isArray(context?.files) && context.files.length > 0) ||
(Array.isArray(context?.blobs) && context.blobs.length > 0)
? {
contextFiles: [
...context.files,
...(await context.getBlobMetadata()),
],
}
: {};
const lastParams = latestMessage
? {
...latestMessage.params,
content: latestMessage.content,
attachments: latestMessage.attachments,
}
: {};
const finalMessage = session.finish({
...params,
...lastParams,
...contextParams,
});
return {
provider,
model,
session,
finalMessage,
};
private toAttachmentEvent(messageId: string | undefined, data: string) {
return { type: 'attachment' as const, id: messageId, data };
}
@Sse('/chat/:sessionId/stream')
@@ -250,19 +106,6 @@ export class CopilotController implements BeforeApplicationShutdown {
const info: any = { sessionId, params: query, throwInStream: false };
try {
const { provider, model, session, finalMessage } =
await this.prepareChatSession(
user,
sessionId,
query,
ModelOutputType.Text
);
info.model = model;
info.finalMessage = finalMessage.filter(m => m.role !== 'system');
metrics.ai.counter('chat_stream_calls').add(1, { model });
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const { signal, onConnectionClosed } = getSignal(req);
let endBeforePromiseResolve = false;
onConnectionClosed(isAborted => {
@@ -271,51 +114,23 @@ export class CopilotController implements BeforeApplicationShutdown {
}
});
const { messageId, reasoning, webSearch, toolsConfig } =
ChatQuerySchema.parse(query);
const prepared = await this.orchestrator.streamText(
user.id,
sessionId,
query,
signal,
() => endBeforePromiseResolve
);
const source$ = from(
provider.streamText({ modelId: model }, finalMessage, {
...session.config.promptConfig,
signal,
user: user.id,
session: session.config.sessionId,
workspace: session.config.workspaceId,
reasoning,
webSearch,
tools: getTools(session.config.promptConfig?.tools, toolsConfig),
})
).pipe(
connect(shared$ =>
merge(
// actual chat event stream
shared$.pipe(
map(data => ({ type: 'message' as const, id: messageId, data }))
),
// save the generated text to the session
shared$.pipe(
reduce((acc, chunk) => acc + chunk, ''),
tap(buffer => {
session.push({
role: 'assistant',
content: endBeforePromiseResolve
? '> Request aborted'
: buffer,
createdAt: new Date(),
});
void session
.save()
.catch(err =>
this.logger.error(
'Failed to save session in sse stream',
err
)
);
}),
ignoreElements()
)
)
),
info.model = prepared.model;
info.finalMessage = prepared.finalMessage.filter(
m => m.role !== 'system'
);
metrics.ai.counter('chat_stream_calls').add(1, { model: prepared.model });
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const source$ = from(prepared.stream).pipe(
map(data => this.toMessageEvent(prepared.messageId, data)),
catchError(e => {
metrics.ai.counter('chat_stream_errors').add(1);
info.throwInStream = true;
@@ -326,7 +141,7 @@ export class CopilotController implements BeforeApplicationShutdown {
})
);
return this.mergePingStream(messageId || '', source$);
return this.mergePingStream(prepared.messageId || '', source$);
} catch (err) {
metrics.ai.counter('chat_stream_errors').add(1, info);
return mapSseError(err, info);
@@ -344,19 +159,6 @@ export class CopilotController implements BeforeApplicationShutdown {
const info: any = { sessionId, params: query, throwInStream: false };
try {
const { provider, model, session, finalMessage } =
await this.prepareChatSession(
user,
sessionId,
query,
ModelOutputType.Object
);
info.model = model;
info.finalMessage = finalMessage.filter(m => m.role !== 'system');
metrics.ai.counter('chat_object_stream_calls').add(1, { model });
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const { signal, onConnectionClosed } = getSignal(req);
let endBeforePromiseResolve = false;
onConnectionClosed(isAborted => {
@@ -365,55 +167,25 @@ export class CopilotController implements BeforeApplicationShutdown {
}
});
const { messageId, reasoning, webSearch, toolsConfig } =
ChatQuerySchema.parse(query);
const prepared = await this.orchestrator.streamObject(
user.id,
sessionId,
query,
signal,
() => endBeforePromiseResolve
);
const source$ = from(
provider.streamObject({ modelId: model }, finalMessage, {
...session.config.promptConfig,
signal,
user: user.id,
session: session.config.sessionId,
workspace: session.config.workspaceId,
reasoning,
webSearch,
tools: getTools(session.config.promptConfig?.tools, toolsConfig),
})
).pipe(
connect(shared$ =>
merge(
// actual chat event stream
shared$.pipe(
map(data => ({ type: 'message' as const, id: messageId, data }))
),
// save the generated text to the session
shared$.pipe(
reduce((acc, chunk) => acc.concat([chunk]), [] as StreamObject[]),
tap(result => {
const parser = new StreamObjectParser();
const streamObjects = parser.mergeTextDelta(result);
const content = parser.mergeContent(streamObjects);
session.push({
role: 'assistant',
content: endBeforePromiseResolve
? '> Request aborted'
: content,
streamObjects: endBeforePromiseResolve ? null : streamObjects,
createdAt: new Date(),
});
void session
.save()
.catch(err =>
this.logger.error(
'Failed to save session in sse stream',
err
)
);
}),
ignoreElements()
)
)
),
info.model = prepared.model;
info.finalMessage = prepared.finalMessage.filter(
m => m.role !== 'system'
);
metrics.ai.counter('chat_object_stream_calls').add(1, {
model: prepared.model,
});
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const source$ = from(prepared.stream).pipe(
map(data => this.toMessageEvent(prepared.messageId, data)),
catchError(e => {
metrics.ai.counter('chat_object_stream_errors').add(1);
info.throwInStream = true;
@@ -424,16 +196,16 @@ export class CopilotController implements BeforeApplicationShutdown {
})
);
return this.mergePingStream(messageId || '', source$);
return this.mergePingStream(prepared.messageId || '', source$);
} catch (err) {
metrics.ai.counter('chat_object_stream_errors').add(1, info);
return mapSseError(err, info);
}
}
@Sse('/chat/:sessionId/workflow')
@CallMetric('ai', 'chat_workflow', { timer: true })
async chatWorkflow(
@Sse('/actions/:sessionId/stream')
@CallMetric('ai', 'action_stream', { timer: true })
async actionStream(
@CurrentUser() user: CurrentUser,
@Req() req: Request,
@Param('sessionId') sessionId: string,
@@ -441,103 +213,26 @@ export class CopilotController implements BeforeApplicationShutdown {
): Promise<Observable<ChatEvent>> {
const info: any = { sessionId, params: query, throwInStream: false };
try {
let { messageId, params } = ChatQuerySchema.parse(query);
const { signal } = getSignal(req);
const [, session] = await this.appendSessionMessage(sessionId, messageId);
info.model = session.model;
metrics.ai.counter('workflow_calls').add(1, { model: session.model });
const latestMessage = session.stashMessages.findLast(
m => m.role === 'user'
const prepared = await this.actionStreams.stream(
user.id,
sessionId,
query,
signal
);
if (latestMessage) {
params = Object.assign({}, params, latestMessage.params, {
content: latestMessage.content,
attachments: latestMessage.attachments,
});
}
info.actionId = prepared.actionId;
info.actionVersion = prepared.actionVersion;
metrics.ai.counter('action_stream_calls').add(1, {
actionId: prepared.actionId,
actionVersion: prepared.actionVersion,
});
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const { signal, onConnectionClosed } = getSignal(req);
let endBeforePromiseResolve = false;
onConnectionClosed(isAborted => {
if (isAborted) {
endBeforePromiseResolve = true;
}
});
const source$ = from(
this.workflow.runGraph(params, session.model, {
...session.config.promptConfig,
signal,
user: user.id,
session: session.config.sessionId,
workspace: session.config.workspaceId,
})
).pipe(
connect(shared$ =>
merge(
// actual chat event stream
shared$.pipe(
map(data => {
switch (data.status) {
case GraphExecutorState.EmitContent:
return {
type: 'message' as const,
id: messageId,
data: data.content,
};
case GraphExecutorState.EmitAttachment:
return {
type: 'attachment' as const,
id: messageId,
data: data.attachment,
};
default:
return {
type: 'event' as const,
id: messageId,
data: {
status: data.status,
id: data.node.id,
type: data.node.config.nodeType,
},
};
}
})
),
// save the generated text to the session
shared$.pipe(
reduce((acc, chunk) => {
if (chunk.status === GraphExecutorState.EmitContent) {
acc += chunk.content;
}
return acc;
}, ''),
tap(content => {
session.push({
role: 'assistant',
content: endBeforePromiseResolve
? '> Request aborted'
: content,
createdAt: new Date(),
});
void session
.save()
.catch(err =>
this.logger.error(
'Failed to save session in sse stream',
err
)
);
}),
ignoreElements()
)
)
),
const source$ = from(prepared.stream).pipe(
map(data => projectActionEventToChatEvent(prepared.messageId, data)),
catchError(e => {
metrics.ai.counter('workflow_errors').add(1, info);
metrics.ai.counter('action_stream_errors').add(1, info);
info.throwInStream = true;
return mapSseError(e, info);
}),
@@ -546,9 +241,9 @@ export class CopilotController implements BeforeApplicationShutdown {
)
);
return this.mergePingStream(messageId || '', source$);
return this.mergePingStream(prepared.messageId || '', source$);
} catch (err) {
metrics.ai.counter('workflow_errors').add(1, info);
metrics.ai.counter('action_stream_errors').add(1, info);
return mapSseError(err, info);
}
}
@@ -563,36 +258,6 @@ export class CopilotController implements BeforeApplicationShutdown {
): Promise<Observable<ChatEvent>> {
const info: any = { sessionId, params: query, throwInStream: false };
try {
let { messageId, params } = ChatQuerySchema.parse(query);
const { provider, model, hasAttachment } = await this.chooseProvider(
ModelOutputType.Image,
user.id,
sessionId,
messageId
);
const [latestMessage, session] = await this.appendSessionMessage(
sessionId,
messageId
);
info.model = model;
metrics.ai.counter('images_stream_calls').add(1, { model });
if (latestMessage) {
params = Object.assign({}, params, latestMessage.params, {
content: latestMessage.content,
attachments: latestMessage.attachments,
});
}
const handleRemoteLink = this.storage.handleRemoteLink.bind(
this.storage,
user.id,
sessionId
);
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const { signal, onConnectionClosed } = getSignal(req);
let endBeforePromiseResolve = false;
onConnectionClosed(isAborted => {
@@ -601,59 +266,22 @@ export class CopilotController implements BeforeApplicationShutdown {
}
});
const source$ = from(
provider.streamImages(
{
modelId: model,
inputTypes: hasAttachment
? [ModelInputType.Image]
: [ModelInputType.Text],
},
session.finish(params),
{
...session.config.promptConfig,
quality: params.quality || undefined,
seed: this.parseNumber(params.seed),
signal,
user: user.id,
session: session.config.sessionId,
workspace: session.config.workspaceId,
}
)
).pipe(
mergeMap(handleRemoteLink),
connect(shared$ =>
merge(
// actual chat event stream
shared$.pipe(
map(attachment => ({
type: 'attachment' as const,
id: messageId,
data: attachment,
}))
),
// save the generated text to the session
shared$.pipe(
reduce((acc, chunk) => acc.concat([chunk]), [] as string[]),
tap(attachments => {
session.push({
role: 'assistant',
content: endBeforePromiseResolve ? '> Request aborted' : '',
attachments: endBeforePromiseResolve ? [] : attachments,
createdAt: new Date(),
});
void session
.save()
.catch(err =>
this.logger.error(
'Failed to save session in sse stream',
err
)
);
}),
ignoreElements()
)
)
const prepared = await this.orchestrator.streamImages(
user.id,
sessionId,
query,
signal,
() => endBeforePromiseResolve
);
info.model = prepared.model;
metrics.ai.counter('images_stream_calls').add(1, {
model: prepared.model,
});
this.ongoingStreamCount$.next(this.ongoingStreamCount$.value + 1);
const source$ = from(prepared.stream).pipe(
map(attachment =>
this.toAttachmentEvent(prepared.messageId, attachment)
),
catchError(e => {
metrics.ai.counter('images_stream_errors').add(1, info);
@@ -665,7 +293,7 @@ export class CopilotController implements BeforeApplicationShutdown {
)
);
return this.mergePingStream(messageId || '', source$);
return this.mergePingStream(prepared.messageId || '', source$);
} catch (err) {
metrics.ai.counter('images_stream_errors').add(1, info);
return mapSseError(err, info);
@@ -0,0 +1,95 @@
import { createHash } from 'node:crypto';
import { BadRequestException, Injectable } from '@nestjs/common';
import {
type FileUpload,
ImageFormatNotSupported,
sniffMime,
} from '../../../base';
import { WorkspacePolicyService } from '../../../core/permission';
import { processImage } from '../../../native';
import { CompatSubmissionStore } from '../compat/submission-store';
import type { PromptMessage } from '../providers/types';
import { ChatSessionService } from '../session';
import { CopilotStorage } from '../storage';
const COPILOT_IMAGE_MAX_EDGE = 1536;
type CreateInboxMessage = {
sessionId: string;
content?: string;
attachments?: string[];
blob?: Promise<FileUpload>;
blobs?: Promise<FileUpload>[];
params?: Record<string, any>;
};
@Injectable()
export class ConversationInboxService {
constructor(
private readonly chatSession: ChatSessionService,
private readonly policy: WorkspacePolicyService,
private readonly storage: CopilotStorage,
private readonly submissions: CompatSubmissionStore
) {}
async createMessage(
userId: string,
options: CreateInboxMessage
): Promise<string> {
const session = await this.chatSession.get(options.sessionId);
if (!session || session.config.userId !== userId) {
throw new BadRequestException('Session not found');
}
const attachments: PromptMessage['attachments'] = options.attachments || [];
const blobs = await Promise.all(
options.blob ? [options.blob] : options.blobs || []
);
if (blobs.length) {
await this.policy.assertCanUploadBlob(userId, session.config.workspaceId);
}
for (const blob of blobs) {
const uploaded = await this.storage.handleUpload(userId, blob);
const detectedMime =
sniffMime(uploaded.buffer, blob.mimetype)?.toLowerCase() ||
blob.mimetype;
let attachmentBuffer = uploaded.buffer;
let attachmentMimeType = detectedMime;
if (detectedMime.startsWith('image/')) {
try {
attachmentBuffer = await processImage(
uploaded.buffer,
COPILOT_IMAGE_MAX_EDGE,
true
);
attachmentMimeType = 'image/webp';
} catch {
throw new ImageFormatNotSupported({ format: detectedMime });
}
}
const filename = createHash('sha256')
.update(attachmentBuffer)
.digest('base64url');
const attachment = await this.storage.put(
userId,
session.config.workspaceId,
filename,
attachmentBuffer
);
attachments.push({ attachment, mimeType: attachmentMimeType });
}
return await this.submissions.create({
sessionId: options.sessionId,
content: options.content,
attachments,
params: options.params,
});
}
}
@@ -0,0 +1,68 @@
import { Injectable } from '@nestjs/common';
import { CopilotQuotaExceeded } from '../../../base';
import { QuotaService } from '../../../core/quota';
import { Models } from '../../../models';
import type { Turn } from '../core';
import type { ResolvedPrompt } from '../prompt';
@Injectable()
export class ConversationPolicy {
constructor(
private readonly models: Models,
private readonly quota: QuotaService
) {}
async getQuota(userId: string) {
const isCopilotUser = await this.models.userFeature.has(
userId,
'unlimited_copilot'
);
let limit: number | undefined;
if (!isCopilotUser) {
const quota = await this.quota.getUserQuota(userId);
limit = quota.copilotActionLimit;
}
const used = await this.models.copilotSession.countUserMessages(userId);
return { limit, used };
}
async checkQuota(userId: string) {
const { limit, used } = await this.getQuota(userId);
if (limit && Number.isFinite(limit) && used >= limit) {
throw new CopilotQuotaExceeded();
}
}
shouldScheduleTitle(prompt: Pick<ResolvedPrompt, 'action'>) {
return !prompt.action;
}
shouldGenerateTitle(input: { title: string | null; turns: Turn[] }) {
if (input.title || !input.turns.length) {
return false;
}
let hasUser = false;
let hasAssistant = false;
for (const turn of input.turns) {
if (turn.role === 'user') {
hasUser = true;
} else if (turn.role === 'assistant') {
hasAssistant = true;
}
if (hasUser && hasAssistant) {
return true;
}
}
return false;
}
buildTitlePromptContent(turns: Turn[]) {
return turns.map(turn => `[${turn.role}]: ${turn.content}`).join('\n');
}
}
@@ -0,0 +1,256 @@
import { Injectable } from '@nestjs/common';
import {
CleanupSessionOptions,
ListSessionOptions,
Models,
UpdateChatSessionOptions,
} from '../../../models';
import {
chatMessageFromTurn,
type Conversation,
type Turn,
turnFromChatMessage,
} from '../core';
import { type ChatMessage, ChatMessageSchema } from '../types';
type SessionRecord = NonNullable<
Awaited<ReturnType<Models['copilotSession']['get']>>
>;
type ConversationSeed = Parameters<
Models['copilotSession']['createWithPrompt']
>[0];
type ForkConversationSeed = Parameters<Models['copilotSession']['fork']>[0];
type ForkTurnsInput = Omit<ForkConversationSeed, 'messages'> & {
turns: Turn[];
};
@Injectable()
export class ConversationStore {
constructor(private readonly models: Models) {}
/**
* Durable-history boundary only.
*
* This store intentionally does not own:
* - quota / model / pin policy
* - title generation
* - prompt preload or rendering
* - compat ChatHistory / SSE projection
*/
private toConversation(session: SessionRecord): Conversation {
return {
id: session.id,
userId: session.userId,
workspaceId: session.workspaceId,
docId: session.docId,
pinned: session.pinned,
parentId: session.parentSessionId,
title: session.title,
createdAt: session.createdAt,
updatedAt: session.updatedAt,
};
}
private toTurns(session: SessionRecord): Turn[] {
return this.toMessages(session.messages).map(message =>
turnFromChatMessage(message, session.id)
);
}
private toMessages(messages: unknown): ChatMessage[] {
const parsed = ChatMessageSchema.array().safeParse(messages ?? []);
if (!parsed.success) return [];
return parsed.data;
}
async create(
seed: ConversationSeed,
reuseLatestChat = false
): Promise<string> {
return await this.models.copilotSession.createWithPrompt(
seed,
reuseLatestChat
);
}
async get(sessionId: string): Promise<
| {
conversation: Conversation;
turns: Turn[];
promptName: string;
tokenCost: number;
}
| undefined
> {
const session = await this.models.copilotSession.get(sessionId);
if (!session) {
return;
}
return {
conversation: this.toConversation(session),
turns: this.toTurns(session),
promptName: session.promptName,
tokenCost: session.tokenCost,
};
}
async getMeta(sessionId: string): Promise<
| {
conversation: Conversation;
promptName: string;
tokenCost: number;
}
| undefined
> {
const session = await this.models.copilotSession.getMeta(sessionId);
if (!session) return;
return {
conversation: {
id: session.id,
userId: session.userId,
workspaceId: session.workspaceId,
docId: session.docId,
pinned: session.pinned,
parentId: session.parentSessionId,
title: session.title,
createdAt: session.createdAt,
updatedAt: session.updatedAt,
},
promptName: session.promptName,
tokenCost: session.tokenCost,
};
}
async list(options: ListSessionOptions) {
const sessions = await this.models.copilotSession.list(options);
return sessions.map(session => ({
conversation: {
id: session.id,
userId: session.userId,
workspaceId: session.workspaceId,
docId: session.docId,
pinned: session.pinned,
parentId: session.parentSessionId,
title: session.title,
createdAt: session.createdAt,
updatedAt: session.updatedAt,
} satisfies Conversation,
turns: this.toMessages(session.messages).map(message =>
turnFromChatMessage(message, session.id)
),
promptName: session.promptName,
tokenCost: session.tokenCost,
}));
}
async listMeta(options: ListSessionOptions) {
const sessions = await this.models.copilotSession.list({
...options,
withMessages: false,
});
return sessions.map(session => ({
conversation: {
id: session.id,
userId: session.userId,
workspaceId: session.workspaceId,
docId: session.docId,
pinned: session.pinned,
parentId: session.parentSessionId,
title: session.title,
createdAt: session.createdAt,
updatedAt: session.updatedAt,
} satisfies Conversation,
promptName: session.promptName,
tokenCost: session.tokenCost,
}));
}
async appendTurns(input: {
sessionId: string;
userId: string;
prompt: { model: string };
turns: Turn[];
}) {
return await this.models.copilotSession.updateMessages({
...input,
messages: input.turns.map(turn => {
const { id: _id, ...message } = chatMessageFromTurn(turn);
return message;
}),
});
}
async appendTurn(input: {
sessionId: string;
userId: string;
prompt: { model: string };
turn: Turn;
compatSubmissionId?: string;
}) {
const message = await this.models.copilotSession.appendMessage({
sessionId: input.sessionId,
userId: input.userId,
prompt: input.prompt,
message: (() => {
const { id: _id, ...message } = chatMessageFromTurn(input.turn);
return { ...message, compatSubmissionId: input.compatSubmissionId };
})(),
});
return turnFromChatMessage(message, input.sessionId);
}
async findTurnByCompatSubmissionId(
sessionId: string,
compatSubmissionId: string
): Promise<Turn | undefined> {
const message =
await this.models.copilotSession.findMessageByCompatSubmissionId(
sessionId,
compatSubmissionId
);
if (!message) return;
return turnFromChatMessage(message, sessionId);
}
async update(options: UpdateChatSessionOptions): Promise<string> {
return await this.models.copilotSession.update(options);
}
async fork(seed: ForkTurnsInput): Promise<string> {
return await this.models.copilotSession.fork({
...seed,
messages: seed.turns.map(turn => {
const { id: _id, ...message } = chatMessageFromTurn(turn);
return message;
}),
});
}
async revertLatestTurn(sessionId: string, removeLatestUserMessage: boolean) {
return await this.models.copilotSession.revertLatestMessage(
sessionId,
removeLatestUserMessage
);
}
async cleanup(options: CleanupSessionOptions): Promise<string[]> {
return await this.models.copilotSession.cleanup(options);
}
async count(options: ListSessionOptions): Promise<number> {
return await this.models.copilotSession.count(options);
}
async unpin(workspaceId: string, userId: string) {
return await this.models.copilotSession.unpin(workspaceId, userId);
}
}
@@ -0,0 +1,108 @@
import type { PromptMessage, StreamObject } from '../providers/types';
import {
streamObjectToToolEvent,
toolEventToStreamObject,
} from '../runtime/contracts/runtime-event-contract';
import type { ChatMessage } from '../types';
import { type ToolEvent, type Turn, TurnSchema } from './types';
const normalizeRenderTrace = (
streamObjects: StreamObject[]
): StreamObject[] => {
return streamObjects.reduce((acc, current) => {
const previous = acc.at(-1);
switch (current.type) {
case 'reasoning':
case 'text-delta': {
if (previous?.type === current.type) {
previous.textDelta += current.textDelta;
} else {
acc.push({ ...current });
}
break;
}
case 'tool-result': {
const index = acc.findIndex(
candidate =>
candidate.type === 'tool-call' &&
candidate.toolCallId === current.toolCallId &&
candidate.toolName === current.toolName
);
if (index !== -1) {
acc[index] = { ...current };
} else {
acc.push({ ...current });
}
break;
}
default: {
acc.push({ ...current });
break;
}
}
return acc;
}, [] as StreamObject[]);
};
const deriveToolEvents = (renderTrace: StreamObject[]): ToolEvent[] =>
renderTrace
.map(streamObjectToToolEvent)
.filter((event): event is ToolEvent => !!event);
export const canonicalizeTurnTrace = (trace: {
renderTrace?: StreamObject[];
toolEvents?: ToolEvent[];
}) => {
const renderTrace =
trace.renderTrace && trace.renderTrace.length
? normalizeRenderTrace(trace.renderTrace)
: trace.toolEvents?.length
? trace.toolEvents.map(toolEventToStreamObject)
: [];
return { renderTrace, toolEvents: deriveToolEvents(renderTrace) };
};
export const turnFromChatMessage = (
message: ChatMessage,
conversationId: string
): Turn => {
const trace = canonicalizeTurnTrace({
renderTrace: message.streamObjects ?? [],
});
return TurnSchema.parse({
id: message.id,
conversationId,
role: message.role,
content: message.content,
attachments: message.attachments ?? [],
renderTrace: trace.renderTrace,
toolEvents: trace.toolEvents,
metadata: message.params ?? {},
createdAt: message.createdAt,
});
};
export const chatMessageFromTurn = (turn: Turn): ChatMessage => {
const { renderTrace } = canonicalizeTurnTrace(turn);
return {
id: turn.id,
role: turn.role,
content: turn.content,
attachments: turn.attachments.length ? turn.attachments : undefined,
params: turn.metadata,
streamObjects: renderTrace.length ? renderTrace : undefined,
createdAt: turn.createdAt,
};
};
export const promptMessageFromTurn = (turn: Turn): PromptMessage => ({
role: turn.role,
content: turn.content,
attachments: turn.attachments.length ? turn.attachments : undefined,
params: Object.keys(turn.metadata).length ? turn.metadata : undefined,
});
@@ -0,0 +1,2 @@
export * from './adapters';
export * from './types';
@@ -0,0 +1,58 @@
import { z } from 'zod';
import { ChatMessageAttachment } from '../providers/types';
import {
StreamObjectSchema,
type ToolEvent,
ToolEventSchema,
} from '../runtime/contracts/runtime-event-contract';
const CanonicalDateSchema = z.coerce.date();
export const ConversationSchema = z
.object({
id: z.string(),
userId: z.string(),
workspaceId: z.string(),
docId: z.string().nullable(),
pinned: z.boolean(),
parentId: z.string().nullable(),
title: z.string().nullable(),
createdAt: CanonicalDateSchema,
updatedAt: CanonicalDateSchema,
})
.strict();
export type Conversation = z.infer<typeof ConversationSchema>;
export const TurnSchema = z
.object({
id: z.string().optional(),
conversationId: z.string(),
role: z.enum(['system', 'assistant', 'user']),
content: z.string(),
attachments: z.array(ChatMessageAttachment).default([]),
renderTrace: z.array(StreamObjectSchema).default([]),
toolEvents: z.array(ToolEventSchema).default([]),
metadata: z.record(z.string(), z.any()).default({}),
createdAt: CanonicalDateSchema,
})
.strict();
export type Turn = z.infer<typeof TurnSchema>;
export const ValidatedStructuredValueSchema = z
.object({
value: z.any(),
schemaHash: z.string(),
schemaValidationVersion: z.string(),
provider: z.string(),
model: z.string(),
})
.strict();
export type ValidatedStructuredValue = z.infer<
typeof ValidatedStructuredValueSchema
>;
export type { ToolEvent };
@@ -25,14 +25,6 @@ export class CopilotCronJobs {
private readonly jobs: JobQueue
) {}
async triggerCleanupTrashedDocEmbeddings() {
await this.jobs.add(
'copilot.workspace.cleanupTrashedDocEmbeddings',
{},
{ jobId: 'daily-copilot-cleanup-trashed-doc-embeddings' }
);
}
@Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT)
async dailyCleanupJob() {
await this.jobs.add(
@@ -1,41 +1,32 @@
import { Logger } from '@nestjs/common';
import type { ModuleRef } from '@nestjs/core';
import { createHash } from 'node:crypto';
import { Injectable, Logger } from '@nestjs/common';
import { Config, CopilotProviderNotSupported } from '../../../base';
import { CopilotFailedToGenerateEmbedding } from '../../../base/error/errors.gen';
import {
ChunkSimilarity,
Embedding,
EMBEDDING_DIMENSIONS,
} from '../../../models';
import { CopilotProviderFactory } from '../providers/factory';
import type { CopilotProvider } from '../providers/provider';
import {
type CopilotRerankRequest,
type ModelFullConditions,
ModelInputType,
ModelOutputType,
} from '../providers/types';
import { type CopilotRerankRequest } from '../providers/types';
import { CapabilityRuntime } from '../runtime/capability-runtime';
import { TaskPolicy } from '../runtime/task-policy';
import { EmbeddingClient, type ReRankResult } from './types';
const EMBEDDING_MODEL = 'gemini-embedding-001';
const RERANK_MODEL = 'gpt-4o-mini';
class ProductionEmbeddingClient extends EmbeddingClient {
private readonly logger = new Logger(ProductionEmbeddingClient.name);
constructor(
private readonly config: Config,
private readonly providerFactory: CopilotProviderFactory
private readonly taskPolicy: TaskPolicy,
private readonly runtime: CapabilityRuntime
) {
super();
}
override async configured(): Promise<boolean> {
const embedding = await this.providerFactory.getProvider({
modelId: this.getEmbeddingModelId(),
outputType: ModelOutputType.Embedding,
});
const result = Boolean(embedding);
const result = await this.runtime.embeddingConfigured(
this.taskPolicy.resolveEmbeddingModelId()
);
if (!result) {
this.logger.warn(
'Copilot embedding client is not configured properly, please check your configuration.'
@@ -44,42 +35,14 @@ class ProductionEmbeddingClient extends EmbeddingClient {
return result;
}
private async getProvider(
cond: ModelFullConditions
): Promise<CopilotProvider> {
const provider = await this.providerFactory.getProvider(cond);
if (!provider) {
throw new CopilotProviderNotSupported({
provider: 'embedding',
kind: cond.outputType || 'embedding',
});
}
return provider;
}
private getEmbeddingModelId() {
return this.config.copilot?.scenarios?.override_enabled
? this.config.copilot.scenarios.scenarios?.embedding || EMBEDDING_MODEL
: EMBEDDING_MODEL;
}
async getEmbeddings(input: string[]): Promise<Embedding[]> {
const provider = await this.getProvider({
modelId: this.getEmbeddingModelId(),
outputType: ModelOutputType.Embedding,
const modelId = this.taskPolicy.resolveEmbeddingModelId();
const embeddings = await this.runtime.embed(modelId, input, {
dimensions: EMBEDDING_DIMENSIONS,
});
this.logger.verbose(
`Using provider ${provider.type} for embedding: ${input.join(', ')}`
);
const embeddings = await provider.embedding(
{ inputTypes: [ModelInputType.Text] },
input,
{ dimensions: EMBEDDING_DIMENSIONS }
);
if (embeddings.length !== input.length) {
throw new CopilotFailedToGenerateEmbedding({
provider: provider.type,
provider: modelId,
message: `Expected ${input.length} embeddings, got ${embeddings.length}`,
});
}
@@ -108,11 +71,6 @@ class ProductionEmbeddingClient extends EmbeddingClient {
): Promise<ReRankResult> {
if (!embeddings.length) return [];
const provider = await this.getProvider({
modelId: RERANK_MODEL,
outputType: ModelOutputType.Rerank,
});
const rerankRequest: CopilotRerankRequest = {
query,
candidates: embeddings.map((embedding, index) => ({
@@ -121,8 +79,8 @@ class ProductionEmbeddingClient extends EmbeddingClient {
})),
};
const ranks = await provider.rerank(
{ modelId: RERANK_MODEL },
const ranks = await this.runtime.rerank(
this.taskPolicy.resolveRerankModelId(),
rerankRequest,
{ signal }
);
@@ -211,32 +169,40 @@ class ProductionEmbeddingClient extends EmbeddingClient {
}
}
let EMBEDDING_CLIENT: EmbeddingClient | undefined;
export async function getEmbeddingClient(
moduleRef: ModuleRef
): Promise<EmbeddingClient | undefined> {
if (EMBEDDING_CLIENT) {
return EMBEDDING_CLIENT;
@Injectable()
export class CopilotEmbeddingClientService {
private client: EmbeddingClient | undefined;
constructor(
private readonly taskPolicy: TaskPolicy,
private readonly runtime: CapabilityRuntime
) {}
async refresh() {
const client = new ProductionEmbeddingClient(this.taskPolicy, this.runtime);
this.client = (await client.configured()) ? client : undefined;
return this.client;
}
const config = moduleRef.get(Config, { strict: false });
const providerFactory = moduleRef.get(CopilotProviderFactory, {
strict: false,
});
const client = new ProductionEmbeddingClient(config, providerFactory);
if (await client.configured()) {
EMBEDDING_CLIENT = client;
getClient() {
return this.client;
}
return EMBEDDING_CLIENT;
}
export class MockEmbeddingClient extends EmbeddingClient {
private embed(content: string) {
const seed = createHash('sha256').update(content).digest();
return Array.from({ length: EMBEDDING_DIMENSIONS }, (_, index) => {
const byte = seed[index % seed.length];
return byte / 255;
});
}
async getEmbeddings(input: string[]): Promise<Embedding[]> {
return input.map((_, i) => ({
return input.map((content, i) => ({
index: i,
content: input[i],
embedding: Array.from({ length: EMBEDDING_DIMENSIONS }, () =>
Math.random()
),
content,
embedding: this.embed(content),
}));
}
}
@@ -1,4 +1,4 @@
export { getEmbeddingClient, MockEmbeddingClient } from './client';
export { CopilotEmbeddingClientService, MockEmbeddingClient } from './client';
export { CopilotEmbeddingJob } from './job';
export type { Chunk, DocFragment } from './types';
export { EmbeddingClient } from './types';
@@ -1,5 +1,4 @@
import { Injectable, Logger } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import {
BlobNotFound,
@@ -18,7 +17,7 @@ import { readAllDocIdsFromWorkspaceSnapshot } from '../../../core/utils/blocksui
import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
import { getEmbeddingClient } from './client';
import { CopilotEmbeddingClientService } from './client';
import type { Chunk, DocFragment } from './types';
import { EmbeddingClient } from './types';
@@ -32,12 +31,13 @@ export class CopilotEmbeddingJob {
private client: EmbeddingClient | undefined;
constructor(
private readonly moduleRef: ModuleRef,
private readonly embeddingClients: CopilotEmbeddingClientService,
private readonly doc: DocReader,
private readonly event: EventBus,
private readonly models: Models,
private readonly queue: JobQueue,
private readonly storage: CopilotStorage
private readonly storage: CopilotStorage,
private readonly workspaceStorage: WorkspaceBlobStorage
) {}
@OnEvent('config.init')
@@ -54,7 +54,7 @@ export class CopilotEmbeddingJob {
this.supportEmbedding =
await this.models.copilotContext.checkEmbeddingAvailable();
if (this.supportEmbedding) {
this.client = await getEmbeddingClient(this.moduleRef);
this.client = await this.embeddingClients.refresh();
}
}
@@ -64,10 +64,15 @@ export class CopilotEmbeddingJob {
}
@CallMetric('ai', 'addFileEmbeddingQueue')
async addFileEmbeddingQueue(file: Jobs['copilot.embedding.files']) {
async addFileEmbeddingQueue(
file: Jobs['copilot.embedding.files'],
options?: { priority?: number }
) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.files', file);
await this.queue.add('copilot.embedding.files', file, {
priority: options?.priority,
});
}
@CallMetric('ai', 'addBlobEmbeddingQueue')
@@ -231,10 +236,7 @@ export class CopilotEmbeddingJob {
blobId: string,
fileName: string
) {
const workspaceStorage = this.moduleRef.get(WorkspaceBlobStorage, {
strict: false,
});
const { body } = await workspaceStorage.get(workspaceId, blobId);
const { body } = await this.workspaceStorage.get(workspaceId, blobId);
if (!body) throw new BlobNotFound({ spaceId: workspaceId, blobId });
const buffer = await readStream(body);
return new File([buffer], fileName);
@@ -445,6 +447,12 @@ export class CopilotEmbeddingJob {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no content change, skipping embedding.`
);
if (contextId) {
this.event.emit('workspace.doc.embed.finished', {
contextId,
docId,
});
}
return;
}
@@ -487,6 +495,12 @@ export class CopilotEmbeddingJob {
);
}
}
if (contextId) {
this.event.emit('workspace.doc.embed.finished', {
contextId,
docId,
});
}
} catch (error: any) {
if (contextId) {
this.event.emit('workspace.doc.embed.failed', {
@@ -36,6 +36,11 @@ declare global {
docId: string;
};
'workspace.doc.embed.finished': {
contextId: string;
docId: string;
};
'workspace.file.embed.finished': {
contextId: string;
fileId: string;
@@ -7,82 +7,55 @@ import { DocStorageModule } from '../../core/doc';
import { FeatureModule } from '../../core/features';
import { PermissionModule } from '../../core/permission';
import { QuotaModule } from '../../core/quota';
import { StorageModule } from '../../core/storage';
import { WorkspaceModule } from '../../core/workspaces';
import { IndexerModule } from '../indexer';
import {
CopilotContextResolver,
CopilotContextRootResolver,
CopilotContextService,
} from './context';
import { CopilotController } from './controller';
import { CopilotCronJobs } from './cron';
import { CopilotEmbeddingJob } from './embedding';
import { WorkspaceMcpController } from './mcp/controller';
import { WorkspaceMcpProvider } from './mcp/provider';
import { ChatMessageCache } from './message';
import { PromptService } from './prompt';
import { CopilotProviderFactory, CopilotProviders } from './providers';
import {
CopilotResolver,
PromptsManagementResolver,
UserCopilotResolver,
} from './resolver';
import { ChatSessionService } from './session';
import { CopilotStorage } from './storage';
import {
CopilotTranscriptionResolver,
CopilotTranscriptionService,
} from './transcript';
import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow';
import {
CopilotWorkspaceEmbeddingConfigResolver,
CopilotWorkspaceEmbeddingResolver,
CopilotWorkspaceService,
} from './workspace';
COPILOT_API_PROVIDERS,
COPILOT_FEATURE_PROVIDERS,
COPILOT_KERNEL_PROVIDERS,
} from './module-providers';
const COPILOT_SHARED_IMPORTS = [
DocStorageModule,
FeatureModule,
QuotaModule,
PermissionModule,
ServerConfigModule,
StorageModule,
WorkspaceModule,
IndexerModule,
];
@Module({
imports: [...COPILOT_SHARED_IMPORTS],
providers: [...COPILOT_KERNEL_PROVIDERS],
exports: [...COPILOT_KERNEL_PROVIDERS],
})
export class CopilotKernelModule {}
@Module({
imports: [...COPILOT_SHARED_IMPORTS, CopilotKernelModule],
providers: [...COPILOT_FEATURE_PROVIDERS],
exports: [...COPILOT_FEATURE_PROVIDERS],
})
export class CopilotFeatureModule {}
@Module({
imports: [
DocStorageModule,
FeatureModule,
QuotaModule,
PermissionModule,
ServerConfigModule,
WorkspaceModule,
IndexerModule,
],
providers: [
// providers
...CopilotProviders,
CopilotProviderFactory,
// services
ChatSessionService,
CopilotResolver,
ChatMessageCache,
PromptService,
CopilotStorage,
// workflow
CopilotWorkflowService,
...CopilotWorkflowExecutors,
// context
CopilotContextResolver,
CopilotContextService,
// jobs
CopilotEmbeddingJob,
CopilotCronJobs,
// transcription
CopilotTranscriptionService,
CopilotTranscriptionResolver,
// workspace embeddings
CopilotWorkspaceService,
CopilotWorkspaceEmbeddingResolver,
CopilotWorkspaceEmbeddingConfigResolver,
// gql resolvers
UserCopilotResolver,
PromptsManagementResolver,
CopilotContextRootResolver,
// mcp
WorkspaceMcpProvider,
...COPILOT_SHARED_IMPORTS,
CopilotKernelModule,
CopilotFeatureModule,
],
providers: [...COPILOT_API_PROVIDERS],
exports: [...COPILOT_API_PROVIDERS],
})
export class CopilotApiModule {}
@Module({
imports: [CopilotKernelModule, CopilotFeatureModule, CopilotApiModule],
controllers: [CopilotController, WorkspaceMcpController],
})
export class CopilotModule {}
@@ -1,27 +0,0 @@
import { randomUUID } from 'node:crypto';
import { Injectable } from '@nestjs/common';
import { SessionCache } from '../../base';
import { SubmittedMessage, SubmittedMessageSchema } from './types';
const CHAT_MESSAGE_KEY = 'chat-message';
const CHAT_MESSAGE_TTL = 3600 * 1 * 1000; // 1 hours
@Injectable()
export class ChatMessageCache {
constructor(private readonly cache: SessionCache) {}
async get(id: string): Promise<SubmittedMessage | undefined> {
return await this.cache.get(`${CHAT_MESSAGE_KEY}:${id}`);
}
async set(message: SubmittedMessage): Promise<string> {
const parsedMessage = SubmittedMessageSchema.parse(message);
const id = randomUUID();
await this.cache.set(`${CHAT_MESSAGE_KEY}:${id}`, parsedMessage, {
ttl: CHAT_MESSAGE_TTL,
});
return id;
}
}
@@ -0,0 +1,139 @@
import { HistoryAttachmentUrlProjector } from './compat/history-attachment-url-projector';
import { CompatHistoryProjector } from './compat/history-projector';
import { HistoryPromptPreloadProjector } from './compat/history-prompt-preload-projector';
import { HistoryVisibilityPolicy } from './compat/history-visibility-policy';
import { CompatSubmissionStore } from './compat/submission-store';
import {
CopilotContextResolver,
CopilotContextRootResolver,
CopilotContextService,
} from './context';
import { ConversationInboxService } from './conversation/inbox';
import { ConversationPolicy } from './conversation/policy';
import { ConversationStore } from './conversation/store';
import { CopilotCronJobs } from './cron';
import {
CopilotEmbeddingClientService,
CopilotEmbeddingJob,
} from './embedding';
import { WorkspaceMcpProvider } from './mcp/provider';
import { PromptService } from './prompt';
import {
CopilotProviderFactory,
CopilotProviderLifecycleService,
CopilotProviderRegistryService,
CopilotProviders,
} from './providers';
import { CopilotResolver, UserCopilotResolver } from './resolver';
import { ActionRuntimeBridge } from './runtime/action-runtime-bridge';
import { CapabilityRuntime } from './runtime/capability-runtime';
import { CopilotExecutionMetrics } from './runtime/execution-metrics';
import { ExecutionPlanBuilder } from './runtime/execution-plan';
import { ActionStreamHost } from './runtime/hosts/action-stream-host';
import { AttachmentAdmissionHost } from './runtime/hosts/attachment-admission';
import { AttachmentMaterializer } from './runtime/hosts/attachment-materializer';
import { CapabilityPolicyHost } from './runtime/hosts/capability-policy-host';
import { ConversationHost } from './runtime/hosts/conversation-host';
import { ImageResultHost } from './runtime/hosts/image-result-host';
import { ResponsePostprocessor } from './runtime/hosts/response-postprocessor';
import { ToolExecutorHost } from './runtime/hosts/tool-executor-host';
import { TurnPersistence } from './runtime/hosts/turn-persistence';
import { ModelSelectionPolicy } from './runtime/model-selection-policy';
import { NativeExecutionEngine } from './runtime/native-execution-engine';
import { PromptRuntime } from './runtime/prompt-runtime';
import { TaskPolicy } from './runtime/task-policy';
import { ToolRuntime } from './runtime/tool-runtime';
import { TurnOrchestrator } from './runtime/turn-orchestrator';
import { ChatSessionService } from './session';
import { CopilotStorage } from './storage';
import {
CopilotTranscriptionResolver,
CopilotTranscriptionService,
} from './transcript';
import {
CopilotWorkspaceEmbeddingConfigResolver,
CopilotWorkspaceEmbeddingResolver,
CopilotWorkspaceService,
} from './workspace';
export const COPILOT_PROVIDER_PROVIDERS = [
...CopilotProviders,
CopilotProviderRegistryService,
CopilotProviderFactory,
CopilotProviderLifecycleService,
];
export const COPILOT_RUNTIME_PROVIDERS = [
ChatSessionService,
ConversationStore,
ConversationInboxService,
ConversationPolicy,
HistoryAttachmentUrlProjector,
CompatHistoryProjector,
HistoryPromptPreloadProjector,
CompatSubmissionStore,
HistoryVisibilityPolicy,
CopilotContextService,
CopilotEmbeddingClientService,
PromptService,
ModelSelectionPolicy,
ActionRuntimeBridge,
CopilotExecutionMetrics,
ExecutionPlanBuilder,
PromptRuntime,
CapabilityPolicyHost,
ConversationHost,
CapabilityRuntime,
NativeExecutionEngine,
TaskPolicy,
ToolRuntime,
ToolExecutorHost,
AttachmentMaterializer,
AttachmentAdmissionHost,
ActionStreamHost,
ImageResultHost,
ResponsePostprocessor,
CopilotStorage,
TurnPersistence,
];
export const COPILOT_CONTEXT_PROVIDERS = [CopilotContextResolver];
export const COPILOT_TRANSCRIPT_PROVIDERS = [
CopilotTranscriptionService,
CopilotTranscriptionResolver,
];
export const COPILOT_WORKSPACE_PROVIDERS = [
CopilotWorkspaceService,
CopilotWorkspaceEmbeddingResolver,
CopilotWorkspaceEmbeddingConfigResolver,
];
export const COPILOT_RESOLVER_PROVIDERS = [
CopilotResolver,
UserCopilotResolver,
CopilotContextRootResolver,
];
export const COPILOT_JOB_PROVIDERS = [CopilotEmbeddingJob, CopilotCronJobs];
export const COPILOT_MCP_PROVIDERS = [WorkspaceMcpProvider];
export const COPILOT_KERNEL_PROVIDERS = [
...COPILOT_PROVIDER_PROVIDERS,
...COPILOT_RUNTIME_PROVIDERS,
];
export const COPILOT_FEATURE_PROVIDERS = [
TurnOrchestrator,
...COPILOT_CONTEXT_PROVIDERS,
...COPILOT_TRANSCRIPT_PROVIDERS,
...COPILOT_WORKSPACE_PROVIDERS,
...COPILOT_JOB_PROVIDERS,
];
export const COPILOT_API_PROVIDERS = [
...COPILOT_RESOLVER_PROVIDERS,
...COPILOT_MCP_PROVIDERS,
];
@@ -1,183 +0,0 @@
import { type Tokenizer } from '@affine/server-native';
import { Logger } from '@nestjs/common';
import { AiPrompt } from '@prisma/client';
import Mustache from 'mustache';
import { getTokenEncoder } from '../../../native';
import type {
PromptConfig,
PromptMessage,
PromptParams,
} from '../providers/types';
// disable escaping
Mustache.escape = (text: string) => text;
function extractMustacheParams(template: string) {
const regex = /\{\{\s*([^{}]+)\s*\}\}/g;
const params = [];
let match;
while ((match = regex.exec(template)) !== null) {
params.push(match[1]);
}
return Array.from(new Set(params));
}
export class ChatPrompt {
private readonly logger = new Logger(ChatPrompt.name);
public readonly encoder: Tokenizer | null;
private readonly promptTokenSize: number;
private readonly templateParamKeys: string[] = [];
private readonly templateParams: PromptParams = {};
static createFromPrompt(
options: Omit<
AiPrompt,
'id' | 'createdAt' | 'updatedAt' | 'modified' | 'config'
> & {
messages: PromptMessage[];
config: PromptConfig | undefined;
}
) {
return new ChatPrompt(
options.name,
options.action || undefined,
options.model,
options.optionalModels,
options.config,
options.messages
);
}
constructor(
public readonly name: string,
public readonly action: string | undefined,
public readonly model: string,
public readonly optionalModels: string[],
public readonly config: PromptConfig | undefined,
private readonly messages: PromptMessage[]
) {
this.encoder = getTokenEncoder(model);
this.promptTokenSize = this.encode(messages.map(m => m.content).join(''));
this.templateParamKeys = extractMustacheParams(
messages.map(m => m.content).join('')
);
this.templateParams = messages.reduce(
(acc, m) => Object.assign(acc, m.params),
{} as PromptParams
);
}
/**
* get prompt token size
*/
get tokens() {
return this.promptTokenSize;
}
/**
* get prompt param keys in template
*/
get paramKeys() {
return this.templateParamKeys.slice();
}
/**
* get prompt params
*/
get params() {
return { ...this.templateParams };
}
encode(message: string) {
return this.encoder?.count(message) || 0;
}
private checkParams(params: PromptParams, sessionId?: string) {
const selfParams = this.templateParams;
for (const key of Object.keys(selfParams)) {
const options = selfParams[key];
const income = params[key];
if (
typeof income !== 'string' ||
(Array.isArray(options) && !options.includes(income))
) {
if (sessionId) {
const prefix = income
? `Invalid param value: ${key}=${income}`
: `Missing param value: ${key}`;
this.logger.warn(
`${prefix} in session ${sessionId}, use default options: ${Array.isArray(options) ? options[0] : options}`
);
}
if (Array.isArray(options)) {
// use the first option if income is not in options
params[key] = options[0];
} else {
params[key] = options;
}
}
}
}
private preDefinedParams(params: PromptParams) {
const {
language,
timezone,
docs,
contextFiles: files,
selectedMarkdown,
selectedSnapshot,
html,
currentDocId,
} = params;
return {
'affine::date': new Date().toLocaleDateString(),
'affine::language': language || 'same language as the user query',
'affine::timezone': timezone || 'no preference',
'affine::hasDocsRef': Array.isArray(docs) && docs.length > 0,
'affine::hasFilesRef': Array.isArray(files) && files.length > 0,
'affine::hasSelected': !!selectedMarkdown || !!selectedSnapshot || !!html,
'affine::hasCurrentDoc':
typeof currentDocId === 'string' && currentDocId.trim().length > 0,
};
}
/**
* render prompt messages with params
* @param params record of params, e.g. { name: 'Alice' }
* @returns e.g. [{ role: 'system', content: 'Hello, {{name}}' }] => [{ role: 'system', content: 'Hello, Alice' }]
*/
finish(params: PromptParams, sessionId?: string): PromptMessage[] {
this.checkParams(params, sessionId);
const { attachments: attach, ...restParams } = Object.fromEntries(
Object.entries(params).filter(([k]) => !k.startsWith('affine::'))
);
const paramsAttach = Array.isArray(attach) ? attach : [];
return this.messages.map(
({ attachments: attach, content, params: _, ...rest }) => {
const result: PromptMessage = {
...rest,
params,
content: Mustache.render(
content,
Object.assign({}, restParams, this.preDefinedParams(restParams))
),
};
const attachments = [
...(Array.isArray(attach) ? attach : []),
...paramsAttach,
];
if (attachments.length && rest.role === 'user') {
result.attachments = attachments;
}
return result;
}
);
}
}
@@ -1,3 +1,2 @@
export { ChatPrompt } from './chat-prompt';
export { prompts } from './prompts';
export { PromptService } from './service';
export type { ResolvedPrompt } from './spec';
@@ -0,0 +1,260 @@
import {
llmCollectPromptMetadata,
llmCountPromptTokens,
llmGetBuiltInPromptSpec,
llmListBuiltInPromptSpecs,
llmRenderBuiltInPrompt,
llmRenderBuiltInSessionPrompt,
llmRenderPrompt,
llmRenderSessionPrompt,
type NativeBuiltInPromptRenderRequest as NativeBuiltInPromptRenderContract,
type NativeBuiltInPromptSessionRenderRequest as NativeBuiltInPromptSessionContract,
type NativePromptCountTokensRequest as NativePromptTokenCountContract,
type NativePromptCountTokensResponse as NativePromptTokenCountResult,
type NativePromptMetadataRequest as NativePromptMetadataContract,
type NativePromptMetadataResponse as NativePromptMetadataResult,
type NativePromptRenderRequest as NativePromptRenderContract,
type NativePromptRenderResponse as NativePromptRenderResult,
type NativePromptSessionRenderRequest as NativePromptSessionContract,
type NativePromptSessionRenderResponse as NativePromptSessionResult,
} from '../../../native';
import type { PromptMessage, PromptParams } from '../providers/types';
import { projectPromptMessageForNative } from '../runtime/contracts';
import type { PromptSpec } from './spec';
export type NativePromptRenderRequest = Omit<
NativePromptRenderContract,
'messages' | 'templateParams' | 'renderParams'
> & {
messages: PromptMessage[];
templateParams: PromptParams;
renderParams: PromptParams;
};
export type NativePromptRenderResponse = Omit<
NativePromptRenderResult,
'messages'
> & {
messages: PromptMessage[];
};
export type NativeBuiltInPromptRenderRequest = Omit<
NativeBuiltInPromptRenderContract,
'renderParams'
> & {
renderParams: PromptParams;
};
export type NativePromptCountTokensRequest = Omit<
NativePromptTokenCountContract,
'messages' | 'model'
> & {
model?: string | null;
messages: Pick<PromptMessage, 'content'>[];
};
export type NativePromptCountTokensResponse = NativePromptTokenCountResult;
export type NativePromptMetadataRequest = Omit<
NativePromptMetadataContract,
'messages'
> & {
messages: PromptMessage[];
};
export type NativePromptMetadataResponse = Omit<
NativePromptMetadataResult,
'templateParams'
> & {
templateParams: PromptParams;
};
export type NativePromptSessionRenderRequest = Omit<
NativePromptSessionContract,
'prompt' | 'turns' | 'renderParams'
> & {
prompt: Omit<
NativePromptSessionContract['prompt'],
'templateParams' | 'messages' | 'model'
> & {
model?: string | null;
templateParams: PromptParams;
messages: PromptMessage[];
};
turns: PromptMessage[];
renderParams: PromptParams;
};
export type NativePromptSessionRenderResponse = Omit<
NativePromptSessionResult,
'messages'
> & {
messages: PromptMessage[];
};
export type NativeBuiltInPromptSessionRenderRequest = Omit<
NativeBuiltInPromptSessionContract,
'turns' | 'renderParams'
> & {
turns: PromptMessage[];
renderParams: PromptParams;
};
type NativePromptContractMessage =
NativePromptRenderContract['messages'][number];
function toNativePromptMessage(
message: PromptMessage
): NativePromptContractMessage {
return projectPromptMessageForNative(message).message;
}
function fromNativePromptMessage(
message: NativePromptContractMessage
): PromptMessage {
return {
role: message.role,
content: message.content,
...(message.attachments ? { attachments: message.attachments } : {}),
...(message.params ? { params: message.params } : {}),
...(message.responseFormat
? { responseFormat: message.responseFormat }
: {}),
};
}
export function renderPromptNative(
request: NativePromptRenderRequest
): NativePromptRenderResponse {
const normalizedMessages = request.messages.map(toNativePromptMessage);
const rendered = llmRenderPrompt({
messages: normalizedMessages,
templateParams: request.templateParams,
renderParams: request.renderParams,
});
return {
...rendered,
messages: rendered.messages.map(fromNativePromptMessage),
};
}
export function renderBuiltInPromptNative(
request: NativeBuiltInPromptRenderRequest
): NativePromptRenderResponse {
const rendered = llmRenderBuiltInPrompt({
name: request.name,
renderParams: request.renderParams,
});
return {
...rendered,
messages: rendered.messages.map(fromNativePromptMessage),
};
}
export function renderPromptSessionNative(
request: NativePromptSessionRenderRequest
): NativePromptSessionRenderResponse {
const rendered = llmRenderSessionPrompt({
...request,
prompt: {
...request.prompt,
messages: request.prompt.messages.map(toNativePromptMessage),
model: request.prompt.model ?? undefined,
},
turns: request.turns.map(toNativePromptMessage),
renderParams: request.renderParams,
});
return {
...rendered,
messages: rendered.messages.map(fromNativePromptMessage),
};
}
export function renderBuiltInPromptSessionNative(
request: NativeBuiltInPromptSessionRenderRequest
): NativePromptSessionRenderResponse {
const rendered = llmRenderBuiltInSessionPrompt({
...request,
turns: request.turns.map(toNativePromptMessage),
renderParams: request.renderParams,
});
return {
...rendered,
messages: rendered.messages.map(fromNativePromptMessage),
};
}
export function countPromptTokensNative(
request: NativePromptCountTokensRequest
): NativePromptCountTokensResponse {
return llmCountPromptTokens({
...request,
model: request.model ?? undefined,
});
}
export function collectPromptMetadataNative(
request: NativePromptMetadataRequest
): NativePromptMetadataResponse {
return llmCollectPromptMetadata({
messages: request.messages.map(toNativePromptMessage),
});
}
export function listBuiltInPromptSpecsNative(): PromptSpec[] {
return llmListBuiltInPromptSpecs().map(spec => ({
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels,
config: spec.config,
params: spec.params
? Object.fromEntries(
Object.entries(spec.params).map(([key, value]) => [
key,
{
default: value.default,
enum: value.enumValues,
},
])
)
: undefined,
messages: spec.messages.map(message => ({
role: message.role,
template: message.template,
})),
}));
}
export function getBuiltInPromptSpecNative(name: string): PromptSpec | null {
const spec = llmGetBuiltInPromptSpec(name);
if (!spec) {
return null;
}
return {
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels,
config: spec.config,
params: spec.params
? Object.fromEntries(
Object.entries(spec.params).map(([key, value]) => [
key,
{
default: value.default,
enum: value.enumValues,
},
])
)
: undefined,
messages: spec.messages.map(message => ({
role: message.role,
template: message.template,
})),
};
}
File diff suppressed because it is too large Load Diff
@@ -1,323 +1,110 @@
import { Injectable, Logger, OnApplicationBootstrap } from '@nestjs/common';
import { Transactional } from '@nestjs-cls/transactional';
import { Prisma, PrismaClient } from '@prisma/client';
import { Injectable, Logger } from '@nestjs/common';
import { Config, OnEvent } from '../../../base';
import type { PromptMessage, PromptParams } from '../providers/types';
import {
PromptConfig,
PromptConfigSchema,
PromptMessage,
PromptMessageSchema,
} from '../providers/types';
import { ChatPrompt } from './chat-prompt';
import {
CopilotPromptScenario,
type Prompt,
prompts,
refreshPrompts,
Scenario,
} from './prompts';
collectPromptMetadataNative,
countPromptTokensNative,
getBuiltInPromptSpecNative,
renderBuiltInPromptNative,
renderBuiltInPromptSessionNative,
renderPromptNative,
renderPromptSessionNative,
} from './native-contract';
import type { Prompt, PromptSpec, ResolvedPrompt } from './spec';
@Injectable()
export class PromptService implements OnApplicationBootstrap {
private readonly logger = new Logger(PromptService.name);
private readonly cache = new Map<string, ChatPrompt>();
private readonly inMemoryPrompts = new Map<string, Prompt>();
constructor(
private readonly config: Config,
private readonly db: PrismaClient
) {}
async onApplicationBootstrap() {
this.resetInMemoryPrompts();
await refreshPrompts(this.db);
export class PromptService {
protected readonly logger = new Logger(PromptService.name);
constructor() {
this.logger.log('Using native built-in prompt catalog.');
}
@OnEvent('config.init')
async onConfigInit() {
await this.setup(this.config.copilot?.scenarios);
}
@OnEvent('config.changed')
async onConfigChanged(event: Events['config.changed']) {
if ('copilot' in event.updates) {
await this.setup(event.updates.copilot?.scenarios);
}
}
protected async setup(scenarios?: CopilotPromptScenario) {
this.ensureInMemoryPrompts();
if (!!scenarios && scenarios.override_enabled && scenarios.scenarios) {
this.logger.log('Updating prompts based on scenarios...');
for (const [scenario, model] of Object.entries(scenarios.scenarios)) {
const promptNames = Scenario[scenario as keyof typeof Scenario] || [];
if (!promptNames.length) continue;
for (const name of promptNames) {
const prompt = prompts.find(p => p.name === name);
if (prompt && model) {
await this.update(
prompt.name,
{ model, modified: true },
{ model: { not: model } }
);
}
}
}
} else {
this.logger.log('No scenarios enabled, using default prompts.');
const prompts = Object.values(Scenario).flat();
for (const prompt of prompts) {
await this.update(prompt, { modified: false });
}
}
}
/**
* list prompt names
* @returns prompt names
*/
async listNames() {
this.ensureInMemoryPrompts();
return Array.from(this.inMemoryPrompts.keys());
}
async list() {
this.ensureInMemoryPrompts();
return Array.from(this.inMemoryPrompts.values())
.map(prompt => ({
name: prompt.name,
action: prompt.action ?? null,
model: prompt.model,
config: prompt.config ? structuredClone(prompt.config) : null,
messages: prompt.messages.map(message => ({
role: message.role,
content: message.content,
params: message.params ?? null,
})),
}))
.sort((a, b) => {
if (a.action === null && b.action !== null) return -1;
if (a.action !== null && b.action === null) return 1;
return (a.action ?? '').localeCompare(b.action ?? '');
});
}
/**
* get prompt messages by prompt name
* @param name prompt name
* @returns prompt messages
*/
async get(name: string): Promise<ChatPrompt | null> {
this.ensureInMemoryPrompts();
// skip cache in dev mode to ensure the latest prompt is always fetched
if (!env.dev) {
const cached = this.cache.get(name);
if (cached) return cached;
async get(name: string): Promise<ResolvedPrompt | null> {
const compatPrompt = this.lookupCompatPrompt(name);
if (compatPrompt) {
return this.describeCompatPrompt(this.clonePrompt(compatPrompt));
}
const prompt = this.inMemoryPrompts.get(name);
if (!prompt) return null;
const builtInPromptSpec = this.lookupBuiltInPromptSpec(name);
if (!builtInPromptSpec) return null;
const messages = PromptMessageSchema.array().safeParse(prompt.messages);
const config = PromptConfigSchema.safeParse(prompt.config);
if (messages.success && config.success) {
const chatPrompt = ChatPrompt.createFromPrompt({
...this.clonePrompt(prompt),
action: prompt.action ?? null,
optionalModels: prompt.optionalModels ?? [],
config: config.data,
messages: messages.data,
});
this.cache.set(name, chatPrompt);
return chatPrompt;
}
return this.describeBuiltInPromptSpec(builtInPromptSpec);
}
finish(
prompt: ResolvedPrompt,
params: PromptParams,
sessionId?: string
): PromptMessage[] {
const rendered =
prompt.source === 'built_in'
? renderBuiltInPromptNative({
name: prompt.name,
renderParams: params,
})
: renderPromptNative({
messages: this.requireCompatMessages(prompt),
templateParams: prompt.params,
renderParams: params,
});
this.logWarnings(rendered.warnings, sessionId);
return rendered.messages;
}
renderSession(
prompt: ResolvedPrompt,
turns: PromptMessage[],
params: PromptParams,
maxTokenSize = prompt.config?.maxTokens || 128 * 1024,
sessionId?: string
): PromptMessage[] {
const rendered =
prompt.source === 'built_in'
? renderBuiltInPromptSessionNative({
name: prompt.name,
turns,
renderParams: params,
maxTokenSize,
})
: renderPromptSessionNative({
prompt: {
action: prompt.action,
model: prompt.model,
promptTokens: this.countCompatPromptTokens(prompt),
templateParams: prompt.params,
messages: this.requireCompatMessages(prompt),
},
turns,
renderParams: params,
maxTokenSize,
});
this.logWarnings(rendered.warnings, sessionId);
return rendered.messages;
}
protected lookupCompatPrompt(_name: string): Prompt | null {
return null;
}
async set(
name: string,
model: string,
messages: PromptMessage[],
config?: PromptConfig | null,
extraConfig?: { optionalModels: string[] }
) {
this.ensureInMemoryPrompts();
const existing = this.inMemoryPrompts.get(name);
const mergedOptionalModels = existing?.optionalModels
? [...existing.optionalModels, ...(extraConfig?.optionalModels ?? [])]
: extraConfig?.optionalModels;
const inMemoryConfig = (!!config && structuredClone(config)) || undefined;
const dbConfig = this.toDbConfig(config);
this.inMemoryPrompts.set(name, {
name,
model,
action: existing?.action,
optionalModels: mergedOptionalModels,
config: inMemoryConfig,
messages: this.cloneMessages(messages),
});
this.cache.delete(name);
try {
return await this.db.aiPrompt
.upsert({
where: { name },
create: {
name,
action: existing?.action,
model,
optionalModels: mergedOptionalModels,
config: dbConfig,
messages: {
create: messages.map((m, idx) => ({
idx,
...m,
attachments: m.attachments || undefined,
params: m.params || undefined,
})),
},
},
update: {
model,
optionalModels: mergedOptionalModels,
config: dbConfig,
updatedAt: new Date(),
messages: {
deleteMany: {},
create: messages.map((m, idx) => ({
idx,
...m,
attachments: m.attachments || undefined,
params: m.params || undefined,
})),
},
},
})
.then(ret => ret.id);
} catch (error) {
this.logger.warn(
`Compat prompt upsert failed for "${name}": ${this.stringifyError(error)}`
);
return -1;
}
protected lookupBuiltInPromptSpec(name: string): PromptSpec | null {
const spec = getBuiltInPromptSpecNative(name);
return spec ? this.clonePromptSpec(spec) : null;
}
@Transactional()
async update(
name: string,
data: {
messages?: PromptMessage[];
model?: string;
modified?: boolean;
config?: PromptConfig | null;
},
where?: Prisma.AiPromptWhereInput
) {
this.ensureInMemoryPrompts();
const { config, messages, model, modified } = data;
const current = this.inMemoryPrompts.get(name);
if (current) {
const next = this.clonePrompt(current);
if (model !== undefined) {
next.model = model;
}
if (config === null) {
next.config = undefined;
} else if (config !== undefined) {
next.config = structuredClone(config);
}
if (messages) {
next.messages = this.cloneMessages(messages);
}
this.inMemoryPrompts.set(name, next);
this.cache.delete(name);
}
try {
const existing = await this.db.aiPrompt
.count({ where: { ...where, name } })
.then(count => count > 0);
if (existing) {
await this.db.aiPrompt.update({
where: { name },
data: {
config: this.toDbConfig(config),
updatedAt: new Date(),
modified,
model,
messages: messages
? {
// cleanup old messages
deleteMany: {},
create: messages.map((m, idx) => ({
idx,
...m,
attachments: m.attachments || undefined,
params: m.params || undefined,
})),
}
: undefined,
},
});
}
} catch (error) {
this.logger.warn(
`Compat prompt update failed for "${name}": ${this.stringifyError(error)}`
);
}
}
async delete(name: string) {
this.inMemoryPrompts.delete(name);
this.cache.delete(name);
try {
const { id } = await this.db.aiPrompt.delete({ where: { name } });
return id;
} catch (error) {
this.logger.warn(
`Compat prompt delete failed for "${name}": ${this.stringifyError(error)}`
);
return -1;
}
}
private resetInMemoryPrompts() {
this.cache.clear();
this.inMemoryPrompts.clear();
for (const prompt of prompts) {
this.inMemoryPrompts.set(prompt.name, this.clonePrompt(prompt));
}
}
private ensureInMemoryPrompts() {
if (!this.inMemoryPrompts.size) {
this.resetInMemoryPrompts();
}
}
private toDbConfig(
config: PromptConfig | null | undefined
): Prisma.InputJsonValue | Prisma.NullableJsonNullValueInput | undefined {
if (config === null) return Prisma.DbNull;
if (config === undefined) return undefined;
return config as Prisma.InputJsonValue;
}
private cloneMessages(messages: PromptMessage[]) {
protected cloneMessages(messages: PromptMessage[]) {
return messages.map(message => ({
...message,
attachments: message.attachments ? [...message.attachments] : undefined,
params: message.params ? structuredClone(message.params) : undefined,
responseFormat: message.responseFormat
? structuredClone(message.responseFormat)
: undefined,
}));
}
private clonePrompt(prompt: Prompt): Prompt {
protected clonePrompt(prompt: Prompt): Prompt {
return {
...prompt,
optionalModels: prompt.optionalModels
@@ -328,7 +115,93 @@ export class PromptService implements OnApplicationBootstrap {
};
}
private stringifyError(error: unknown) {
return error instanceof Error ? error.message : String(error);
protected clonePromptSpec(spec: PromptSpec): PromptSpec {
return {
...spec,
optionalModels: spec.optionalModels
? [...spec.optionalModels]
: undefined,
config: spec.config ? structuredClone(spec.config) : undefined,
params: spec.params ? structuredClone(spec.params) : undefined,
messages: spec.messages.map(message => ({ ...message })),
};
}
private describeBuiltInPromptSpec(spec: PromptSpec): ResolvedPrompt {
const params = this.normalizePromptSpecParams(spec.params);
return {
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels ?? [],
config: spec.config ? structuredClone(spec.config) : undefined,
paramKeys: Object.keys(params),
params,
source: 'built_in',
};
}
private describeCompatPrompt(prompt: Prompt): ResolvedPrompt {
const metadata = collectPromptMetadataNative({ messages: prompt.messages });
return {
name: prompt.name,
action: prompt.action,
model: prompt.model,
optionalModels: prompt.optionalModels ?? [],
config: prompt.config ? structuredClone(prompt.config) : undefined,
paramKeys: metadata.paramKeys,
params: metadata.templateParams,
source: 'compat',
messages: prompt.messages,
};
}
private normalizePromptSpecParams(
params?: PromptSpec['params']
): PromptParams {
if (!params) return {};
return Object.fromEntries(
Object.entries(params).map(([key, value]) => {
if (value.enum?.length) {
const normalized = value.default
? [
value.default,
...value.enum.filter(option => option !== value.default),
]
: [...value.enum];
return [key, normalized];
}
return [key, value.default ?? ''];
})
);
}
private countCompatPromptTokens(prompt: ResolvedPrompt): number {
return countPromptTokensNative({
model: prompt.model,
messages: this.requireCompatMessages(prompt).map(message => ({
content: message.content,
})),
}).tokens;
}
private requireCompatMessages(prompt: ResolvedPrompt): PromptMessage[] {
if (prompt.source === 'compat' && prompt.messages) {
return this.cloneMessages(prompt.messages);
}
throw new Error(`Prompt ${prompt.name} does not expose compat messages`);
}
private logWarnings(warnings: string[], sessionId?: string) {
if (!sessionId) {
return;
}
for (const warning of warnings) {
this.logger.warn(`${warning} in session ${sessionId}`);
}
}
}
@@ -0,0 +1,46 @@
import type {
PromptConfig,
PromptMessage,
PromptParams,
} from '../providers/types';
export type Prompt = {
name: string;
model: string;
optionalModels?: string[];
action?: string;
messages: PromptMessage[];
config?: PromptConfig;
};
export type ResolvedPrompt = {
name: string;
model: string;
optionalModels: string[];
action?: string;
config?: PromptConfig;
paramKeys: string[];
params: PromptParams;
source: 'built_in' | 'compat';
messages?: PromptMessage[];
};
type PromptParamSpec = {
default?: string;
enum?: string[];
};
type PromptSpecMessage = {
role: 'system' | 'assistant' | 'user';
template: string;
};
export type PromptSpec = {
name: string;
action?: string;
model: string;
optionalModels?: string[];
config?: PromptConfig;
params?: Record<string, PromptParamSpec>;
messages: PromptSpecMessage[];
};
@@ -1,24 +1,15 @@
import { CopilotProviderSideError, UserFriendlyError } from '../../../../base';
import {
CopilotProviderSideError,
metrics,
UserFriendlyError,
} from '../../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
type LlmBackendConfig,
llmResolveRequestIntentOptions,
} from '../../../../native';
import type { NodeTextMiddleware } from '../../config';
import type { CopilotToolSet } from '../../tools';
import { buildNativeRequest, NativeProviderAdapter } from '../native';
import { CopilotProvider } from '../provider';
import type {
CopilotChatOptions,
ModelConditions,
PromptMessage,
StreamObject,
} from '../types';
import { CopilotProviderType, ModelOutputType } from '../types';
import { hasProviderModelBehaviorFlag } from '../provider-model-runtime';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import {
getGoogleAuth,
getVertexAnthropicBaseUrl,
@@ -26,6 +17,51 @@ import {
} from '../utils';
export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
protected resolveModelBackendKind() {
return this.type === CopilotProviderType.AnthropicVertex
? ('anthropic_vertex' as const)
: ('anthropic' as const);
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveRequestOptions: async context => {
const requestIntent = await llmResolveRequestIntentOptions({
protocol: context.protocol,
backendConfig: context.backendConfig,
reasoning: {
enabled: context.options.reasoning,
supported: hasProviderModelBehaviorFlag(
context.model,
'reasoning_budget_12000'
),
budgetTokens: hasProviderModelBehaviorFlag(
context.model,
'reasoning_budget_12000'
)
? 12000
: undefined,
},
});
return {
attachmentCapability: this.getAttachCapability(
context.model,
context.outputType
),
reasoning: requestIntent.reasoning,
};
},
},
structured: false,
embedding: false,
rerank: false,
};
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
@@ -37,198 +73,28 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
});
}
private async createNativeConfig(): Promise<NativeLlmBackendConfig> {
private async createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig> {
const config = this.getConfig(execution);
if (this.type === CopilotProviderType.AnthropicVertex) {
const config = this.config as VertexAnthropicProviderConfig;
const auth = await getGoogleAuth(config, 'anthropic');
const vertexConfig = config as VertexAnthropicProviderConfig;
const auth = await getGoogleAuth(vertexConfig, 'anthropic');
const { Authorization: authHeader } = auth.headers();
const token = authHeader.replace(/^Bearer\s+/i, '');
const baseUrl = getVertexAnthropicBaseUrl(config) || auth.baseUrl;
const baseUrl = getVertexAnthropicBaseUrl(vertexConfig) || auth.baseUrl;
return {
base_url: baseUrl || '',
auth_token: token,
request_layer: 'vertex_anthropic',
headers: { Authorization: authHeader },
};
}
const config = this.config as { apiKey: string; baseURL?: string };
const baseUrl = config.baseURL || 'https://api.anthropic.com/v1';
const officialConfig = config as { apiKey: string; baseURL?: string };
const baseUrl = officialConfig.baseURL || 'https://api.anthropic.com/v1';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: config.apiKey,
auth_token: officialConfig.apiKey,
};
}
private createAdapter(
backendConfig: NativeLlmBackendConfig,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream('anthropic', backendConfig, request, signal),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
);
}
private getReasoning(
options: NonNullable<CopilotChatOptions>,
model: string
): Record<string, unknown> | undefined {
if (options.reasoning && this.isReasoningModel(model)) {
return { budget_tokens: 12000, include_thought: true };
}
return undefined;
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const reasoning = this.getReasoning(options, model.id);
const cap = this.getAttachCapability(model, ModelOutputType.Text);
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
attachmentCapability: cap,
reasoning,
middleware,
});
const adapter = this.createAdapter(
backendConfig,
tools,
middleware.node?.text
);
return await adapter.text(request, options.signal, messages);
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const cap = this.getAttachCapability(model, ModelOutputType.Text);
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
attachmentCapability: cap,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createAdapter(
backendConfig,
tools,
middleware.node?.text
);
for await (const chunk of adapter.streamText(
request,
options.signal,
messages
)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<StreamObject> {
const fullCond = { ...cond, outputType: ModelOutputType.Object };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_object_stream_calls')
.add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const cap = this.getAttachCapability(model, ModelOutputType.Object);
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
attachmentCapability: cap,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createAdapter(
backendConfig,
tools,
middleware.node?.text
);
for await (const chunk of adapter.streamObject(
request,
options.signal,
messages
)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_object_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
private isReasoningModel(model: string) {
// claude 3.5 sonnet doesn't support reasoning config
return model.includes('sonnet') && !model.startsWith('claude-3-5-sonnet');
}
}
@@ -1,7 +1,5 @@
import z from 'zod';
import { IMAGE_ATTACHMENT_CAPABILITY } from '../attachments';
import { CopilotProviderType, ModelInputType, ModelOutputType } from '../types';
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { AnthropicProvider } from './anthropic';
export type AnthropicOfficialConfig = {
@@ -9,74 +7,10 @@ export type AnthropicOfficialConfig = {
baseURL?: string;
};
const ModelListSchema = z.object({
data: z.array(z.object({ id: z.string() })),
});
export class AnthropicOfficialProvider extends AnthropicProvider<AnthropicOfficialConfig> {
override readonly type = CopilotProviderType.Anthropic;
override readonly models = [
{
name: 'Claude Opus 4',
id: 'claude-opus-4-20250514',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
attachments: IMAGE_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Claude Sonnet 4',
id: 'claude-sonnet-4-5-20250929',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
attachments: IMAGE_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Claude Sonnet 4',
id: 'claude-sonnet-4-20250514',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
attachments: IMAGE_ATTACHMENT_CAPABILITY,
},
],
},
];
override configured(): boolean {
return !!this.config.apiKey;
}
override setup() {
super.setup();
}
override async refreshOnlineModels() {
try {
const baseUrl = this.config.baseURL || 'https://api.anthropic.com/v1';
if (baseUrl && !this.onlineModelList.length) {
const { data } = await fetch(`${baseUrl}/models`, {
headers: {
'x-api-key': this.config.apiKey,
'anthropic-version': '2023-06-01',
'Content-Type': 'application/json',
},
})
.then(r => r.json())
.then(r => ModelListSchema.parse(r));
this.onlineModelList = data.map(model => model.id);
}
} catch (e) {
this.logger.error('Failed to fetch available models', e);
}
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
}
@@ -1,11 +1,6 @@
import { IMAGE_ATTACHMENT_CAPABILITY } from '../attachments';
import { CopilotProviderType, ModelInputType, ModelOutputType } from '../types';
import {
getGoogleAuth,
getVertexAnthropicBaseUrl,
VertexModelListSchema,
type VertexProviderConfig,
} from '../utils';
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { getVertexAnthropicBaseUrl, type VertexProviderConfig } from '../utils';
import { AnthropicProvider } from './anthropic';
export type AnthropicVertexConfig = VertexProviderConfig;
@@ -13,67 +8,9 @@ export type AnthropicVertexConfig = VertexProviderConfig;
export class AnthropicVertexProvider extends AnthropicProvider<AnthropicVertexConfig> {
override readonly type = CopilotProviderType.AnthropicVertex;
override readonly models = [
{
name: 'Claude Opus 4',
id: 'claude-opus-4@20250514',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
attachments: IMAGE_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Claude Sonnet 4.5',
id: 'claude-sonnet-4-5@20250929',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
attachments: IMAGE_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Claude Sonnet 4',
id: 'claude-sonnet-4@20250514',
capabilities: [
{
input: [ModelInputType.Text, ModelInputType.Image],
output: [ModelOutputType.Text, ModelOutputType.Object],
attachments: IMAGE_ATTACHMENT_CAPABILITY,
},
],
},
];
override configured(): boolean {
if (!this.config.location || !this.config.googleAuthOptions) return false;
return !!this.config.project || !!getVertexAnthropicBaseUrl(this.config);
}
override async refreshOnlineModels() {
try {
const { baseUrl, headers } = await getGoogleAuth(
this.config,
'anthropic'
);
if (baseUrl && !this.onlineModelList.length) {
const { publisherModels } = await fetch(`${baseUrl}/models`, {
headers: headers(),
})
.then(r => r.json())
.then(r => VertexModelListSchema.parse(r));
this.onlineModelList = publisherModels.map(
model =>
model.name.replace('publishers/anthropic/models/', '') +
(model.versionId !== 'default' ? `@${model.versionId}` : '')
);
}
} catch (e) {
this.logger.error('Failed to fetch available models', e);
}
override configured(execution?: CopilotProviderExecution): boolean {
const config = this.getConfig(execution);
if (!config.location || !config.googleAuthOptions) return false;
return !!config.project || !!getVertexAnthropicBaseUrl(config);
}
}
@@ -1,11 +1,8 @@
import type {
ModelAttachmentCapability,
PromptAttachment,
PromptAttachmentKind,
PromptAttachmentSourceKind,
PromptMessage,
} from './types';
import { inferMimeType } from './utils';
export const IMAGE_ATTACHMENT_CAPABILITY: ModelAttachmentCapability = {
kinds: ['image'],
@@ -19,75 +16,6 @@ export const GEMINI_ATTACHMENT_CAPABILITY: ModelAttachmentCapability = {
allowRemoteUrls: true,
};
export type CanonicalPromptAttachment = {
kind: PromptAttachmentKind;
sourceKind: PromptAttachmentSourceKind;
mediaType?: string;
source: Record<string, unknown>;
isRemote: boolean;
};
function parseDataUrl(url: string) {
if (!url.startsWith('data:')) {
return null;
}
const commaIndex = url.indexOf(',');
if (commaIndex === -1) {
return null;
}
const meta = url.slice(5, commaIndex);
const payload = url.slice(commaIndex + 1);
const parts = meta.split(';');
const mediaType = parts[0] || 'text/plain;charset=US-ASCII';
const isBase64 = parts.includes('base64');
return {
mediaType,
data: isBase64
? payload
: Buffer.from(decodeURIComponent(payload), 'utf8').toString('base64'),
};
}
function attachmentTypeFromMediaType(mediaType: string): PromptAttachmentKind {
if (mediaType.startsWith('image/')) {
return 'image';
}
if (mediaType.startsWith('audio/')) {
return 'audio';
}
return 'file';
}
function attachmentKindFromHintOrMediaType(
hint: PromptAttachmentKind | undefined,
mediaType: string | undefined
): PromptAttachmentKind {
if (hint) return hint;
return attachmentTypeFromMediaType(mediaType || '');
}
function toBase64Data(data: string, encoding: 'base64' | 'utf8' = 'base64') {
return encoding === 'base64'
? data
: Buffer.from(data, 'utf8').toString('base64');
}
function appendAttachMetadata(
source: Record<string, unknown>,
attachment: Exclude<PromptAttachment, string> & Record<string, unknown>
) {
if (attachment.fileName) {
source.file_name = attachment.fileName;
}
if (attachment.providerHint) {
source.provider_hint = attachment.providerHint;
}
return source;
}
export function promptAttachmentHasSource(
attachment: PromptAttachment
): boolean {
@@ -110,124 +38,36 @@ export function promptAttachmentHasSource(
}
}
export async function canonicalizePromptAttachment(
export function applyPromptAttachmentMimeTypeHintForNative(
attachment: PromptAttachment,
message: Pick<PromptMessage, 'params'>
): Promise<CanonicalPromptAttachment> {
): PromptAttachment {
const fallbackMimeType =
typeof message.params?.mimetype === 'string'
? message.params.mimetype
: undefined;
if (typeof attachment === 'string') {
const dataUrl = parseDataUrl(attachment);
const mediaType =
fallbackMimeType ??
dataUrl?.mediaType ??
(await inferMimeType(attachment));
const kind = attachmentKindFromHintOrMediaType(undefined, mediaType);
if (dataUrl) {
return {
kind,
sourceKind: 'data',
mediaType,
isRemote: false,
source: {
media_type: mediaType || dataUrl.mediaType,
data: dataUrl.data,
},
};
}
return {
kind,
sourceKind: 'url',
mediaType,
isRemote: /^https?:\/\//.test(attachment),
source: { url: attachment, media_type: mediaType },
};
if (attachment.startsWith('data:')) return attachment;
return fallbackMimeType
? { attachment, mimeType: fallbackMimeType }
: attachment;
}
if ('attachment' in attachment) {
return await canonicalizePromptAttachment(
{
kind: 'url',
url: attachment.attachment,
mimeType: attachment.mimeType,
},
message
);
if (attachment.mimeType || !fallbackMimeType) return attachment;
return { ...attachment, mimeType: fallbackMimeType };
}
if (attachment.kind === 'url') {
const dataUrl = parseDataUrl(attachment.url);
const mediaType =
attachment.mimeType ??
fallbackMimeType ??
dataUrl?.mediaType ??
(await inferMimeType(attachment.url));
const kind = attachmentKindFromHintOrMediaType(
attachment.providerHint?.kind,
mediaType
);
if (dataUrl) {
return {
kind,
sourceKind: 'data',
mediaType,
isRemote: false,
source: appendAttachMetadata(
{ media_type: mediaType || dataUrl.mediaType, data: dataUrl.data },
attachment
),
};
}
if (attachment.kind !== 'url') return attachment;
return {
kind,
sourceKind: 'url',
mediaType,
isRemote: /^https?:\/\//.test(attachment.url),
source: appendAttachMetadata(
{ url: attachment.url, media_type: mediaType },
attachment
),
};
if (
attachment.url.startsWith('data:') ||
attachment.mimeType ||
!fallbackMimeType
) {
return attachment;
}
if (attachment.kind === 'data' || attachment.kind === 'bytes') {
return {
kind: attachmentKindFromHintOrMediaType(
attachment.providerHint?.kind,
attachment.mimeType
),
sourceKind: attachment.kind,
mediaType: attachment.mimeType,
isRemote: false,
source: appendAttachMetadata(
{
media_type: attachment.mimeType,
data: toBase64Data(
attachment.data,
attachment.kind === 'data' ? attachment.encoding : 'base64'
),
},
attachment
),
};
}
return {
kind: attachmentKindFromHintOrMediaType(
attachment.providerHint?.kind,
attachment.mimeType
),
sourceKind: 'file_handle',
mediaType: attachment.mimeType,
isRemote: false,
source: appendAttachMetadata(
{ file_handle: attachment.fileHandle, media_type: attachment.mimeType },
attachment
),
};
return { ...attachment, mimeType: fallbackMimeType };
}
@@ -1,34 +1,12 @@
import {
CopilotProviderSideError,
metrics,
UserFriendlyError,
} from '../../../base';
import {
llmDispatchStream,
llmRerankDispatch,
type NativeLlmBackendConfig,
type NativeLlmRequest,
type NativeLlmRerankRequest,
type NativeLlmRerankResponse,
} from '../../../native';
import type { NodeTextMiddleware } from '../config';
import type { CopilotTool, CopilotToolSet } from '../tools';
import {
buildNativeRequest,
buildNativeRerankRequest,
NativeProviderAdapter,
} from './native';
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import type { CopilotTool } from '../tools';
import { CopilotProvider } from './provider';
import type {
CopilotChatOptions,
CopilotChatTools,
CopilotProviderModel,
CopilotRerankRequest,
ModelConditions,
PromptMessage,
StreamObject,
} from './types';
import { CopilotProviderType, ModelInputType, ModelOutputType } from './types';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { type CopilotChatTools, CopilotProviderType } from './types';
export type CloudflareWorkersAIConfig = {
apiToken: string;
@@ -36,77 +14,17 @@ export type CloudflareWorkersAIConfig = {
baseURL?: string;
};
function rerankOnlyModel(
id: string,
name: string,
defaultForOutputType = false
): CopilotProviderModel {
return {
name,
id,
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Rerank],
...(defaultForOutputType ? { defaultForOutputType } : {}),
},
],
};
}
function chatAndRerankModel(
id: string,
name: string,
defaultForRerank = false
): CopilotProviderModel {
return {
name,
id,
capabilities: [
{
input: [ModelInputType.Text],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Rerank,
],
...(defaultForRerank ? { defaultForOutputType: true } : {}),
},
],
};
}
export class CloudflareWorkersAIProvider extends CopilotProvider<CloudflareWorkersAIConfig> {
override readonly type = CopilotProviderType.CloudflareWorkersAi;
override readonly models = [
rerankOnlyModel('@cf/baai/bge-reranker-base', 'BGE Reranker Base', true),
chatAndRerankModel('@cf/moonshotai/kimi-k2.5', 'Kimi K2.5'),
chatAndRerankModel(
'@cf/ibm-granite/granite-4.0-h-micro',
'Granite 4.0 H Micro'
),
chatAndRerankModel(
'@cf/aisingapore/gemma-sea-lion-v4-27b-it',
'Gemma Sea Lion V4 27B IT'
),
chatAndRerankModel(
'@cf/nvidia/nemotron-3-120b-a12b',
'Nemotron 3 120B A12B'
),
chatAndRerankModel('@cf/zai-org/glm-4.7-flash', 'GLM 4.7 Flash'),
chatAndRerankModel('@cf/qwen/qwen3-30b-a3b-fp8', 'Qwen3 30B A3B FP8'),
];
override configured(): boolean {
return (
!!this.config.apiToken &&
(!!this.config.accountId || !!this.config.baseURL)
);
protected resolveModelBackendKind() {
return 'cloudflare_workers_ai' as const;
}
override async refreshOnlineModels() {}
override configured(execution?: CopilotProviderExecution): boolean {
const config = this.getConfig(execution);
return !!config.apiToken && (!!config.accountId || !!config.baseURL);
}
override getProviderSpecificTools(
toolName: CopilotChatTools,
_model: string
@@ -128,178 +46,31 @@ export class CloudflareWorkersAIProvider extends CopilotProvider<CloudflareWorke
});
}
private createNativeConfig(): NativeLlmBackendConfig {
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
const config = this.getConfig(execution);
return {
base_url: this.resolveBaseUrl(),
auth_token: this.config.apiToken,
request_layer: 'cloudflare_workers_ai',
base_url: this.resolveBaseUrl(execution),
auth_token: config.apiToken,
};
}
private createNativeDispatch(
backendConfig: NativeLlmBackendConfig,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream('openai_chat', backendConfig, request, signal),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
);
}
private createNativeRerankDispatch(backendConfig: NativeLlmBackendConfig) {
return (
request: NativeLlmRerankRequest
): Promise<NativeLlmRerankResponse> =>
llmRerankDispatch('openai_chat', backendConfig, request);
}
private resolveBaseUrl() {
if (this.config.baseURL) {
return this.config.baseURL.replace(/\/v1\/?$/, '').replace(/\/$/, '');
private resolveBaseUrl(execution?: CopilotProviderExecution) {
const config = this.getConfig(execution);
if (config.baseURL) {
return config.baseURL.replace(/\/v1\/?$/, '').replace(/\/$/, '');
}
const accountId = this.config.accountId ?? '';
const accountId = config.accountId ?? '';
return `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai`;
}
override async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const normalizedCond = await this.checkParams({
messages,
cond: { ...cond, outputType: ModelOutputType.Text },
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
middleware,
});
return await this.createNativeDispatch(
this.createNativeConfig(),
tools,
middleware.node?.text
).text(request, options.signal, messages);
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<string> {
const normalizedCond = await this.checkParams({
messages,
cond: { ...cond, outputType: ModelOutputType.Text },
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
middleware,
});
for await (const chunk of this.createNativeDispatch(
this.createNativeConfig(),
tools,
middleware.node?.text
).streamText(request, options.signal, messages)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<StreamObject> {
const normalizedCond = await this.checkParams({
messages,
cond: { ...cond, outputType: ModelOutputType.Object },
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_object_stream_calls')
.add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
middleware,
});
for await (const chunk of this.createNativeDispatch(
this.createNativeConfig(),
tools,
middleware.node?.text
).streamObject(request, options.signal, messages)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_object_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async rerank(
cond: ModelConditions,
request: CopilotRerankRequest,
options: CopilotChatOptions = {}
): Promise<number[]> {
const normalizedCond = await this.checkParams({
messages: [],
cond: { ...cond, outputType: ModelOutputType.Rerank },
options,
});
const model = this.selectModel(normalizedCond);
try {
const response = await this.createNativeRerankDispatch(
this.createNativeConfig()
)(buildNativeRerankRequest(model.id, request));
return response.scores;
} catch (e: any) {
throw this.handleError(e);
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
structured: false,
embedding: false,
};
}
}
@@ -1,39 +1,76 @@
import { Injectable, Logger } from '@nestjs/common';
import { Config } from '../../../base';
import { ServerFeature, ServerService } from '../../../core';
import type { RequiredStructuredOutputContract } from '../runtime/contracts';
import { getProviderRuntimeHost } from '../runtime/provider-runtime-context';
import type { CopilotProvider } from './provider';
import {
buildProviderRegistry,
type NormalizedCopilotProviderProfile,
resolveModel,
stripProviderPrefix,
} from './provider-registry';
import { CopilotProviderType, ModelFullConditions } from './types';
import type {
CopilotProviderExecution,
PreparedNativeEmbeddingExecution,
PreparedNativeExecution,
PreparedNativeImageExecution,
PreparedNativeRerankExecution,
PreparedNativeStructuredExecution,
} from './provider-runtime-contract';
import { CopilotProviderRegistryService } from './registry-service';
import {
type CopilotChatOptions,
type CopilotEmbeddingOptions,
type CopilotImageOptions,
CopilotProviderType,
type CopilotRerankRequest,
type CopilotStructuredOptions,
ModelFullConditions,
ModelOutputType,
type PromptMessage,
} from './types';
function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
return (
value !== null &&
value !== undefined &&
typeof (value as AsyncIterable<unknown>)[Symbol.asyncIterator] ===
'function'
);
}
export type ResolvedCopilotProvider = {
providerId: string;
provider: CopilotProvider;
execution: CopilotProviderExecution;
profile: NormalizedCopilotProviderProfile;
rawModelId?: string;
modelId?: string;
explicitProviderId?: string;
prepared?: PreparedNativeExecution;
preparedStructured?: PreparedNativeStructuredExecution;
preparedEmbedding?: PreparedNativeEmbeddingExecution;
preparedRerank?: PreparedNativeRerankExecution;
preparedImage?: PreparedNativeImageExecution;
};
type RoutePreparationResult = Partial<
Pick<
ResolvedCopilotProvider,
| 'prepared'
| 'preparedStructured'
| 'preparedEmbedding'
| 'preparedRerank'
| 'preparedImage'
| 'modelId'
>
>;
@Injectable()
export class CopilotProviderFactory {
constructor(
private readonly server: ServerService,
private readonly config: Config
private readonly registries: CopilotProviderRegistryService
) {}
private readonly logger = new Logger(CopilotProviderFactory.name);
readonly #providers = new Map<string, CopilotProvider>();
readonly #boundProviders = new Map<string, CopilotProvider>();
readonly #providerIdsByType = new Map<CopilotProviderType, Set<string>>();
private getRegistry() {
return buildProviderRegistry(this.config.copilot.providers);
return this.registries.getRegistry();
}
private getPreferredProviderIds(type?: CopilotProviderType) {
@@ -50,91 +87,235 @@ export class CopilotProviderFactory {
return { ...cond, modelId };
}
private normalizeMethodArgs(providerId: string, args: unknown[]) {
const [first, ...rest] = args;
if (
!first ||
typeof first !== 'object' ||
Array.isArray(first) ||
!('modelId' in first)
) {
return args;
}
private filterPreparedRoutes(routes: Array<ResolvedCopilotProvider | null>) {
return routes.filter(
(route): route is ResolvedCopilotProvider => route !== null
);
}
const cond = first as Record<string, unknown>;
if (typeof cond.modelId !== 'string') return args;
private async prepareResolvedRoutes(
routes: ResolvedCopilotProvider[],
prepare: (
route: ResolvedCopilotProvider
) => Promise<RoutePreparationResult | null | undefined>
) {
const preparedRoutes = await Promise.all(
routes.map(async route => {
const prepared = await prepare(route);
return prepared ? { ...route, ...prepared } : null;
})
);
return this.filterPreparedRoutes(preparedRoutes);
}
async resolveProvider(
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider | null> {
return (await this.resolveRoutes(cond, filter))[0] ?? null;
}
async resolveRoutes(
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
this.logger.debug(
`Resolving copilot provider for output type: ${cond.outputType}`
);
const registry = this.getRegistry();
const modelId = stripProviderPrefix(registry, providerId, cond.modelId);
return [{ ...cond, modelId }, ...rest];
}
const route = resolveModel({
registry,
modelId: cond.modelId,
outputType: cond.outputType,
availableProviderIds: this.#providers.keys(),
preferredProviderIds: this.getPreferredProviderIds(filter.prefer),
});
private wrapAsyncIterable<T>(
provider: CopilotProvider,
providerId: string,
iterable: AsyncIterable<T>
): AsyncIterableIterator<T> {
const iterator = iterable[Symbol.asyncIterator]();
const resolved: ResolvedCopilotProvider[] = [];
for (const providerId of route.candidateProviderIds) {
const provider = this.#providers.get(providerId);
const profile = registry.profiles.get(providerId);
if (!provider || !profile) continue;
return {
next: value =>
provider.runWithProfile(providerId, () => iterator.next(value)),
return: value =>
provider.runWithProfile(providerId, async () => {
if (typeof iterator.return === 'function') {
return iterator.return(value as never);
}
return { done: true, value: value as T };
}),
throw: error =>
provider.runWithProfile(providerId, async () => {
if (typeof iterator.throw === 'function') {
return iterator.throw(error);
}
throw error;
}),
[Symbol.asyncIterator]() {
return this;
},
};
}
const normalizedCond = this.normalizeCond(providerId, cond);
if (
normalizedCond.modelId &&
profile.models?.length &&
!profile.models.includes(normalizedCond.modelId)
) {
continue;
}
private getBoundProvider(providerId: string, provider: CopilotProvider) {
const cached = this.#boundProviders.get(providerId);
if (cached) {
return cached;
const execution = { providerId, profile };
const matched = await provider.match(normalizedCond, execution);
if (!matched) continue;
this.logger.debug(
`Copilot provider candidate found: ${provider.type} (${providerId})`
);
resolved.push({
providerId,
provider,
execution,
profile,
rawModelId: route.rawModelId,
modelId: normalizedCond.modelId,
explicitProviderId: route.explicitProviderId,
});
}
const wrapped = new Proxy(provider, {
get: (target, prop, receiver) => {
if (prop === 'providerId') {
return providerId;
}
return resolved;
}
const value = Reflect.get(target, prop, receiver);
if (typeof value !== 'function') {
return value;
}
async prepareRoutes(
kind: 'text' | 'streamText' | 'streamObject',
cond: ModelFullConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {},
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
return await this.prepareResolvedRoutes(routes, async route => {
const prepared = await getProviderRuntimeHost(
route.provider
).prepare.chat(
kind,
{ ...cond, modelId: route.modelId },
messages,
options,
route.execution
);
const normalizedPrepared = prepared?.route ? prepared : undefined;
if (!normalizedPrepared) {
return null;
}
return (...args: unknown[]) => {
const normalizedArgs = this.normalizeMethodArgs(providerId, args);
const result = provider.runWithProfile(providerId, () =>
Reflect.apply(value, provider, normalizedArgs)
);
if (isAsyncIterable(result)) {
return this.wrapAsyncIterable(
provider,
providerId,
result as AsyncIterable<unknown>
);
}
return result;
};
},
}) as CopilotProvider;
return {
modelId: normalizedPrepared.route.model,
prepared: normalizedPrepared,
};
});
}
this.#boundProviders.set(providerId, wrapped);
return wrapped;
async prepareStructuredRoutes(
cond: ModelFullConditions,
messages: PromptMessage[],
options: CopilotStructuredOptions = {},
filter: {
prefer?: CopilotProviderType;
} = {},
responseContract?: RequiredStructuredOutputContract
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedStructured =
(await getProviderRuntimeHost(route.provider).prepare.structured(
{ ...cond, modelId: route.modelId },
messages,
options,
responseContract,
route.execution
)) ?? undefined;
if (!preparedStructured) {
return null;
}
return {
modelId: preparedStructured.route.model,
preparedStructured,
};
});
}
async prepareEmbeddingRoutes(
modelId: string,
input: string | string[],
options: CopilotEmbeddingOptions = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes({
modelId,
outputType: ModelOutputType.Embedding,
});
return await this.prepareResolvedRoutes(routes, async route => {
const preparedEmbedding =
(await getProviderRuntimeHost(route.provider).prepare.embedding(
{ modelId: route.modelId },
input,
options,
route.execution
)) ?? undefined;
if (!preparedEmbedding) {
return null;
}
return {
modelId: preparedEmbedding.route.model,
preparedEmbedding,
};
});
}
async prepareRerankRoutes(
modelId: string,
request: CopilotRerankRequest,
options: CopilotChatOptions = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes({
modelId,
outputType: ModelOutputType.Rerank,
});
return await this.prepareResolvedRoutes(routes, async route => {
const preparedRerank =
(await getProviderRuntimeHost(route.provider).prepare.rerank(
{ modelId: route.modelId },
request,
options,
route.execution
)) ?? undefined;
if (!preparedRerank) {
return null;
}
return {
modelId: preparedRerank.route.model,
preparedRerank,
};
});
}
async prepareImageRoutes(
cond: ModelFullConditions,
messages: PromptMessage[],
options: CopilotImageOptions = {},
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedImage =
(await getProviderRuntimeHost(route.provider).prepare.image(
{ ...cond, modelId: route.modelId },
messages,
options,
route.execution
)) ?? undefined;
if (!preparedImage) {
return null;
}
return {
modelId: preparedImage.route.model,
preparedImage,
};
});
}
async getProvider(
@@ -143,44 +324,7 @@ export class CopilotProviderFactory {
prefer?: CopilotProviderType;
} = {}
): Promise<CopilotProvider | null> {
this.logger.debug(
`Resolving copilot provider for output type: ${cond.outputType}`
);
const route = resolveModel({
registry: this.getRegistry(),
modelId: cond.modelId,
outputType: cond.outputType,
availableProviderIds: this.#providers.keys(),
preferredProviderIds: this.getPreferredProviderIds(filter.prefer),
});
const registry = this.getRegistry();
for (const providerId of route.candidateProviderIds) {
const provider = this.#providers.get(providerId);
if (!provider) continue;
const profile = registry.profiles.get(providerId);
const normalizedCond = this.normalizeCond(providerId, cond);
if (
normalizedCond.modelId &&
profile?.models?.length &&
!profile.models.includes(normalizedCond.modelId)
) {
continue;
}
const matched = await provider.runWithProfile(providerId, () =>
provider.match(normalizedCond)
);
if (!matched) continue;
this.logger.debug(
`Copilot provider candidate found: ${provider.type} (${providerId})`
);
return this.getBoundProvider(providerId, provider);
}
return null;
return (await this.resolveProvider(cond, filter))?.provider ?? null;
}
async getProviderByModel(
@@ -204,7 +348,6 @@ export class CopilotProviderFactory {
}
this.#providers.set(providerId, provider);
this.#boundProviders.delete(providerId);
const ids = this.#providerIdsByType.get(provider.type) ?? new Set<string>();
ids.add(providerId);
@@ -223,7 +366,6 @@ export class CopilotProviderFactory {
}
this.#providers.delete(providerId);
this.#boundProviders.delete(providerId);
const ids = this.#providerIdsByType.get(provider.type);
ids?.delete(providerId);
@@ -1,228 +1,46 @@
import {
config as falConfig,
stream as falStream,
} from '@fal-ai/serverless-client';
import { Injectable } from '@nestjs/common';
import { z, ZodType } from 'zod';
import {
CopilotPromptInvalid,
CopilotProviderSideError,
metrics,
UserFriendlyError,
} from '../../../base';
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { CopilotProvider } from './provider';
import type {
CopilotChatOptions,
CopilotImageOptions,
ModelConditions,
PromptMessage,
} from './types';
import { CopilotProviderType, ModelInputType, ModelOutputType } from './types';
import { promptAttachmentMimeType, promptAttachmentToUrl } from './utils';
CopilotProviderExecution,
ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType } from './types';
export type FalConfig = {
apiKey: string;
};
const FalImageSchema = z
.object({
url: z.string(),
seed: z.number().nullable().optional(),
content_type: z.string(),
file_name: z.string().nullable().optional(),
file_size: z.number().nullable().optional(),
width: z.number(),
height: z.number(),
})
.optional();
type FalImage = z.infer<typeof FalImageSchema>;
const FalResponseSchema = z.object({
detail: z
.union([
z.array(z.object({ type: z.string(), msg: z.string() })),
z.string(),
])
.optional(),
images: z.array(FalImageSchema).nullable().optional(),
image: FalImageSchema.nullable().optional(),
output: z.string().nullable().optional(),
});
type FalResponse = z.infer<typeof FalResponseSchema>;
const FalStreamOutputSchema = z.object({
type: z.literal('output'),
output: FalResponseSchema,
});
type FalPrompt = {
model_name?: string;
image_url?: string;
prompt?: string;
loras?: { path: string; scale?: number }[];
controlnets?: {
image_url: string;
start_percentage?: number;
end_percentage?: number;
}[];
};
@Injectable()
export class FalProvider extends CopilotProvider<FalConfig> {
override type = CopilotProviderType.FAL;
override readonly models = [
{
id: 'flux-1/schnell',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Image],
defaultForOutputType: true,
},
],
},
// image to image models
{
id: 'lcm-sd15-i2i',
capabilities: [
{
input: [ModelInputType.Image],
output: [ModelOutputType.Image],
defaultForOutputType: true,
},
],
},
{
id: 'clarity-upscaler',
capabilities: [
{
input: [ModelInputType.Image],
output: [ModelOutputType.Image],
},
],
},
{
id: 'face-to-sticker',
capabilities: [
{
input: [ModelInputType.Image],
output: [ModelOutputType.Image],
},
],
},
{
id: 'imageutils/rembg',
capabilities: [
{
input: [ModelInputType.Image],
output: [ModelOutputType.Image],
},
],
},
{
id: 'workflowutils/teed',
capabilities: [
{
input: [ModelInputType.Image],
output: [ModelOutputType.Image],
},
],
},
{
id: 'lora/image-to-image',
capabilities: [
{
input: [ModelInputType.Image],
output: [ModelOutputType.Image],
},
],
},
];
override configured(): boolean {
return !!this.config.apiKey;
protected resolveModelBackendKind() {
return 'fal' as const;
}
protected override setup() {
super.setup();
falConfig({ credentials: this.config.apiKey });
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
private extractArray<T>(value: T | T[] | undefined): T[] {
if (Array.isArray(value)) return value;
return value ? [value] : [];
}
private extractPrompt(
message?: PromptMessage,
options: CopilotImageOptions = {}
): FalPrompt {
if (!message) throw new CopilotPromptInvalid('Prompt is empty');
const { content, attachments, params } = message;
// prompt attachments require at least one
if (!content && (!Array.isArray(attachments) || !attachments.length)) {
throw new CopilotPromptInvalid('Prompt or Attachments is empty');
}
if (Array.isArray(attachments) && attachments.length > 1) {
throw new CopilotPromptInvalid('Only one attachment is allowed');
}
const lora = [
...this.extractArray(params?.lora),
...this.extractArray(options.loras),
].filter(
(v): v is { path: string; scale?: number } =>
!!v && typeof v === 'object' && typeof v.path === 'string'
);
const controlnets = this.extractArray(params?.controlnets).filter(
(v): v is { image_url: string } =>
!!v && typeof v === 'object' && typeof v.image_url === 'string'
);
private createNativeConfig(execution?: CopilotProviderExecution) {
return {
model_name: options.modelName || undefined,
image_url: attachments
?.map(v => {
const url = promptAttachmentToUrl(v);
const mediaType = promptAttachmentMimeType(
v,
typeof params?.mimetype === 'string' ? params.mimetype : undefined
);
return url && mediaType?.startsWith('image/') ? url : undefined;
})
.find(v => !!v),
prompt: content.trim(),
loras: lora.length ? lora : undefined,
controlnets: controlnets.length ? controlnets : undefined,
base_url: 'https://fal.run',
auth_token: this.getConfig(execution).apiKey,
};
}
private extractFalError(
resp: FalResponse,
message?: string
): CopilotProviderSideError {
if (Array.isArray(resp.detail) && resp.detail.length) {
const error = resp.detail[0].msg;
return new CopilotProviderSideError({
provider: this.type,
kind: resp.detail[0].type,
message: message ? `${message}: ${error}` : error,
});
} else if (typeof resp.detail === 'string') {
const error = resp.detail;
return new CopilotProviderSideError({
provider: this.type,
kind: resp.detail,
message: message ? `${message}: ${error}` : error,
});
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unknown',
message: 'No content generated',
});
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: false,
structured: false,
embedding: false,
rerank: false,
image: {},
};
}
private handleError(e: any) {
@@ -238,152 +56,4 @@ export class FalProvider extends CopilotProvider<FalConfig> {
return error;
}
}
private parseSchema<R>(schema: ZodType<R>, data: unknown): R {
const result = schema.safeParse(data);
if (result.success) return result.data;
const errors = JSON.stringify(result.error.errors);
throw new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: `Unexpected fal response: ${errors}`,
});
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const model = this.selectModel(cond);
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
// by default, image prompt assumes there is only one message
const prompt = this.extractPrompt(messages[messages.length - 1]);
const response = await fetch(`https://fal.run/fal-ai/${model.id}`, {
method: 'POST',
headers: {
Authorization: `key ${this.config.apiKey}`,
'Content-Type': 'application/json',
},
body: JSON.stringify({
...prompt,
sync_mode: true,
enable_safety_checks: false,
}),
signal: options.signal,
});
const data = this.parseSchema(FalResponseSchema, await response.json());
if (!data.output) {
throw this.extractFalError(data, 'Failed to generate text');
}
return data.output;
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions | CopilotImageOptions = {}
): AsyncIterable<string> {
const model = this.selectModel(cond);
try {
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const result = await this.text(cond, messages, options);
yield result;
} catch (e) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw e;
}
}
override async *streamImages(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotImageOptions = {}
): AsyncIterable<string> {
const model = this.selectModel({
...cond,
outputType: ModelOutputType.Image,
});
try {
metrics.ai
.counter('generate_images_stream_calls')
.add(1, this.metricLabels(model.id));
// by default, image prompt assumes there is only one message
const prompt = this.extractPrompt(
messages[messages.length - 1],
options as CopilotImageOptions
);
let data: FalResponse;
if (model.id.startsWith('workflows/')) {
const stream = await falStream(model.id, { input: prompt });
data = this.parseSchema(
FalStreamOutputSchema,
await stream.done()
).output;
} else {
const response = await fetch(`https://fal.run/fal-ai/${model.id}`, {
method: 'POST',
headers: {
Authorization: `key ${this.config.apiKey}`,
'Content-Type': 'application/json',
},
body: JSON.stringify({
...prompt,
sync_mode: true,
seed: (options as CopilotImageOptions)?.seed || 42,
enable_safety_checks: false,
}),
signal: options.signal,
});
data = this.parseSchema(FalResponseSchema, await response.json());
}
if (!data.images?.length && !data.image?.url) {
throw this.extractFalError(data, 'Failed to generate images');
}
if (data.image?.url) {
yield data.image.url;
return;
}
const imageUrls =
data.images
?.filter((image): image is NonNullable<FalImage> => !!image)
.map(image => image.url) || [];
for (const url of imageUrls) {
yield url;
if (options.signal?.aborted) {
break;
}
}
return;
} catch (e) {
metrics.ai
.counter('generate_images_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
}
@@ -1,47 +1,34 @@
import { setTimeout as delay } from 'node:timers/promises';
import { Inject } from '@nestjs/common';
import { ZodError } from 'zod';
import {
CopilotProviderSideError,
metrics,
OneMB,
readResponseBufferWithLimit,
safeFetch,
UserFriendlyError,
} from '../../../../base';
import { sniffMime } from '../../../../base/storage/providers/utils';
import {
llmDispatchStream,
llmEmbeddingDispatch,
llmStructuredDispatch,
type NativeLlmBackendConfig,
type NativeLlmEmbeddingRequest,
type NativeLlmRequest,
type NativeLlmStructuredRequest,
isInvalidStructuredOutputError,
type LlmBackendConfig,
llmResolveRequestIntentOptions,
} from '../../../../native';
import type { NodeTextMiddleware } from '../../config';
import type { CopilotToolSet } from '../../tools';
import {
buildNativeEmbeddingRequest,
buildNativeRequest,
buildNativeStructuredRequest,
NativeProviderAdapter,
parseNativeStructuredOutput,
StructuredResponseParseError,
} from '../native';
admittedAttachmentToPromptAttachment,
AttachmentAdmissionHost,
} from '../../runtime/hosts/attachment-admission';
import {
planAdmittedAttachmentMaterialization,
planHostUrlAttachmentMaterialization,
} from '../../runtime/hosts/attachment-materialization-planner';
import { AttachmentMaterializer } from '../../runtime/hosts/attachment-materializer';
import { CopilotProvider } from '../provider';
import type {
CopilotChatOptions,
CopilotEmbeddingOptions,
CopilotImageOptions,
CopilotStructuredOptions,
ModelConditions,
PromptAttachment,
PromptMessage,
StreamObject,
} from '../types';
import { ModelOutputType } from '../types';
import { hasProviderModelBehaviorFlag } from '../provider-model-runtime';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from '../provider-runtime-contract';
import type { PromptAttachment, PromptMessage } from '../types';
import { promptAttachmentMimeType, promptAttachmentToUrl } from '../utils';
export const DEFAULT_DIMENSIONS = 256;
@@ -53,38 +40,20 @@ function normalizeMimeType(mediaType?: string) {
return mediaType?.split(';', 1)[0]?.trim() || 'application/octet-stream';
}
function isYoutubeUrl(url: URL) {
const hostname = url.hostname.toLowerCase();
if (hostname === 'youtu.be') {
return /^\/[\w-]+$/.test(url.pathname);
}
if (hostname !== 'youtube.com' && hostname !== 'www.youtube.com') {
return false;
}
if (url.pathname !== '/watch') {
return false;
}
return !!url.searchParams.get('v');
}
function isGeminiFileUrl(url: URL, baseUrl: string) {
try {
const base = new URL(baseUrl);
const basePath = base.pathname.replace(/\/+$/, '');
return (
url.origin === base.origin &&
url.pathname.startsWith(`${basePath}/files/`)
);
} catch {
return false;
}
}
export abstract class GeminiProvider<T> extends CopilotProvider<T> {
protected abstract createNativeConfig(): Promise<NativeLlmBackendConfig>;
@Inject() protected readonly attachmentMaterializer!: AttachmentMaterializer;
@Inject()
protected readonly attachmentAdmissionHost?: AttachmentAdmissionHost;
protected resolveModelBackendKind() {
return this.type === 'geminiVertex'
? ('gemini_vertex' as const)
: ('gemini_api' as const);
}
protected abstract createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig>;
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
@@ -98,158 +67,27 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
}
}
protected createNativeDispatch(backendConfig: NativeLlmBackendConfig) {
return (request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream('gemini', backendConfig, request, signal);
}
protected createNativeStructuredDispatch(
backendConfig: NativeLlmBackendConfig
) {
return (request: NativeLlmStructuredRequest) =>
llmStructuredDispatch('gemini', backendConfig, request);
}
protected createNativeEmbeddingDispatch(
backendConfig: NativeLlmBackendConfig
) {
return (request: NativeLlmEmbeddingRequest) =>
llmEmbeddingDispatch('gemini', backendConfig, request);
}
protected createNativeAdapter(
backendConfig: NativeLlmBackendConfig,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
this.createNativeDispatch(backendConfig),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
private getAttachmentAdmissionHost() {
return (
this.attachmentAdmissionHost ??
new AttachmentAdmissionHost(this.attachmentMaterializer)
);
}
protected async fetchRemoteAttach(url: string, signal?: AbortSignal) {
const parsed = new URL(url);
const response = await safeFetch(
parsed,
{ method: 'GET', signal },
this.buildAttachFetchOptions(parsed)
);
if (!response.ok) {
throw new Error(
`Failed to fetch attachment: ${response.status} ${response.statusText}`
);
}
const buffer = await readResponseBufferWithLimit(
response,
GEMINI_REMOTE_ATTACHMENT_MAX_BYTES
);
const headerMimeType = normalizeMimeType(
response.headers.get('content-type') || ''
);
return {
data: buffer.toString('base64'),
mimeType: normalizeMimeType(sniffMime(buffer, headerMimeType)),
};
}
private buildAttachFetchOptions(url: URL) {
const baseOptions = { timeoutMs: 15_000, maxRedirects: 3 } as const;
if (!env.prod) {
return { ...baseOptions, allowPrivateOrigins: new Set([url.origin]) };
}
const trustedOrigins = new Set<string>();
const protocol = this.AFFiNEConfig.server.https ? 'https:' : 'http:';
const port = this.AFFiNEConfig.server.port;
const isDefaultPort =
(protocol === 'https:' && port === 443) ||
(protocol === 'http:' && port === 80);
const addHostOrigin = (host: string) => {
if (!host) return;
try {
const parsed = new URL(`${protocol}//${host}`);
if (!parsed.port && !isDefaultPort) {
parsed.port = String(port);
}
trustedOrigins.add(parsed.origin);
} catch {
// ignore invalid host config entries
}
};
if (this.AFFiNEConfig.server.externalUrl) {
try {
trustedOrigins.add(
new URL(this.AFFiNEConfig.server.externalUrl).origin
);
} catch {
// ignore invalid external URL
}
}
addHostOrigin(this.AFFiNEConfig.server.host);
for (const host of this.AFFiNEConfig.server.hosts) {
addHostOrigin(host);
}
const hostname = url.hostname.toLowerCase();
const trustedByHost = TRUSTED_ATTACHMENT_HOST_SUFFIXES.some(
suffix => hostname === suffix || hostname.endsWith(`.${suffix}`)
);
if (trustedOrigins.has(url.origin) || trustedByHost) {
return { ...baseOptions, allowPrivateOrigins: new Set([url.origin]) };
}
return baseOptions;
}
private shouldInlineRemoteAttach(url: URL, config: NativeLlmBackendConfig) {
switch (config.request_layer) {
case 'gemini_api':
if (url.protocol !== 'http:' && url.protocol !== 'https:') return false;
return !(isGeminiFileUrl(url, config.base_url) || isYoutubeUrl(url));
case 'gemini_vertex':
return false;
default:
return false;
}
}
private toInlineAttach(
attachment: PromptAttachment,
mimeType: string,
data: string
): PromptAttachment {
if (typeof attachment === 'string' || !('kind' in attachment)) {
return { kind: 'bytes', data, mimeType };
}
if (attachment.kind !== 'url') {
return attachment;
}
return {
kind: 'bytes',
data,
mimeType,
fileName: attachment.fileName,
providerHint: attachment.providerHint,
};
}
protected async prepareMessages(
messages: PromptMessage[],
backendConfig: NativeLlmBackendConfig,
signal?: AbortSignal
backendConfig: LlmBackendConfig,
options?: {
signal?: AbortSignal;
user?: string;
workspace?: string;
session?: string;
}
): Promise<PromptMessage[]> {
const prepared: PromptMessage[] = [];
for (const message of messages) {
signal?.throwIfAborted();
options?.signal?.throwIfAborted();
if (!Array.isArray(message.attachments) || !message.attachments.length) {
prepared.push(message);
continue;
@@ -258,41 +96,60 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
const attachments: PromptAttachment[] = [];
let changed = false;
for (const attachment of message.attachments) {
signal?.throwIfAborted();
options?.signal?.throwIfAborted();
const rawUrl = promptAttachmentToUrl(attachment);
if (!rawUrl || rawUrl.startsWith('data:')) {
attachments.push(attachment);
continue;
}
let parsed: URL;
try {
parsed = new URL(rawUrl);
new URL(rawUrl);
} catch {
attachments.push(attachment);
continue;
}
if (!this.shouldInlineRemoteAttach(parsed, backendConfig)) {
attachments.push(attachment);
continue;
}
const declaredMimeType = promptAttachmentMimeType(
attachment,
typeof message.params?.mimetype === 'string'
? message.params.mimetype
: undefined
);
const downloaded = await this.fetchRemoteAttach(rawUrl, signal);
attachments.push(
this.toInlineAttach(
attachment,
declaredMimeType
const referencePlan = await planHostUrlAttachmentMaterialization(
'gemini',
backendConfig,
{
attachmentId: rawUrl,
url: rawUrl,
expectedMime: declaredMimeType
? normalizeMimeType(declaredMimeType)
: downloaded.mimeType,
downloaded.data
)
: undefined,
maxSize: GEMINI_REMOTE_ATTACHMENT_MAX_BYTES,
}
);
if (referencePlan.mode === 'remote_reference') {
attachments.push(attachment);
continue;
}
const admitted =
await this.getAttachmentAdmissionHost().admitPromptAttachment(
attachment,
{
userId: options?.user ?? 'provider-runtime',
workspaceId: options?.workspace ?? 'provider-runtime',
sessionId: options?.session,
signal: options?.signal,
maxBytes: referencePlan.request.maxSize,
trustedHostSuffixes: TRUSTED_ATTACHMENT_HOST_SUFFIXES,
}
);
const materialization = planAdmittedAttachmentMaterialization(admitted);
attachments.push(
materialization.mode === 'inline'
? materialization.attachment
: admittedAttachmentToPromptAttachment(admitted)
);
changed = true;
}
@@ -310,291 +167,84 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
await delay(delayMs, undefined, signal ? { signal } : undefined);
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
prepareMessages: async context =>
await this.prepareMessages(
context.input.messages,
context.backendConfig,
context.options
),
resolveRequestOptions: async context => {
const requestIntent = await llmResolveRequestIntentOptions({
protocol: context.protocol,
backendConfig: context.backendConfig,
reasoning: {
enabled: context.options.reasoning,
supported:
hasProviderModelBehaviorFlag(
context.model,
'reasoning_medium'
) ||
hasProviderModelBehaviorFlag(context.model, 'reasoning_high'),
effort: hasProviderModelBehaviorFlag(
context.model,
'reasoning_high'
)
? 'high'
: 'medium',
includeReasoning:
hasProviderModelBehaviorFlag(
context.model,
'reasoning_medium'
) ||
hasProviderModelBehaviorFlag(context.model, 'reasoning_high'),
},
});
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const msg = await this.prepareMessages(
messages,
backendConfig,
options.signal
);
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const cap = this.getAttachCapability(model, ModelOutputType.Text);
const { request } = await buildNativeRequest({
model: model.id,
messages: msg,
options,
tools,
attachmentCapability: cap,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createNativeAdapter(
backendConfig,
tools,
middleware.node?.text
);
return await adapter.text(request, options.signal, messages);
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async structure(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotStructuredOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Structured };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const msg = await this.prepareMessages(
messages,
backendConfig,
options.signal
);
const structuredDispatch =
this.createNativeStructuredDispatch(backendConfig);
const middleware = this.getActiveProviderMiddleware();
const cap = this.getAttachCapability(model, ModelOutputType.Structured);
const { request, schema } = await buildNativeStructuredRequest({
model: model.id,
messages: msg,
options,
attachmentCapability: cap,
reasoning: this.getReasoning(options, model.id),
responseSchema: options.schema,
middleware,
});
const maxRetries = Math.max(options.maxRetries ?? 3, 0);
for (let attempt = 0; ; attempt++) {
try {
const response = await structuredDispatch(request);
const parsed = parseNativeStructuredOutput(response);
const validated = schema.parse(parsed);
return JSON.stringify(validated);
} catch (error) {
return {
attachmentCapability: this.getAttachCapability(
context.model,
context.outputType
),
include: requestIntent.include,
reasoning: requestIntent.reasoning,
};
},
},
structured: {
prepareMessages: (inputMessages, backendConfig, structuredOptions) =>
this.prepareMessages(inputMessages, backendConfig, structuredOptions),
shouldRetry: async ({ error, attempt, options: structuredOptions }) => {
const isParsingError =
error instanceof StructuredResponseParseError ||
error instanceof ZodError;
isInvalidStructuredOutputError(error) || error instanceof ZodError;
const retryableError =
isParsingError || !(error instanceof UserFriendlyError);
const maxRetries = Math.max(structuredOptions.maxRetries ?? 3, 0);
if (!retryableError || attempt >= maxRetries) {
throw error;
return false;
}
if (!isParsingError) {
await this.waitForStructuredRetry(
GEMINI_RETRY_INITIAL_DELAY_MS * 2 ** attempt,
options.signal
structuredOptions.signal
);
}
}
}
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions | CopilotImageOptions = {}
): AsyncIterable<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const preparedMessages = await this.prepareMessages(
messages,
backendConfig,
options.signal
);
const tools = await this.getTools(
options as CopilotChatOptions,
model.id
);
const middleware = this.getActiveProviderMiddleware();
const cap = this.getAttachCapability(model, ModelOutputType.Text);
const { request } = await buildNativeRequest({
model: model.id,
messages: preparedMessages,
options: options as CopilotChatOptions,
tools,
attachmentCapability: cap,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createNativeAdapter(
backendConfig,
tools,
middleware.node?.text
);
for await (const chunk of adapter.streamText(
request,
options.signal,
messages
)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<StreamObject> {
const fullCond = { ...cond, outputType: ModelOutputType.Object };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_object_stream_calls')
.add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const msg = await this.prepareMessages(
messages,
backendConfig,
options.signal
);
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const cap = this.getAttachCapability(model, ModelOutputType.Object);
const { request } = await buildNativeRequest({
model: model.id,
messages: msg,
options,
tools,
attachmentCapability: cap,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createNativeAdapter(
backendConfig,
tools,
middleware.node?.text
);
for await (const chunk of adapter.streamObject(
request,
options.signal,
messages
)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_object_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
override async embedding(
cond: ModelConditions,
messages: string | string[],
options: CopilotEmbeddingOptions = { dimensions: DEFAULT_DIMENSIONS }
): Promise<number[][]> {
const values = Array.isArray(messages) ? messages : [messages];
const fullCond = { ...cond, outputType: ModelOutputType.Embedding };
const normalizedCond = await this.checkParams({
embeddings: values,
cond: fullCond,
options,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('generate_embedding_calls')
.add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const response = await this.createNativeEmbeddingDispatch(backendConfig)(
buildNativeEmbeddingRequest({
model: model.id,
inputs: values,
dimensions: options.dimensions || DEFAULT_DIMENSIONS,
taskType: 'RETRIEVAL_DOCUMENT',
})
);
return response.embeddings;
} catch (e: any) {
metrics.ai
.counter('generate_embedding_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
protected getReasoning(
options: CopilotChatOptions | CopilotImageOptions,
model: string
): Record<string, unknown> | undefined {
if (
options &&
'reasoning' in options &&
options.reasoning &&
this.isReasoningModel(model)
) {
return this.isGemini3Model(model)
? { include_thoughts: true, thinking_level: 'high' }
: { include_thoughts: true, thinking_budget: 12000 };
}
return undefined;
}
private isGemini3Model(model: string) {
return model.startsWith('gemini-3');
}
private isReasoningModel(model: string) {
return model.startsWith('gemini-2.5') || this.isGemini3Model(model);
return true;
},
},
embedding: {
defaultDimensions: DEFAULT_DIMENSIONS,
taskType: 'RETRIEVAL_DOCUMENT',
},
rerank: false,
image: {
prepareMessages: (inputMessages, backendConfig, imageOptions) =>
this.prepareMessages(inputMessages, backendConfig, imageOptions),
},
};
}
}
@@ -1,8 +1,6 @@
import z from 'zod';
import type { NativeLlmBackendConfig } from '../../../../native';
import { GEMINI_ATTACHMENT_CAPABILITY } from '../attachments';
import { CopilotProviderType, ModelInputType, ModelOutputType } from '../types';
import type { LlmBackendConfig } from '../../../../native';
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { GeminiProvider } from './gemini';
export type GeminiGenerativeConfig = {
@@ -10,142 +8,21 @@ export type GeminiGenerativeConfig = {
baseURL?: string;
};
const ModelListSchema = z.object({
models: z.array(z.object({ name: z.string() })),
});
export class GeminiGenerativeProvider extends GeminiProvider<GeminiGenerativeConfig> {
override readonly type = CopilotProviderType.Gemini;
readonly models = [
{
name: 'Gemini 2.5 Flash',
id: 'gemini-2.5-flash',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini 2.5 Pro',
id: 'gemini-2.5-pro',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini 3.1 Pro Preview',
id: 'gemini-3.1-pro-preview',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini 3.1 Flash Lite Preview',
id: 'gemini-3.1-flash-lite-preview',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini Embedding',
id: 'gemini-embedding-001',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Embedding],
defaultForOutputType: true,
},
],
},
];
override configured(): boolean {
return !!this.config.apiKey;
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
override async refreshOnlineModels() {
try {
const baseUrl =
this.config.baseURL ||
'https://generativelanguage.googleapis.com/v1beta';
if (baseUrl && !this.onlineModelList.length) {
const { models } = await fetch(
`${baseUrl}/models?key=${this.config.apiKey}`
)
.then(r => r.json())
.then(r => ModelListSchema.parse(r));
this.onlineModelList = models.map(model =>
model.name.replace('models/', '')
);
}
} catch (e) {
this.logger.error('Failed to fetch available models', e);
}
}
protected override async createNativeConfig(): Promise<NativeLlmBackendConfig> {
protected override async createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig> {
const config = this.getConfig(execution);
return {
base_url: (
this.config.baseURL ||
'https://generativelanguage.googleapis.com/v1beta'
config.baseURL || 'https://generativelanguage.googleapis.com/v1beta'
).replace(/\/$/, ''),
auth_token: this.config.apiKey,
request_layer: 'gemini_api',
auth_token: config.apiKey,
};
}
}
@@ -1,149 +1,30 @@
import type { NativeLlmBackendConfig } from '../../../../native';
import { GEMINI_ATTACHMENT_CAPABILITY } from '../attachments';
import { CopilotProviderType, ModelInputType, ModelOutputType } from '../types';
import {
getGoogleAuth,
VertexModelListSchema,
type VertexProviderConfig,
} from '../utils';
import type { LlmBackendConfig } from '../../../../native';
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { getGoogleAuth, type VertexProviderConfig } from '../utils';
import { GeminiProvider } from './gemini';
export type GeminiVertexConfig = VertexProviderConfig;
export class GeminiVertexProvider extends GeminiProvider<GeminiVertexConfig> {
override readonly type = CopilotProviderType.GeminiVertex;
readonly models = [
{
name: 'Gemini 2.5 Flash',
id: 'gemini-2.5-flash',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini 2.5 Pro',
id: 'gemini-2.5-pro',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini 3.1 Pro Preview',
id: 'gemini-3.1-pro-preview',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini 3.1 Flash Lite Preview',
id: 'gemini-3.1-flash-lite-preview',
capabilities: [
{
input: [
ModelInputType.Text,
ModelInputType.Image,
ModelInputType.Audio,
ModelInputType.File,
],
output: [
ModelOutputType.Text,
ModelOutputType.Object,
ModelOutputType.Structured,
],
attachments: GEMINI_ATTACHMENT_CAPABILITY,
structuredAttachments: GEMINI_ATTACHMENT_CAPABILITY,
},
],
},
{
name: 'Gemini Embedding',
id: 'gemini-embedding-001',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Embedding],
defaultForOutputType: true,
},
],
},
];
override configured(): boolean {
return !!this.config.location && !!this.config.googleAuthOptions;
override configured(execution?: CopilotProviderExecution): boolean {
const config = this.getConfig(execution);
return !!config.location && !!config.googleAuthOptions;
}
protected async resolveVertexAuth(execution?: CopilotProviderExecution) {
return await getGoogleAuth(this.getConfig(execution), 'google');
}
override async refreshOnlineModels() {
try {
const { baseUrl, headers } = await this.resolveVertexAuth();
if (baseUrl && !this.onlineModelList.length) {
const { publisherModels } = await fetch(`${baseUrl}/models`, {
headers: headers(),
})
.then(r => r.json())
.then(r => VertexModelListSchema.parse(r));
this.onlineModelList = publisherModels.map(model =>
model.name.replace('publishers/google/models/', '')
);
}
} catch (e) {
this.logger.error('Failed to fetch available models', e);
}
}
protected async resolveVertexAuth() {
return await getGoogleAuth(this.config, 'google');
}
protected override async createNativeConfig(): Promise<NativeLlmBackendConfig> {
const auth = await this.resolveVertexAuth();
protected override async createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig> {
const auth = await this.resolveVertexAuth(execution);
const { Authorization: authHeader } = auth.headers();
return {
base_url: auth.baseUrl || '',
auth_token: authHeader.replace(/^Bearer\s+/i, ''),
request_layer: 'gemini_vertex',
};
}
}
@@ -1,26 +1,3 @@
import {
AnthropicOfficialProvider,
AnthropicVertexProvider,
} from './anthropic';
import { CloudflareWorkersAIProvider } from './cloudflare';
import { FalProvider } from './fal';
import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
import { MorphProvider } from './morph';
import { OpenAIProvider } from './openai';
import { PerplexityProvider } from './perplexity';
export const CopilotProviders = [
OpenAIProvider,
CloudflareWorkersAIProvider,
FalProvider,
GeminiGenerativeProvider,
GeminiVertexProvider,
PerplexityProvider,
AnthropicOfficialProvider,
AnthropicVertexProvider,
MorphProvider,
];
export {
AnthropicOfficialProvider,
AnthropicVertexProvider,
@@ -29,7 +6,10 @@ export { CloudflareWorkersAIProvider } from './cloudflare';
export { CopilotProviderFactory } from './factory';
export { FalProvider } from './fal';
export { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
export { CopilotProviderLifecycleService } from './lifecycle-service';
export { OpenAIProvider } from './openai';
export { PerplexityProvider } from './perplexity';
export type { CopilotProvider } from './provider';
export { CopilotProviders } from './provider-tokens';
export { CopilotProviderRegistryService } from './registry-service';
export * from './types';
@@ -0,0 +1,90 @@
import { Injectable, Type } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import { OnEvent } from '../../../base';
import { CopilotProviderFactory } from './factory';
import type { CopilotProvider } from './provider';
import type { CopilotProviderExecution } from './provider-runtime-contract';
import { CopilotProviders } from './provider-tokens';
import { CopilotProviderRegistryService } from './registry-service';
@Injectable()
export class CopilotProviderLifecycleService {
private readonly registeredByProvider = new WeakMap<
CopilotProvider,
Set<string>
>();
constructor(
private readonly moduleRef: ModuleRef,
private readonly factory: CopilotProviderFactory,
private readonly registries: CopilotProviderRegistryService
) {}
private getProviders(): CopilotProvider[] {
return CopilotProviders.flatMap(token => {
const provider = this.moduleRef.get(token as Type<CopilotProvider>, {
strict: false,
});
return provider ? [provider] : [];
});
}
private getRegisteredProviderIds(provider: CopilotProvider) {
const current = this.registeredByProvider.get(provider);
if (current) {
return current;
}
const next = new Set<string>();
this.registeredByProvider.set(provider, next);
return next;
}
private async syncProvider(provider: CopilotProvider) {
const registry = this.registries.getRegistry();
const configuredIds = new Set<string>();
for (const providerId of registry.byType.get(provider.type) ?? []) {
const profile = registry.profiles.get(providerId);
if (!profile) {
continue;
}
const execution: CopilotProviderExecution = { providerId, profile };
if (!provider.configured(execution)) {
this.factory.unregister(providerId, provider);
continue;
}
configuredIds.add(providerId);
this.factory.register(providerId, provider);
}
const previous = this.getRegisteredProviderIds(provider);
for (const providerId of previous) {
if (!configuredIds.has(providerId)) {
this.factory.unregister(providerId, provider);
}
}
this.registeredByProvider.set(provider, configuredIds);
}
async syncProviders() {
for (const provider of this.getProviders()) {
await this.syncProvider(provider);
}
}
@OnEvent('config.init')
async onConfigInit() {
await this.syncProviders();
}
@OnEvent('config.changed')
async onConfigChanged(event: Events['config.changed']) {
if ('copilot' in event.updates) {
await this.syncProviders();
}
}
}
@@ -1,479 +0,0 @@
import { z } from 'zod';
import type {
NativeLlmRequest,
NativeLlmStreamEvent,
NativeLlmToolDefinition,
} from '../../../native';
import type {
CopilotTool,
CopilotToolExecuteOptions,
CopilotToolSet,
} from '../tools';
export type NativeDispatchFn = (
request: NativeLlmRequest,
signal?: AbortSignal
) => AsyncIterableIterator<NativeLlmStreamEvent>;
export type NativeToolCall = {
id: string;
name: string;
args: Record<string, unknown>;
rawArgumentsText?: string;
argumentParseError?: string;
thought?: string;
};
type ToolCallState = {
name?: string;
argumentsText: string;
};
type ToolExecutionResult = {
callId: string;
name: string;
args: Record<string, unknown>;
rawArgumentsText?: string;
argumentParseError?: string;
output: unknown;
isError?: boolean;
};
type ParsedToolArguments = {
args: Record<string, unknown>;
rawArgumentsText?: string;
argumentParseError?: string;
};
export class ToolCallAccumulator {
readonly #states = new Map<string, ToolCallState>();
feedDelta(event: Extract<NativeLlmStreamEvent, { type: 'tool_call_delta' }>) {
const state = this.#states.get(event.call_id) ?? {
argumentsText: '',
};
if (event.name) {
state.name = event.name;
}
if (event.arguments_delta) {
state.argumentsText += event.arguments_delta;
}
this.#states.set(event.call_id, state);
}
complete(event: Extract<NativeLlmStreamEvent, { type: 'tool_call' }>) {
const state = this.#states.get(event.call_id);
this.#states.delete(event.call_id);
const parsed =
event.arguments_text !== undefined || event.arguments_error !== undefined
? {
args: event.arguments ?? {},
rawArgumentsText: event.arguments_text ?? state?.argumentsText,
argumentParseError: event.arguments_error,
}
: event.arguments
? this.parseArgs(event.arguments, state?.argumentsText)
: this.parseJson(state?.argumentsText ?? '{}');
return {
id: event.call_id,
name: event.name || state?.name || '',
...parsed,
thought: event.thought,
} satisfies NativeToolCall;
}
drainPending() {
const pending: NativeToolCall[] = [];
for (const [callId, state] of this.#states.entries()) {
if (!state.name) {
continue;
}
pending.push({
id: callId,
name: state.name,
...this.parseJson(state.argumentsText),
});
}
this.#states.clear();
return pending;
}
private parseJson(jsonText: string): ParsedToolArguments {
if (!jsonText.trim()) {
return { args: {} };
}
try {
return this.parseArgs(JSON.parse(jsonText), jsonText);
} catch (error) {
return {
args: {},
rawArgumentsText: jsonText,
argumentParseError:
error instanceof Error
? error.message
: 'Invalid tool arguments JSON',
};
}
}
private parseArgs(
value: unknown,
rawArgumentsText?: string
): ParsedToolArguments {
if (value && typeof value === 'object' && !Array.isArray(value)) {
return {
args: value as Record<string, unknown>,
rawArgumentsText,
};
}
return {
args: {},
rawArgumentsText,
argumentParseError: 'Tool arguments must be a JSON object',
};
}
}
export class ToolSchemaExtractor {
static extract(toolSet: CopilotToolSet): NativeLlmToolDefinition[] {
return Object.entries(toolSet).map(([name, tool]) => {
return {
name,
description: tool.description,
parameters: this.toJsonSchema(tool.inputSchema ?? z.object({})),
};
});
}
static toJsonSchema(schema: unknown): Record<string, unknown> {
if (!(schema instanceof z.ZodType)) {
if (schema && typeof schema === 'object' && !Array.isArray(schema)) {
return schema as Record<string, unknown>;
}
return { type: 'object', properties: {} };
}
if (schema instanceof z.ZodObject) {
const shape = schema.shape;
const properties: Record<string, unknown> = {};
const required: string[] = [];
for (const [key, child] of Object.entries(
shape as Record<string, z.ZodTypeAny>
)) {
properties[key] = this.toJsonSchema(child);
if (!this.isOptional(child)) {
required.push(key);
}
}
return {
type: 'object',
properties,
additionalProperties: false,
...(required.length ? { required } : {}),
};
}
if (schema instanceof z.ZodString) {
return { type: 'string' };
}
if (schema instanceof z.ZodNumber) {
return { type: 'number' };
}
if (schema instanceof z.ZodBoolean) {
return { type: 'boolean' };
}
if (schema instanceof z.ZodArray) {
return { type: 'array', items: this.toJsonSchema(schema.element) };
}
if (schema instanceof z.ZodEnum) {
return { type: 'string', enum: schema.options };
}
if (schema instanceof z.ZodLiteral) {
const literal = schema.value;
if (literal === null) {
return { const: null, type: 'null' };
}
if (typeof literal === 'string') {
return { const: literal, type: 'string' };
}
if (typeof literal === 'number') {
return { const: literal, type: 'number' };
}
if (typeof literal === 'boolean') {
return { const: literal, type: 'boolean' };
}
return { const: literal };
}
if (schema instanceof z.ZodUnion) {
return {
anyOf: schema.options.map((option: z.ZodTypeAny) =>
this.toJsonSchema(option)
),
};
}
if (schema instanceof z.ZodRecord) {
return {
type: 'object',
additionalProperties: this.toJsonSchema(schema.valueSchema),
};
}
if (schema instanceof z.ZodNullable) {
const inner = (schema._def as { innerType?: z.ZodTypeAny }).innerType;
return { anyOf: [this.toJsonSchema(inner), { type: 'null' }] };
}
if (schema instanceof z.ZodOptional || schema instanceof z.ZodDefault) {
return this.toJsonSchema(
(schema._def as { innerType?: z.ZodTypeAny }).innerType
);
}
if (schema instanceof z.ZodEffects) {
return this.toJsonSchema(
(schema._def as { schema?: z.ZodTypeAny }).schema
);
}
return { type: 'object', properties: {} };
}
private static isOptional(schema: z.ZodTypeAny): boolean {
if (schema instanceof z.ZodOptional || schema instanceof z.ZodDefault) {
return true;
}
if (schema instanceof z.ZodNullable) {
return this.isOptional(
(schema._def as { innerType: z.ZodTypeAny }).innerType
);
}
if (schema instanceof z.ZodEffects) {
return this.isOptional((schema._def as { schema: z.ZodTypeAny }).schema);
}
return false;
}
}
export class ToolCallLoop {
constructor(
private readonly dispatch: NativeDispatchFn,
private readonly tools: CopilotToolSet,
private readonly maxSteps = 20
) {}
private normalizeToolExecuteOptions(
signalOrOptions?: AbortSignal | CopilotToolExecuteOptions,
maybeMessages?: CopilotToolExecuteOptions['messages']
): CopilotToolExecuteOptions {
if (
signalOrOptions &&
typeof signalOrOptions === 'object' &&
'aborted' in signalOrOptions
) {
return {
signal: signalOrOptions,
messages: maybeMessages,
};
}
if (!signalOrOptions) {
return maybeMessages ? { messages: maybeMessages } : {};
}
return {
...signalOrOptions,
signal: signalOrOptions.signal,
messages: signalOrOptions.messages ?? maybeMessages,
};
}
async *run(
request: NativeLlmRequest,
signalOrOptions?: AbortSignal | CopilotToolExecuteOptions,
maybeMessages?: CopilotToolExecuteOptions['messages']
): AsyncIterableIterator<NativeLlmStreamEvent> {
const toolExecuteOptions = this.normalizeToolExecuteOptions(
signalOrOptions,
maybeMessages
);
const messages = request.messages.map(message => ({
...message,
content: [...message.content],
}));
for (let step = 0; step < this.maxSteps; step++) {
const toolCalls: NativeToolCall[] = [];
const accumulator = new ToolCallAccumulator();
let finalDone: Extract<NativeLlmStreamEvent, { type: 'done' }> | null =
null;
for await (const event of this.dispatch(
{
...request,
stream: true,
messages,
},
toolExecuteOptions.signal
)) {
switch (event.type) {
case 'tool_call_delta': {
accumulator.feedDelta(event);
break;
}
case 'tool_call': {
toolCalls.push(accumulator.complete(event));
yield event;
break;
}
case 'done': {
finalDone = event;
break;
}
case 'error': {
throw new Error(event.message);
}
default: {
yield event;
break;
}
}
}
toolCalls.push(...accumulator.drainPending());
if (toolCalls.length === 0) {
if (finalDone) {
yield finalDone;
}
break;
}
if (step === this.maxSteps - 1) {
throw new Error('ToolCallLoop max steps reached');
}
const toolResults = await this.executeTools(
toolCalls,
toolExecuteOptions
);
messages.push({
role: 'assistant',
content: toolCalls.map(call => ({
type: 'tool_call',
call_id: call.id,
name: call.name,
arguments: call.args,
arguments_text: call.rawArgumentsText,
arguments_error: call.argumentParseError,
thought: call.thought,
})),
});
for (const result of toolResults) {
messages.push({
role: 'tool',
content: [
{
type: 'tool_result',
call_id: result.callId,
name: result.name,
arguments: result.args,
arguments_text: result.rawArgumentsText,
arguments_error: result.argumentParseError,
output: result.output,
is_error: result.isError,
},
],
});
yield {
type: 'tool_result',
call_id: result.callId,
name: result.name,
arguments: result.args,
arguments_text: result.rawArgumentsText,
arguments_error: result.argumentParseError,
output: result.output,
is_error: result.isError,
};
}
}
}
private async executeTools(
calls: NativeToolCall[],
options: CopilotToolExecuteOptions
) {
return await Promise.all(
calls.map(call => this.executeTool(call, options))
);
}
private async executeTool(
call: NativeToolCall,
options: CopilotToolExecuteOptions
): Promise<ToolExecutionResult> {
const tool = this.tools[call.name] as CopilotTool | undefined;
if (!tool?.execute) {
return {
callId: call.id,
name: call.name,
args: call.args,
rawArgumentsText: call.rawArgumentsText,
argumentParseError: call.argumentParseError,
isError: true,
output: {
message: `Tool not found: ${call.name}`,
},
};
}
if (call.argumentParseError) {
return {
callId: call.id,
name: call.name,
args: call.args,
rawArgumentsText: call.rawArgumentsText,
argumentParseError: call.argumentParseError,
isError: true,
output: {
message: 'Invalid tool arguments JSON',
rawArguments: call.rawArgumentsText,
error: call.argumentParseError,
},
};
}
try {
const output = await tool.execute(call.args, options);
return {
callId: call.id,
name: call.name,
args: call.args,
rawArgumentsText: call.rawArgumentsText,
argumentParseError: call.argumentParseError,
output: output ?? null,
};
} catch (error) {
console.error('Tool execution failed', {
callId: call.id,
toolName: call.name,
error,
});
return {
callId: call.id,
name: call.name,
args: call.args,
rawArgumentsText: call.rawArgumentsText,
argumentParseError: call.argumentParseError,
isError: true,
output: {
message: 'Tool execution failed',
},
};
}
}
}
@@ -1,23 +1,11 @@
import {
CopilotProviderSideError,
metrics,
UserFriendlyError,
} from '../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
} from '../../../native';
import type { NodeTextMiddleware } from '../config';
import type { CopilotToolSet } from '../tools';
import { buildNativeRequest, NativeProviderAdapter } from './native';
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import { CopilotProvider } from './provider';
import type {
CopilotChatOptions,
ModelConditions,
PromptMessage,
} from './types';
import { CopilotProviderType, ModelInputType, ModelOutputType } from './types';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType, ModelOutputType } from './types';
export const DEFAULT_DIMENSIONS = 256;
@@ -28,42 +16,12 @@ export type MorphConfig = {
export class MorphProvider extends CopilotProvider<MorphConfig> {
readonly type = CopilotProviderType.Morph;
readonly models = [
{
id: 'morph-v2',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
},
],
},
{
id: 'morph-v3-fast',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
},
],
},
{
id: 'morph-v3-large',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
},
],
},
];
override configured(): boolean {
return !!this.config.apiKey;
protected resolveModelBackendKind() {
return 'morph' as const;
}
protected override setup() {
super.setup();
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
private handleError(e: any) {
@@ -77,106 +35,26 @@ export class MorphProvider extends CopilotProvider<MorphConfig> {
});
}
private createNativeConfig(): NativeLlmBackendConfig {
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
return {
base_url: 'https://api.morphllm.com',
auth_token: this.config.apiKey ?? '',
auth_token: this.getConfig(execution).apiKey ?? '',
};
}
private createNativeAdapter(
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream(
'openai_chat',
this.createNativeConfig(),
request,
signal
),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
);
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const model = this.selectModel(
await this.checkParams({
messages,
cond: fullCond,
options,
})
);
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
middleware,
});
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
return await adapter.text(request, options.signal, messages);
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const model = this.selectModel(
await this.checkParams({
messages,
cond: fullCond,
options,
})
);
try {
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
middleware,
});
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
for await (const chunk of adapter.streamText(
request,
options.signal,
messages
)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveOutputType: kind =>
kind === 'streamObject' ? null : ModelOutputType.Text,
},
structured: false,
embedding: false,
rerank: false,
};
}
}
@@ -1,692 +0,0 @@
import { ZodType } from 'zod';
import { CopilotPromptInvalid } from '../../../base';
import type {
NativeLlmCoreContent,
NativeLlmCoreMessage,
NativeLlmEmbeddingRequest,
NativeLlmRequest,
NativeLlmRerankRequest,
NativeLlmStreamEvent,
NativeLlmStructuredRequest,
NativeLlmStructuredResponse,
} from '../../../native';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import type { CopilotToolSet } from '../tools';
import {
canonicalizePromptAttachment,
type CanonicalPromptAttachment,
} from './attachments';
import { NativeDispatchFn, ToolCallLoop, ToolSchemaExtractor } from './loop';
import type {
CopilotChatOptions,
CopilotRerankRequest,
CopilotStructuredOptions,
ModelAttachmentCapability,
PromptMessage,
StreamObject,
} from './types';
import { CitationFootnoteFormatter, TextStreamParser } from './utils';
type BuildNativeRequestOptions = {
model: string;
messages: PromptMessage[];
options?: CopilotChatOptions | CopilotStructuredOptions;
tools?: CopilotToolSet;
withAttachment?: boolean;
attachmentCapability?: ModelAttachmentCapability;
include?: string[];
reasoning?: Record<string, unknown>;
responseSchema?: unknown;
middleware?: ProviderMiddlewareConfig;
};
type BuildNativeRequestResult = {
request: NativeLlmRequest;
schema?: ZodType;
};
type BuildNativeStructuredRequestResult = {
request: NativeLlmStructuredRequest;
schema: ZodType;
};
type ToolCallMeta = {
name: string;
args: Record<string, unknown>;
};
type NormalizedToolResultEvent = Extract<
NativeLlmStreamEvent,
{ type: 'tool_result' }
> & {
name: string;
arguments: Record<string, unknown>;
};
type AttachmentFootnote = {
blobId: string;
fileName: string;
fileType: string;
};
type NativeProviderAdapterOptions = {
nodeTextMiddleware?: NodeTextMiddleware[];
};
function roleToCore(role: PromptMessage['role']) {
switch (role) {
case 'assistant':
return 'assistant';
case 'system':
return 'system';
default:
return 'user';
}
}
function ensureAttachmentSupported(
attachment: CanonicalPromptAttachment,
attachmentCapability?: ModelAttachmentCapability
) {
if (!attachmentCapability) return;
if (!attachmentCapability.kinds.includes(attachment.kind)) {
throw new CopilotPromptInvalid(
`Native path does not support ${attachment.kind} attachments${
attachment.mediaType ? ` (${attachment.mediaType})` : ''
}`
);
}
if (
attachmentCapability.sourceKinds?.length &&
!attachmentCapability.sourceKinds.includes(attachment.sourceKind)
) {
throw new CopilotPromptInvalid(
`Native path does not support ${attachment.sourceKind} attachment sources`
);
}
if (attachment.isRemote && attachmentCapability.allowRemoteUrls === false) {
throw new CopilotPromptInvalid(
'Native path does not support remote attachment urls'
);
}
}
function resolveResponseSchema(
systemMessage: PromptMessage | undefined,
responseSchema?: unknown
): ZodType | undefined {
if (responseSchema instanceof ZodType) {
return responseSchema;
}
if (systemMessage?.responseFormat?.schema instanceof ZodType) {
return systemMessage.responseFormat.schema;
}
return systemMessage?.params?.schema instanceof ZodType
? systemMessage.params.schema
: undefined;
}
function resolveResponseStrict(
systemMessage: PromptMessage | undefined,
options?: CopilotStructuredOptions
) {
return options?.strict ?? systemMessage?.responseFormat?.strict ?? true;
}
export class StructuredResponseParseError extends Error {}
function normalizeStructuredText(text: string) {
const trimmed = text.replaceAll(/^ny\n/g, ' ').trim();
if (trimmed.startsWith('```') || trimmed.endsWith('```')) {
return trimmed
.replace(/```[\w\s-]*\n/g, '')
.replace(/\n```/g, '')
.trim();
}
return trimmed;
}
export function parseNativeStructuredOutput(
response: Pick<NativeLlmStructuredResponse, 'output_text'> & {
output_json?: unknown;
}
) {
if (response.output_json !== undefined) {
return response.output_json;
}
const normalized = normalizeStructuredText(response.output_text);
const candidates = [
() => normalized,
() => {
const objectStart = normalized.indexOf('{');
const objectEnd = normalized.lastIndexOf('}');
return objectStart !== -1 && objectEnd > objectStart
? normalized.slice(objectStart, objectEnd + 1)
: null;
},
() => {
const arrayStart = normalized.indexOf('[');
const arrayEnd = normalized.lastIndexOf(']');
return arrayStart !== -1 && arrayEnd > arrayStart
? normalized.slice(arrayStart, arrayEnd + 1)
: null;
},
];
for (const candidate of candidates) {
try {
const candidateText = candidate();
if (typeof candidateText === 'string') {
return JSON.parse(candidateText);
}
} catch {
continue;
}
}
throw new StructuredResponseParseError(
`Unexpected structured response: ${normalized.slice(0, 200)}`
);
}
export function buildNativeRerankRequest(
model: string,
request: CopilotRerankRequest
): NativeLlmRerankRequest {
return {
model,
query: request.query,
candidates: request.candidates.map(candidate => ({
...(candidate.id ? { id: candidate.id } : {}),
text: candidate.text,
})),
...(request.topK ? { top_n: request.topK } : {}),
};
}
async function toCoreContents(
message: PromptMessage,
withAttachment: boolean,
attachmentCapability?: ModelAttachmentCapability
): Promise<NativeLlmCoreContent[]> {
const contents: NativeLlmCoreContent[] = [];
if (typeof message.content === 'string' && message.content.length) {
contents.push({ type: 'text', text: message.content });
}
if (!withAttachment || !Array.isArray(message.attachments)) return contents;
for (const entry of message.attachments) {
const normalized = await canonicalizePromptAttachment(entry, message);
ensureAttachmentSupported(normalized, attachmentCapability);
contents.push({
type: normalized.kind,
source: normalized.source,
});
}
return contents;
}
export async function buildNativeRequest({
model,
messages,
options = {},
tools = {},
withAttachment = true,
attachmentCapability,
include,
reasoning,
responseSchema,
middleware,
}: BuildNativeRequestOptions): Promise<BuildNativeRequestResult> {
const copiedMessages = messages.map(message => ({
...message,
attachments: message.attachments
? [...message.attachments]
: message.attachments,
}));
const systemMessage =
copiedMessages[0]?.role === 'system' ? copiedMessages.shift() : undefined;
const schema = resolveResponseSchema(systemMessage, responseSchema);
const coreMessages: NativeLlmCoreMessage[] = [];
if (systemMessage?.content?.length) {
coreMessages.push({
role: 'system',
content: [{ type: 'text', text: systemMessage.content }],
});
}
for (const message of copiedMessages) {
if (message.role === 'system') continue;
const content = await toCoreContents(
message,
withAttachment,
attachmentCapability
);
coreMessages.push({ role: roleToCore(message.role), content });
}
return {
request: {
model,
stream: true,
messages: coreMessages,
max_tokens: options.maxTokens ?? undefined,
temperature: options.temperature ?? undefined,
tools: ToolSchemaExtractor.extract(tools),
tool_choice: Object.keys(tools).length ? 'auto' : undefined,
include,
reasoning,
response_schema: schema
? ToolSchemaExtractor.toJsonSchema(schema)
: undefined,
middleware: middleware?.rust
? { request: middleware.rust.request, stream: middleware.rust.stream }
: undefined,
},
schema,
};
}
export async function buildNativeStructuredRequest({
model,
messages,
options = {},
withAttachment = true,
attachmentCapability,
reasoning,
responseSchema,
middleware,
}: Omit<
BuildNativeRequestOptions,
'tools' | 'include'
>): Promise<BuildNativeStructuredRequestResult> {
const copiedMessages = messages.map(message => ({
...message,
attachments: message.attachments
? [...message.attachments]
: message.attachments,
}));
const systemMessage =
copiedMessages[0]?.role === 'system' ? copiedMessages.shift() : undefined;
const schema = resolveResponseSchema(systemMessage, responseSchema);
const strict = resolveResponseStrict(systemMessage, options);
if (!schema) {
throw new CopilotPromptInvalid('Schema is required');
}
const coreMessages: NativeLlmCoreMessage[] = [];
if (systemMessage?.content?.length) {
coreMessages.push({
role: 'system',
content: [{ type: 'text', text: systemMessage.content }],
});
}
for (const message of copiedMessages) {
if (message.role === 'system') continue;
const content = await toCoreContents(
message,
withAttachment,
attachmentCapability
);
coreMessages.push({ role: roleToCore(message.role), content });
}
return {
request: {
model,
messages: coreMessages,
schema: ToolSchemaExtractor.toJsonSchema(schema),
max_tokens: options.maxTokens ?? undefined,
temperature: options.temperature ?? undefined,
reasoning,
strict,
response_mime_type: 'application/json',
middleware: middleware?.rust
? { request: middleware.rust.request }
: undefined,
},
schema,
};
}
export function buildNativeEmbeddingRequest({
model,
inputs,
dimensions,
taskType = 'RETRIEVAL_DOCUMENT',
}: {
model: string;
inputs: string[];
dimensions?: number;
taskType?: string;
}): NativeLlmEmbeddingRequest {
return {
model,
inputs,
dimensions,
task_type: taskType,
};
}
function ensureToolResultMeta(
event: Extract<NativeLlmStreamEvent, { type: 'tool_result' }>,
toolCalls: Map<string, ToolCallMeta>
): NormalizedToolResultEvent | null {
const name = event.name ?? toolCalls.get(event.call_id)?.name;
const args = event.arguments ?? toolCalls.get(event.call_id)?.args;
if (!name || !args) return null;
return { ...event, name, arguments: args };
}
function pickAttachmentFootnote(value: unknown): AttachmentFootnote | null {
if (!value || typeof value !== 'object') {
return null;
}
const record = value as Record<string, unknown>;
const blobId =
typeof record.blobId === 'string'
? record.blobId
: typeof record.blob_id === 'string'
? record.blob_id
: undefined;
const fileName =
typeof record.fileName === 'string'
? record.fileName
: typeof record.name === 'string'
? record.name
: undefined;
const fileType =
typeof record.fileType === 'string'
? record.fileType
: typeof record.mimeType === 'string'
? record.mimeType
: 'application/octet-stream';
if (!blobId || !fileName) {
return null;
}
return { blobId, fileName, fileType };
}
function collectAttachmentFootnotes(
event: NormalizedToolResultEvent
): AttachmentFootnote[] {
if (event.name === 'blob_read') {
const item = pickAttachmentFootnote(event.output);
return item ? [item] : [];
}
if (event.name === 'doc_semantic_search' && Array.isArray(event.output)) {
return event.output
.map(item => pickAttachmentFootnote(item))
.filter((item): item is AttachmentFootnote => item !== null);
}
return [];
}
function formatAttachmentFootnotes(attachments: AttachmentFootnote[]) {
const references = attachments.map((_, index) => `[^${index + 1}]`).join('');
const definitions = attachments
.map((attachment, index) => {
return `[^${index + 1}]: ${JSON.stringify({
type: 'attachment',
blobId: attachment.blobId,
fileName: attachment.fileName,
fileType: attachment.fileType,
})}`;
})
.join('\n');
return `\n\n${references}\n\n${definitions}`;
}
export class NativeProviderAdapter {
readonly #loop: ToolCallLoop;
readonly #enableCallout: boolean;
readonly #enableCitationFootnote: boolean;
constructor(
dispatch: NativeDispatchFn,
tools: CopilotToolSet,
maxSteps = 20,
options: NativeProviderAdapterOptions = {}
) {
this.#loop = new ToolCallLoop(dispatch, tools, maxSteps);
const enabledNodeTextMiddlewares = new Set(
options.nodeTextMiddleware ?? ['citation_footnote', 'callout']
);
this.#enableCallout =
enabledNodeTextMiddlewares.has('callout') ||
enabledNodeTextMiddlewares.has('thinking_format');
this.#enableCitationFootnote =
enabledNodeTextMiddlewares.has('citation_footnote');
}
async text(
request: NativeLlmRequest,
signal?: AbortSignal,
messages?: PromptMessage[]
) {
let output = '';
for await (const chunk of this.streamText(request, signal, messages)) {
output += chunk;
}
return output.trim();
}
async *streamText(
request: NativeLlmRequest,
signal?: AbortSignal,
messages?: PromptMessage[]
): AsyncIterableIterator<string> {
const textParser = this.#enableCallout ? new TextStreamParser() : null;
const citationFormatter = this.#enableCitationFootnote
? new CitationFootnoteFormatter()
: null;
const toolCalls = new Map<string, ToolCallMeta>();
let streamPartId = 0;
for await (const event of this.#loop.run(request, signal, messages)) {
switch (event.type) {
case 'text_delta': {
if (textParser) {
yield textParser.parse({
type: 'text-delta',
id: String(streamPartId++),
text: event.text,
});
} else {
yield event.text;
}
break;
}
case 'reasoning_delta': {
if (textParser) {
yield textParser.parse({
type: 'reasoning-delta',
id: String(streamPartId++),
text: event.text,
});
} else {
yield event.text;
}
break;
}
case 'tool_call': {
const toolCall = {
name: event.name,
args: event.arguments,
};
toolCalls.set(event.call_id, toolCall);
if (textParser) {
yield textParser.parse({
type: 'tool-call',
toolCallId: event.call_id,
toolName: event.name as never,
input: event.arguments,
});
}
break;
}
case 'tool_result': {
const normalized = ensureToolResultMeta(event, toolCalls);
if (!normalized || !textParser) {
break;
}
yield textParser.parse({
type: 'tool-result',
toolCallId: normalized.call_id,
toolName: normalized.name as never,
input: normalized.arguments,
output: normalized.output,
});
break;
}
case 'citation': {
if (citationFormatter) {
citationFormatter.consume({
type: 'citation',
index: event.index,
url: event.url,
});
}
break;
}
case 'done': {
const footnotes = textParser?.end() ?? '';
const citations = citationFormatter?.end() ?? '';
const tails = [citations, footnotes].filter(Boolean).join('\n');
if (tails) {
yield `\n${tails}`;
}
break;
}
case 'error': {
throw new Error(event.message);
}
default:
break;
}
}
}
async *streamObject(
request: NativeLlmRequest,
signal?: AbortSignal,
messages?: PromptMessage[]
): AsyncIterableIterator<StreamObject> {
const toolCalls = new Map<string, ToolCallMeta>();
const citationFormatter = this.#enableCitationFootnote
? new CitationFootnoteFormatter()
: null;
const fallbackAttachmentFootnotes = new Map<string, AttachmentFootnote>();
let hasFootnoteReference = false;
for await (const event of this.#loop.run(request, signal, messages)) {
switch (event.type) {
case 'text_delta': {
if (event.text.includes('[^')) {
hasFootnoteReference = true;
}
yield {
type: 'text-delta',
textDelta: event.text,
};
break;
}
case 'reasoning_delta': {
yield {
type: 'reasoning',
textDelta: event.text,
};
break;
}
case 'tool_call': {
const toolCall = {
name: event.name,
args: event.arguments,
};
toolCalls.set(event.call_id, toolCall);
yield {
type: 'tool-call',
toolCallId: event.call_id,
toolName: event.name,
args: event.arguments,
};
break;
}
case 'tool_result': {
const normalized = ensureToolResultMeta(event, toolCalls);
if (!normalized) {
break;
}
const attachments = collectAttachmentFootnotes(normalized);
attachments.forEach(attachment => {
fallbackAttachmentFootnotes.set(attachment.blobId, attachment);
});
yield {
type: 'tool-result',
toolCallId: normalized.call_id,
toolName: normalized.name,
args: normalized.arguments,
result: normalized.output,
};
break;
}
case 'citation': {
if (citationFormatter) {
citationFormatter.consume({
type: 'citation',
index: event.index,
url: event.url,
});
}
break;
}
case 'done': {
const citations = citationFormatter?.end() ?? '';
if (citations) {
hasFootnoteReference = true;
yield {
type: 'text-delta',
textDelta: `\n${citations}`,
};
}
if (!hasFootnoteReference && fallbackAttachmentFootnotes.size > 0) {
yield {
type: 'text-delta',
textDelta: formatAttachmentFootnotes(
Array.from(fallbackAttachmentFootnotes.values())
),
};
}
break;
}
case 'error': {
throw new Error(event.message);
}
default:
break;
}
}
}
}
File diff suppressed because it is too large Load Diff
@@ -1,21 +1,12 @@
import { CopilotProviderSideError, metrics } from '../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
} from '../../../native';
import type { NodeTextMiddleware } from '../config';
import type { CopilotToolSet } from '../tools';
import { buildNativeRequest, NativeProviderAdapter } from './native';
import { CopilotProviderSideError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import { CopilotProvider } from './provider';
import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
import {
CopilotChatOptions,
CopilotProviderType,
ModelConditions,
ModelInputType,
ModelOutputType,
PromptMessage,
} from './types';
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType, ModelOutputType } from './types';
export type PerplexityConfig = {
apiKey: string;
@@ -25,166 +16,50 @@ export type PerplexityConfig = {
export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
readonly type = CopilotProviderType.Perplexity;
readonly models = [
{
name: 'Sonar',
id: 'sonar',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
defaultForOutputType: true,
},
],
},
{
name: 'Sonar Pro',
id: 'sonar-pro',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
},
],
},
{
name: 'Sonar Reasoning',
id: 'sonar-reasoning',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
},
],
},
{
name: 'Sonar Reasoning Pro',
id: 'sonar-reasoning-pro',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
},
],
},
];
override configured(): boolean {
return !!this.config.apiKey;
protected resolveModelBackendKind() {
return 'perplexity' as const;
}
protected override setup() {
super.setup();
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
private createNativeConfig(): NativeLlmBackendConfig {
const baseUrl = this.config.endpoint || 'https://api.perplexity.ai';
override getDriverSpec(): ProviderDriverSpec {
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: this.config.apiKey,
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveOutputType: kind =>
kind === 'streamObject' ? null : ModelOutputType.Text,
withAttachment: false,
resolveRequestOptions: async context => ({
withAttachment: !hasProviderModelBehaviorFlag(
context.model,
'no_attachments'
),
include: hasProviderModelBehaviorFlag(
context.model,
'citations_include'
)
? ['citations']
: undefined,
}),
},
structured: false,
embedding: false,
rerank: false,
};
}
private createNativeAdapter(
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream(
'openai_chat',
this.createNativeConfig(),
request,
signal
),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
);
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
withAttachment: false,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
withAttachment: false,
include: ['citations'],
middleware,
});
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
return await adapter.text(request, options.signal, messages);
} catch (e: any) {
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
const normalizedCond = await this.checkParams({
cond: fullCond,
messages,
options,
withAttachment: false,
});
const model = this.selectModel(normalizedCond);
try {
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
withAttachment: false,
include: ['citations'],
middleware,
});
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
for await (const chunk of adapter.streamText(
request,
options.signal,
messages
)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
const config = this.getConfig(execution);
const baseUrl = config.endpoint || 'https://api.perplexity.ai';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: config.apiKey,
};
}
private handleError(e: any) {
@@ -1,81 +1,43 @@
import type { ProviderMiddlewareConfig } from '../config';
import { CopilotProviderType } from './types';
const DEFAULT_NODE_TEXT_MIDDLEWARE: NonNullable<
NonNullable<ProviderMiddlewareConfig['node']>['text']
> = ['citation_footnote', 'callout'];
const DEFAULT_MIDDLEWARE_BY_TYPE: Record<
CopilotProviderType,
ProviderMiddlewareConfig
> = {
[CopilotProviderType.OpenAI]: {
rust: {
request: ['normalize_messages'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.CloudflareWorkersAi]: {
rust: {
request: ['normalize_messages'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Anthropic]: {
rust: {
request: ['normalize_messages', 'tool_schema_rewrite'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.AnthropicVertex]: {
rust: {
request: ['normalize_messages', 'tool_schema_rewrite'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Morph]: {
rust: {
request: ['clamp_max_tokens'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Perplexity]: {
rust: {
request: ['clamp_max_tokens'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Gemini]: {
rust: {
request: ['normalize_messages', 'tool_schema_rewrite'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.GeminiVertex]: {
rust: {
request: ['normalize_messages', 'tool_schema_rewrite'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.FAL]: {},
};
@@ -91,18 +53,26 @@ function mergeArray<T>(base: T[] | undefined, override: T[] | undefined) {
return unique([...(base ?? []), ...(override ?? [])]);
}
function compactMiddlewareSection<T extends Record<string, unknown>>(
section: T
): T | undefined {
return Object.values(section).some(value => value !== undefined)
? section
: undefined;
}
export function mergeProviderMiddleware(
defaults: ProviderMiddlewareConfig,
override?: ProviderMiddlewareConfig
): ProviderMiddlewareConfig {
return {
rust: {
rust: compactMiddlewareSection({
request: mergeArray(defaults.rust?.request, override?.rust?.request),
stream: mergeArray(defaults.rust?.stream, override?.rust?.stream),
},
node: {
}),
node: compactMiddlewareSection({
text: mergeArray(defaults.node?.text, override?.node?.text),
},
}),
};
}
@@ -0,0 +1,385 @@
import { z } from 'zod';
import { CopilotPromptInvalid } from '../../../base';
import {
type LlmBackendConfig,
llmInferPromptModelConditions,
llmMatchModelCapabilities,
llmMatchModelRegistry,
type LlmProtocol,
llmResolveModelRegistryVariant,
} from '../../../native';
import { applyPromptAttachmentMimeTypeHintForNative } from './attachments';
import {
type CopilotChatOptions,
type CopilotImageOptions,
type CopilotModelBackendKind,
type CopilotProviderModel,
type CopilotProviderType,
type CopilotStructuredOptions,
EmbeddingMessage,
type ModelAttachmentCapability,
type ModelCapability,
type ModelFullConditions,
ModelInputType,
ModelOutputType,
type PromptAttachmentKind,
type PromptAttachmentSourceKind,
type PromptMessage,
PromptMessageSchema,
} from './types';
// Owner: backend host model-selection glue.
// Capability matching and catalog lookup are delegated to native/adapter; this
// file keeps provider prefix/default/prefer behavior and Node prompt checks.
export type ProviderModelRuntimeContext = {
type: CopilotProviderType;
backendKind: CopilotModelBackendKind;
};
export type ResolvedProviderModel = CopilotProviderModel & {
backendKind: CopilotModelBackendKind;
canonicalKey: string;
protocol?: LlmProtocol;
requestLayer?: LlmBackendConfig['request_layer'];
routeOverrides?: Partial<
Record<
ModelOutputType,
{
protocol?: LlmProtocol;
requestLayer?: LlmBackendConfig['request_layer'];
}
>
>;
behaviorFlags?: string[];
};
function unique<T>(values: Iterable<T>) {
return Array.from(new Set(values));
}
function resolveAttachmentCapability(
cap: ModelCapability,
outputType?: ModelOutputType
): ModelAttachmentCapability | undefined {
if (outputType === ModelOutputType.Structured) {
return cap.structuredAttachments ?? cap.attachments;
}
return cap.attachments;
}
function toProviderModel(
variant: NonNullable<
ReturnType<typeof llmResolveModelRegistryVariant>['variant']
>
): ResolvedProviderModel {
return {
id: variant.rawModelId,
name: variant.displayName,
backendKind: variant.backendKind,
canonicalKey: variant.canonicalKey,
protocol: variant.protocol,
requestLayer: variant.requestLayer,
routeOverrides: variant.routeOverrides,
behaviorFlags: variant.behaviorFlags,
capabilities: variant.capabilities.map(capability => ({
input: capability.input as ModelInputType[],
output: capability.output as ModelOutputType[],
attachments: capability.attachments
? {
kinds: capability.attachments.kinds as PromptAttachmentKind[],
sourceKinds: capability.attachments.sourceKinds as
| ModelAttachmentCapability['sourceKinds']
| undefined,
allowRemoteUrls: capability.attachments.allowRemoteUrls,
}
: undefined,
structuredAttachments: capability.structuredAttachments
? {
kinds: capability.structuredAttachments
.kinds as PromptAttachmentKind[],
sourceKinds: capability.structuredAttachments.sourceKinds as
| ModelAttachmentCapability['sourceKinds']
| undefined,
allowRemoteUrls: capability.structuredAttachments.allowRemoteUrls,
}
: undefined,
defaultForOutputType: capability.defaultForOutputType,
})),
};
}
export type ProviderModelSelection = {
kind: 'configured';
model: ResolvedProviderModel;
};
export function resolveProviderModelSelection(
context: ProviderModelRuntimeContext,
cond: ModelFullConditions
): ProviderModelSelection | undefined {
if (cond.modelId) {
const resolved = llmResolveModelRegistryVariant({
backendKind: context.backendKind,
modelId: cond.modelId,
}).variant;
if (!resolved) {
return;
}
const model = toProviderModel(resolved);
const matchedModelId = llmMatchModelCapabilities([model], {
...cond,
modelId: model.id,
});
if (!matchedModelId) {
return;
}
return {
kind: 'configured',
model,
};
}
const resolved = llmMatchModelRegistry({
backendKind: context.backendKind,
cond,
}).variant;
if (!resolved) {
return;
}
return {
kind: 'configured',
model: toProviderModel(resolved),
};
}
function isMultimodal(model: CopilotProviderModel) {
return model.capabilities.some(c =>
[ModelInputType.Image, ModelInputType.Audio, ModelInputType.File].some(t =>
c.input.includes(t)
)
);
}
function handleZodError(ret: z.SafeParseReturnType<any, any>) {
if (ret.success) return;
const issues = ret.error.issues.map(i => {
const path =
'root' +
(i.path.length
? `.${i.path.map(seg => (typeof seg === 'number' ? `[${seg}]` : `.${seg}`)).join('')}`
: '');
return `${i.message}${path}`;
});
throw new CopilotPromptInvalid(issues.join('; '));
}
export async function inferModelConditionsFromMessages(
messages?: PromptMessage[],
withAttachment = true
): Promise<Partial<ModelFullConditions>> {
if (!messages?.length || !withAttachment) return {};
const projectedMessages = messages.map(message => ({
role: message.role,
content: message.content,
...(Array.isArray(message.attachments) && message.attachments.length
? {
attachments: message.attachments.map(attachment =>
applyPromptAttachmentMimeTypeHintForNative(attachment, message)
),
}
: {}),
}));
const inferredCond = llmInferPromptModelConditions(projectedMessages);
return {
...(inferredCond.attachmentKinds?.length
? { attachmentKinds: unique(inferredCond.attachmentKinds) }
: {}),
...(inferredCond.attachmentSourceKinds?.length
? {
attachmentSourceKinds: unique(
inferredCond.attachmentSourceKinds
) as PromptAttachmentSourceKind[],
}
: {}),
...(inferredCond.inputTypes?.length
? { inputTypes: unique(inferredCond.inputTypes) as ModelInputType[] }
: {}),
...(inferredCond.hasRemoteAttachments
? { hasRemoteAttachments: true }
: {}),
};
}
export function mergeModelConditions(
cond: ModelFullConditions,
inferredCond: Partial<ModelFullConditions>
): ModelFullConditions {
return {
...inferredCond,
...cond,
inputTypes: unique([
...(inferredCond.inputTypes ?? []),
...(cond.inputTypes ?? []),
]),
attachmentKinds: unique([
...(inferredCond.attachmentKinds ?? []),
...(cond.attachmentKinds ?? []),
]),
attachmentSourceKinds: unique([
...(inferredCond.attachmentSourceKinds ?? []),
...(cond.attachmentSourceKinds ?? []),
]),
hasRemoteAttachments:
cond.hasRemoteAttachments ?? inferredCond.hasRemoteAttachments,
};
}
export function getAttachCapability(
model: CopilotProviderModel,
outputType: ModelOutputType
): ModelAttachmentCapability | undefined {
const capability =
model.capabilities.find(cap => cap.output.includes(outputType)) ??
model.capabilities[0];
if (!capability) {
return;
}
return resolveAttachmentCapability(capability, outputType);
}
export function matchProviderModel(
context: ProviderModelRuntimeContext,
cond: ModelFullConditions
): boolean {
return !!resolveProviderModelSelection(context, cond);
}
export function resolveProviderModel(
context: ProviderModelRuntimeContext,
modelId: string
): ResolvedProviderModel | undefined {
return resolveProviderModelSelection(context, {
modelId,
})?.model;
}
export function hasProviderModelBehaviorFlag(
model: CopilotProviderModel,
flag: string
) {
const behaviorFlags = (model as ResolvedProviderModel).behaviorFlags;
return Array.isArray(behaviorFlags) && behaviorFlags.includes(flag);
}
export function resolveProviderModelRoute(
model: CopilotProviderModel,
outputType: ModelOutputType
) {
const resolved = model as ResolvedProviderModel;
const override = resolved.routeOverrides?.[outputType];
return {
protocol: override?.protocol ?? resolved.protocol,
requestLayer: override?.requestLayer ?? resolved.requestLayer,
};
}
export function requireProviderModelSelection(
context: ProviderModelRuntimeContext,
cond: ModelFullConditions
): ResolvedProviderModel {
const selection = resolveProviderModelSelection(context, cond);
if (selection) return selection.model;
const { modelId, outputType, inputTypes } = cond;
throw new CopilotPromptInvalid(
modelId
? `Model ${modelId} does not support ${outputType ?? '<any>'} output with ${inputTypes ?? '<any>'} input`
: outputType
? `No model supports ${outputType} output with ${inputTypes ?? '<any>'} input for provider ${context.type}`
: 'Output type is required when modelId is not provided'
);
}
export async function checkProviderParams(
context: ProviderModelRuntimeContext,
{
cond,
messages,
embeddings,
options = {},
withAttachment = true,
}: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions;
withAttachment?: boolean;
execution?: unknown;
}
): Promise<ModelFullConditions> {
if (messages) {
const { requireContent = true, requireAttachment = false } = options;
const MessageSchema = z
.array(
PromptMessageSchema.extend({
content: requireContent
? z.string().trim().min(1)
: z.string().optional().nullable(),
})
.passthrough()
.catchall(z.union([z.string(), z.number(), z.date(), z.null()]))
)
.optional();
handleZodError(MessageSchema.safeParse(messages));
const inferredCond = await inferModelConditionsFromMessages(
messages,
withAttachment
);
const mergedCond = mergeModelConditions(cond, inferredCond);
const model = requireProviderModelSelection(context, mergedCond);
const multimodal = isMultimodal(model);
if (
multimodal &&
requireAttachment &&
!messages.some(
message =>
message.role === 'user' &&
Array.isArray(message.attachments) &&
message.attachments.length > 0
)
) {
throw new CopilotPromptInvalid('attachments required in multimodal mode');
}
if (embeddings) {
handleZodError(EmbeddingMessage.safeParse(embeddings));
}
return mergedCond;
}
const inferredCond = await inferModelConditionsFromMessages(
messages,
withAttachment
);
const mergedCond = mergeModelConditions(cond, inferredCond);
if (embeddings) {
handleZodError(EmbeddingMessage.safeParse(embeddings));
}
return mergedCond;
}
@@ -0,0 +1,337 @@
import type {
LlmBackendConfig,
LlmEmbeddingRequest,
LlmProtocol,
LlmRerankRequest,
LlmStructuredRequest,
} from '../../../native';
import {
buildLlmImageRequestFromMessages,
llmEmbeddingDispatch,
llmRerankDispatch,
llmStructuredDispatch,
} from '../../../native';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import {
buildToolContracts,
projectPromptMessageForNative,
} from '../runtime/contracts';
import { buildNativeRequest } from '../runtime/native-request-runtime';
import type { ToolLoopBackend } from '../runtime/tool/bridge';
import type { NativeProviderAdapter } from '../runtime/tool/native-adapter';
import type { CopilotToolSet } from '../tools';
import type {
CopilotProviderExecution,
PreparedNativeEmbeddingExecution,
PreparedNativeExecution,
PreparedNativeImageExecution,
PreparedNativeRequestOptions,
PreparedNativeRerankExecution,
PreparedNativeStructuredExecution,
} from './provider-runtime-contract';
import type {
CopilotChatOptions,
CopilotImageOptions,
PromptMessage,
} from './types';
export type CreateToolAdapterOptions = {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
};
export type CreateNativeAdapter = (
backend: ToolLoopBackend,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[],
options?: CreateToolAdapterOptions
) => NativeProviderAdapter;
export type CreatePreparedExecutionRuntimeInput = {
resolveProviderId: (execution?: CopilotProviderExecution) => string;
getTools: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>;
getActiveProviderMiddleware: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig;
createNativeAdapter: CreateNativeAdapter;
maxSteps: number;
};
export type PreparedExecutionRuntime = ReturnType<
typeof createPreparedExecutionRuntime
>;
export function createPreparedExecutionRuntime(
input: CreatePreparedExecutionRuntimeInput
) {
return {
buildPreparedNativeExecution: async (
prepared: PreparedNativeRequestOptions
) =>
await buildPreparedNativeExecution(
input.resolveProviderId(prepared.execution),
input.getTools,
input.getActiveProviderMiddleware,
input.maxSteps,
prepared
),
createPreparedExecutionAdapter: (prepared: PreparedNativeExecution) =>
createPreparedExecutionAdapter(
input.createNativeAdapter,
input.maxSteps,
prepared
),
buildPreparedNativeStructuredExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmStructuredRequest,
execution?: CopilotProviderExecution
) =>
buildPreparedNativeStructuredExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
request
),
buildPreparedNativeEmbeddingExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmEmbeddingRequest,
execution?: CopilotProviderExecution
) =>
buildPreparedNativeEmbeddingExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
request
),
buildPreparedNativeRerankExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmRerankRequest,
execution?: CopilotProviderExecution
) =>
buildPreparedNativeRerankExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
request
),
buildPreparedNativeImageExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
messages: PromptMessage[],
options: CopilotImageOptions = {},
execution?: CopilotProviderExecution
) =>
buildPreparedNativeImageExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
messages,
options
),
};
}
export function createPreparedExecutionAdapter(
createNativeAdapter: CreateNativeAdapter,
maxSteps: number,
prepared: PreparedNativeExecution
) {
return createNativeAdapter(
{
protocol: prepared.route.protocol,
backendConfig: prepared.route.backendConfig,
},
prepared.tools,
prepared.postprocess?.nodeTextMiddleware,
{
maxSteps,
nodeTextMiddleware: prepared.postprocess?.nodeTextMiddleware,
}
);
}
export function createNativeStructuredDispatch(
backendConfig: LlmBackendConfig,
protocol: LlmProtocol
) {
return (request: LlmStructuredRequest) =>
llmStructuredDispatch(protocol, backendConfig, request);
}
export function createNativeEmbeddingDispatch(
backendConfig: LlmBackendConfig,
protocol: LlmProtocol
) {
return (request: LlmEmbeddingRequest) =>
llmEmbeddingDispatch(protocol, backendConfig, request);
}
export function createNativeRerankDispatch(
backendConfig: LlmBackendConfig,
protocol: LlmProtocol
) {
return (request: LlmRerankRequest) =>
llmRerankDispatch(protocol, backendConfig, request);
}
function buildPreparedRoute(
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string
): PreparedNativeExecution['route'] {
return {
providerId,
protocol,
requestLayer: backendConfig.request_layer,
model,
backendConfig,
};
}
export async function buildPreparedNativeExecution(
providerId: string,
getTools: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>,
getActiveProviderMiddleware: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig,
maxSteps: number,
{
protocol,
backendConfig,
model,
messages,
options = {},
execution,
withAttachment = true,
attachmentCapability,
include,
reasoning,
tools,
middleware,
}: PreparedNativeRequestOptions
): Promise<PreparedNativeExecution> {
const resolvedTools = tools ?? (await getTools(options, model));
const resolvedMiddleware =
middleware ?? getActiveProviderMiddleware(execution);
const { request } = await buildNativeRequest({
model,
messages,
options,
toolContracts: buildToolContracts(resolvedTools),
withAttachment,
attachmentCapability,
include,
reasoning,
middleware: resolvedMiddleware,
});
return {
route: buildPreparedRoute(providerId, protocol, backendConfig, model),
request,
tools: resolvedTools,
maxSteps,
postprocess: {
nodeTextMiddleware: resolvedMiddleware.node?.text,
},
};
}
type BuildPreparedNativeDispatchExecution = <
TRequest extends
| LlmStructuredRequest
| LlmEmbeddingRequest
| LlmRerankRequest,
>(
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: TRequest
) => {
route: PreparedNativeExecution['route'];
request: TRequest;
};
const buildPreparedNativeDispatchExecution: BuildPreparedNativeDispatchExecution =
(providerId, protocol, backendConfig, model, request) => {
return {
route: buildPreparedRoute(providerId, protocol, backendConfig, model),
request,
};
};
export const buildPreparedNativeStructuredExecution =
buildPreparedNativeDispatchExecution as (
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmStructuredRequest
) => PreparedNativeStructuredExecution;
export const buildPreparedNativeEmbeddingExecution =
buildPreparedNativeDispatchExecution as (
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmEmbeddingRequest
) => PreparedNativeEmbeddingExecution;
export const buildPreparedNativeRerankExecution =
buildPreparedNativeDispatchExecution as (
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmRerankRequest
) => PreparedNativeRerankExecution;
export function buildPreparedNativeImageExecution(
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
messages: PromptMessage[],
options: CopilotImageOptions = {}
): PreparedNativeImageExecution {
const nativeMessages = messages.map(
message => projectPromptMessageForNative(message).message
);
return {
route: buildPreparedRoute(providerId, protocol, backendConfig, model),
request: buildLlmImageRequestFromMessages({
model,
protocol,
messages: nativeMessages,
options: projectImageRequestOptions(options),
}),
};
}
function projectImageRequestOptions(options: CopilotImageOptions = {}) {
return {
quality: options.quality,
seed: options.seed,
modelName: options.modelName,
loras: options.loras,
};
}
@@ -240,6 +240,16 @@ export function resolveModel({
};
}
if (modelId) {
return {
rawModelId: modelId,
modelId,
candidateProviderIds: registry.order.filter(providerId =>
isAllowed(providerId)
),
};
}
const defaultProviderId =
outputType && outputType !== ModelOutputType.Rerank
? registry.defaults[outputType]
@@ -0,0 +1,456 @@
import type {
LlmBackendConfig,
LlmEmbeddingRequest,
LlmImageRequest,
LlmProtocol,
LlmRequest,
LlmRerankRequest,
LlmStructuredRequest,
} from '../../../native';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import type { CopilotToolSet } from '../tools';
import {
type ProviderModelRuntimeContext,
resolveProviderModelRoute,
} from './provider-model-runtime';
import type { NormalizedCopilotProviderProfile } from './provider-registry';
import {
CopilotChatOptions,
CopilotImageOptions,
CopilotProviderModel,
CopilotStructuredOptions,
ModelAttachmentCapability,
ModelConditions,
ModelFullConditions,
ModelOutputType,
PromptMessage,
} from './types';
export type NativeExecutionRoute = {
protocol: LlmProtocol;
requestLayer?: LlmBackendConfig['request_layer'];
model: string;
backendConfig: LlmBackendConfig;
};
export type CopilotProviderExecution = {
providerId: string;
profile: NormalizedCopilotProviderProfile;
};
export type PreparedNativeExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmRequest;
tools: CopilotToolSet;
maxSteps?: number;
postprocess?: {
nodeTextMiddleware?: NodeTextMiddleware[];
};
};
export type PreparedNativeStructuredExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmStructuredRequest;
};
export type PreparedNativeEmbeddingExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmEmbeddingRequest;
};
export type PreparedNativeRerankExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmRerankRequest;
};
export type PreparedNativeImageExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmImageRequest;
};
export type PreparedNativeRequestOptions = {
protocol: LlmProtocol;
backendConfig: LlmBackendConfig;
model: string;
messages: PromptMessage[];
options?: CopilotChatOptions;
execution?: CopilotProviderExecution;
withAttachment?: boolean;
attachmentCapability?: ModelAttachmentCapability;
include?: string[];
reasoning?: Record<string, unknown>;
tools?: CopilotToolSet;
middleware?: ProviderMiddlewareConfig;
};
type ProviderChatDriverPrepareResult = Omit<
PreparedNativeRequestOptions,
'execution' | 'options'
>;
type Awaitable<T> = T | Promise<T>;
type NativeBackendConfigResolver = (
execution?: CopilotProviderExecution
) => Awaitable<LlmBackendConfig>;
export type StructuredProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
prepareMessages?: (
messages: PromptMessage[],
backendConfig: LlmBackendConfig,
options: NonNullable<CopilotStructuredOptions>
) => Promise<PromptMessage[]>;
shouldRetry?: (context: {
error: unknown;
attempt: number;
options: NonNullable<CopilotStructuredOptions>;
}) => Awaitable<boolean>;
mapError: (error: unknown) => unknown;
};
export type EmbeddingProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
defaultDimensions?: number;
taskType?: string;
mapError: (error: unknown) => unknown;
};
export type RerankProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
mapError: (error: unknown) => unknown;
};
export type ImageProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
prepareMessages?: (
messages: PromptMessage[],
backendConfig: LlmBackendConfig,
options: NonNullable<CopilotImageOptions>
) => Promise<PromptMessage[]>;
mapError: (error: unknown) => unknown;
};
export type ProviderMetricLabels = Record<
string,
string | number | boolean | undefined
>;
export type ProviderExecutionDrivers = {
chat?: ProviderChatDriver;
structured?: StructuredProviderDriver;
embedding?: EmbeddingProviderDriver;
rerank?: RerankProviderDriver;
image?: ImageProviderDriver;
};
export type ProviderDriverSpec = NativeProviderDriverBase & {
chat?: NativeChatDriverOverrides | false;
structured?: NativeStructuredDriverOverrides | false;
embedding?: NativeEmbeddingDriverOverrides | false;
rerank?: NativeRerankDriverOverrides | false;
image?: NativeImageDriverOverrides | false;
};
export type ProviderRuntimeHostSeed = {
model: ProviderModelRuntimeContext;
resolveExecutionDrivers: () => ProviderExecutionDrivers | undefined;
selectModel: NativeChatDriverBase['selectModel'];
checkParams: NativeChatDriverBase['checkParams'];
getAttachCapability: (
model: CopilotProviderModel,
outputType: ModelOutputType
) => ModelAttachmentCapability | undefined;
getActiveProviderMiddleware: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig;
getTools: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>;
metricLabels: (
model: string,
labels?: ProviderMetricLabels,
execution?: CopilotProviderExecution
) => ProviderMetricLabels;
};
export type ProviderChatDriverPrepareInput = {
kind: 'text' | 'streamText' | 'streamObject';
cond: ModelConditions;
messages: PromptMessage[];
options: CopilotChatOptions;
execution?: CopilotProviderExecution;
};
export type ProviderChatDriver = {
prepare: (input: {
kind: ProviderChatDriverPrepareInput['kind'];
cond: ProviderChatDriverPrepareInput['cond'];
messages: ProviderChatDriverPrepareInput['messages'];
options: ProviderChatDriverPrepareInput['options'];
execution?: ProviderChatDriverPrepareInput['execution'];
}) => Promise<ProviderChatDriverPrepareResult | null>;
mapError: (error: unknown) => unknown;
};
type NativeProviderDriverBase = Pick<
StructuredProviderDriver,
'createBackendConfig' | 'mapError'
>;
type ChatToolingResult = Pick<
ProviderChatDriverPrepareResult,
'tools' | 'middleware'
>;
type NativeChatDriverBase = NativeProviderDriverBase & {
checkParams: (input: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions;
withAttachment?: boolean;
execution?: CopilotProviderExecution;
}) => Promise<ModelFullConditions>;
selectModel: (
cond: ModelFullConditions,
execution?: CopilotProviderExecution
) => CopilotProviderModel;
getTools?: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>;
getActiveProviderMiddleware?: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig;
};
type NativeStructuredDriverOverrides = Partial<StructuredProviderDriver>;
type NativeEmbeddingDriverOverrides = Partial<EmbeddingProviderDriver>;
type NativeRerankDriverOverrides = Partial<RerankProviderDriver>;
type NativeImageDriverOverrides = Partial<ImageProviderDriver>;
type NativeChatDriverContext = {
input: ProviderChatDriverPrepareInput;
outputType: ModelOutputType;
normalizedCond: ModelFullConditions;
model: CopilotProviderModel;
backendConfig: LlmBackendConfig;
protocol: LlmProtocol;
messages: PromptMessage[];
options: NonNullable<CopilotChatOptions>;
execution?: CopilotProviderExecution;
};
type NativeChatDriverOverrides = {
resolveOutputType?: (
kind: ProviderChatDriverPrepareInput['kind']
) => ModelOutputType | null;
withAttachment?: boolean;
prepareMessages?: (
context: Omit<NativeChatDriverContext, 'messages'>
) => Awaitable<PromptMessage[]>;
resolveTooling?: (
context: NativeChatDriverContext
) => Awaitable<ChatToolingResult>;
resolveRequestOptions?: (
context: NativeChatDriverContext
) => Awaitable<
Partial<
Pick<
ProviderChatDriverPrepareResult,
'withAttachment' | 'attachmentCapability' | 'include' | 'reasoning'
>
>
>;
};
export function createNativeProviderDriverFactory(
base: NativeProviderDriverBase
) {
return {
structured(
overrides: NativeStructuredDriverOverrides = {}
): StructuredProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
...(overrides.prepareMessages
? { prepareMessages: overrides.prepareMessages }
: {}),
...(overrides.shouldRetry
? { shouldRetry: overrides.shouldRetry }
: {}),
};
},
embedding(
overrides: NativeEmbeddingDriverOverrides = {}
): EmbeddingProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
...(overrides.defaultDimensions !== undefined
? { defaultDimensions: overrides.defaultDimensions }
: {}),
...(overrides.taskType ? { taskType: overrides.taskType } : {}),
};
},
rerank(overrides: NativeRerankDriverOverrides = {}): RerankProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
};
},
image(overrides: NativeImageDriverOverrides = {}): ImageProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
...(overrides.prepareMessages
? { prepareMessages: overrides.prepareMessages }
: {}),
};
},
};
}
function compileProviderChatDriver(
spec: NativeProviderDriverBase & NativeChatDriverOverrides,
base: NativeChatDriverBase
): ProviderChatDriver {
return {
prepare: async (input: ProviderChatDriverPrepareInput) => {
const options: NonNullable<CopilotChatOptions> = input.options ?? {};
const resolvedOutputType = spec.resolveOutputType?.(input.kind);
const outputType =
resolvedOutputType === undefined
? input.kind === 'streamObject'
? ModelOutputType.Object
: ModelOutputType.Text
: resolvedOutputType;
if (!outputType) {
return null;
}
const normalizedCond = await base.checkParams({
messages: input.messages,
cond: {
...input.cond,
outputType,
},
options,
execution: input.execution,
...(spec.withAttachment !== undefined
? { withAttachment: spec.withAttachment }
: {}),
});
const model = base.selectModel(normalizedCond, input.execution);
const backendConfig = await spec.createBackendConfig(input.execution);
const route = resolveProviderModelRoute(model, outputType);
if (!route.protocol) {
throw new Error(`Missing native protocol for model ${model.id}`);
}
const partialContext = {
input,
outputType,
normalizedCond,
model,
backendConfig:
route.requestLayer === backendConfig.request_layer
? backendConfig
: { ...backendConfig, request_layer: route.requestLayer },
protocol: route.protocol,
options,
execution: input.execution,
};
const messages = spec.prepareMessages
? await spec.prepareMessages(partialContext)
: input.messages;
const context = {
...partialContext,
messages,
};
const tooling = spec.resolveTooling
? await spec.resolveTooling(context)
: {
...(base.getTools
? { tools: await base.getTools(options, model.id) }
: {}),
...(base.getActiveProviderMiddleware
? {
middleware: base.getActiveProviderMiddleware(input.execution),
}
: {}),
};
const requestOptions = spec.resolveRequestOptions
? await spec.resolveRequestOptions(context)
: {};
return {
protocol: context.protocol,
backendConfig: context.backendConfig,
model: model.id,
messages,
...(spec.withAttachment === false ? { withAttachment: false } : {}),
...requestOptions,
...tooling,
};
},
mapError: spec.mapError,
};
}
export function createNativeExecutionDriverSpec(
input: ProviderDriverSpec,
runtimeBase: NativeChatDriverBase
): ProviderExecutionDrivers {
const driverBase = {
createBackendConfig: input.createBackendConfig,
mapError: input.mapError,
};
const nativeDrivers = createNativeProviderDriverFactory(driverBase);
return {
...(input.chat !== false
? {
chat: compileProviderChatDriver(
{ ...driverBase, ...input.chat },
runtimeBase
),
}
: {}),
...(input.structured !== false
? {
structured: nativeDrivers.structured(input.structured ?? undefined),
}
: {}),
...(input.embedding !== false
? {
embedding: nativeDrivers.embedding(input.embedding ?? undefined),
}
: {}),
...(input.rerank !== false
? { rerank: nativeDrivers.rerank(input.rerank ?? undefined) }
: {}),
...(input.image !== false
? { image: nativeDrivers.image(input.image ?? undefined) }
: {}),
};
}
@@ -0,0 +1,22 @@
import {
AnthropicOfficialProvider,
AnthropicVertexProvider,
} from './anthropic';
import { CloudflareWorkersAIProvider } from './cloudflare';
import { FalProvider } from './fal';
import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
import { MorphProvider } from './morph';
import { OpenAIProvider } from './openai';
import { PerplexityProvider } from './perplexity';
export const CopilotProviders = [
OpenAIProvider,
CloudflareWorkersAIProvider,
FalProvider,
GeminiGenerativeProvider,
GeminiVertexProvider,
PerplexityProvider,
AnthropicOfficialProvider,
AnthropicVertexProvider,
MorphProvider,
];
@@ -1,377 +1,214 @@
import { AsyncLocalStorage } from 'node:async_hooks';
import { Inject, Injectable, Logger } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import { z } from 'zod';
import {
Config,
CopilotPromptInvalid,
CopilotProviderNotSupported,
OnEvent,
} from '../../../base';
import { DocReader, DocWriter } from '../../../core/doc';
import { AccessController } from '../../../core/permission';
import { Models } from '../../../models';
import { IndexerService } from '../../indexer';
import type { ProviderMiddlewareConfig } from '../config';
import { CopilotContextService } from '../context/service';
import { PromptService } from '../prompt/service';
import {
buildBlobContentGetter,
buildContentGetter,
buildDocContentGetter,
buildDocCreateHandler,
buildDocKeywordSearchGetter,
buildDocSearchGetter,
buildDocUpdateHandler,
buildDocUpdateMetaHandler,
type CopilotTool,
type CopilotToolSet,
createBlobReadTool,
createCodeArtifactTool,
createConversationSummaryTool,
createDocComposeTool,
createDocCreateTool,
createDocEditTool,
createDocKeywordSearchTool,
createDocReadTool,
createDocSemanticSearchTool,
createDocUpdateMetaTool,
createDocUpdateTool,
createExaCrawlTool,
createExaSearchTool,
createSectionEditTool,
} from '../tools';
import { canonicalizePromptAttachment } from './attachments';
import { CopilotProviderFactory } from './factory';
import { Config } from '../../../base';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import { ToolExecutorHost } from '../runtime/hosts/tool-executor-host';
import { mapNativeSemanticError } from '../runtime/native-errors';
import type { ToolLoopBackend } from '../runtime/tool/bridge';
import type { CopilotTool, CopilotToolSet } from '../tools';
import { resolveProviderMiddleware } from './provider-middleware';
import { buildProviderRegistry } from './provider-registry';
import {
checkProviderParams,
getAttachCapability as getAttachCapabilityHelper,
matchProviderModel as matchProviderModelHelper,
type ProviderModelRuntimeContext,
requireProviderModelSelection,
resolveProviderModel,
} from './provider-model-runtime';
import {
type CopilotProviderExecution,
createNativeExecutionDriverSpec,
type ProviderDriverSpec,
type ProviderExecutionDrivers,
type ProviderRuntimeHostSeed,
} from './provider-runtime-contract';
import {
type CopilotChatOptions,
CopilotChatTools,
type CopilotEmbeddingOptions,
type CopilotImageOptions,
type CopilotModelBackendKind,
CopilotProviderModel,
CopilotProviderType,
type CopilotRerankRequest,
CopilotStructuredOptions,
EmbeddingMessage,
type CopilotStructuredOptions,
type ModelAttachmentCapability,
ModelCapability,
ModelConditions,
ModelFullConditions,
ModelInputType,
ModelOutputType,
type PromptAttachmentKind,
type PromptAttachmentSourceKind,
type PromptMessage,
PromptMessageSchema,
StreamObject,
} from './types';
const providerProfileContext = new AsyncLocalStorage<string>();
export type {
CopilotProviderExecution,
ProviderDriverSpec,
ProviderExecutionDrivers,
ProviderRuntimeHostSeed,
} from './provider-runtime-contract';
@Injectable()
export abstract class CopilotProvider<C = any> {
protected readonly logger = new Logger(this.constructor.name);
protected readonly MAX_STEPS = 20;
protected onlineModelList: string[] = [];
abstract readonly type: CopilotProviderType;
abstract readonly models: CopilotProviderModel[];
abstract configured(): boolean;
protected abstract resolveModelBackendKind(
execution?: CopilotProviderExecution
): CopilotModelBackendKind;
abstract configured(execution?: CopilotProviderExecution): boolean;
@Inject() protected readonly AFFiNEConfig!: Config;
@Inject() protected readonly factory!: CopilotProviderFactory;
@Inject() protected readonly moduleRef!: ModuleRef;
readonly #registeredProviderIds = new Set<string>();
@Inject() protected readonly toolExecutorHost!: ToolExecutorHost;
runWithProfile<T>(providerId: string, callback: () => T): T {
return providerProfileContext.run(providerId, callback);
get maxSteps() {
return this.MAX_STEPS;
}
protected getActiveProviderId() {
return providerProfileContext.getStore() ?? `${this.type}-default`;
protected resolveModelRuntimeContext(
execution?: CopilotProviderExecution
): ProviderModelRuntimeContext {
return {
type: this.type,
backendKind: this.resolveModelBackendKind(execution),
};
}
protected getActiveProviderMiddleware(): ProviderMiddlewareConfig {
const providerId = this.getActiveProviderId();
const registry = buildProviderRegistry(this.AFFiNEConfig.copilot.providers);
const profile = registry.profiles.get(providerId);
return profile?.middleware ?? resolveProviderMiddleware(this.type);
protected get modelRuntimeContext(): ProviderModelRuntimeContext {
return this.resolveModelRuntimeContext();
}
protected metricLabels(
getDriverSpec(): ProviderDriverSpec | undefined {
return undefined;
}
getExecutionDrivers(): ProviderExecutionDrivers | undefined {
const spec = this.getDriverSpec();
return spec ? this.createDriverSpec(spec) : undefined;
}
protected createDriverSpec(
spec: ProviderDriverSpec
): ProviderExecutionDrivers {
return createNativeExecutionDriverSpec(spec, {
createBackendConfig: spec.createBackendConfig,
mapError: error => {
const mapped = mapNativeSemanticError(error);
return mapped === error ? spec.mapError(error) : mapped;
},
checkParams: input =>
checkProviderParams(
this.resolveModelRuntimeContext(input.execution),
input
),
selectModel: (cond, execution) =>
requireProviderModelSelection(
this.resolveModelRuntimeContext(execution),
cond
),
getTools: this.getTools.bind(this),
getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this),
});
}
selectModel(
cond: ModelFullConditions,
execution?: CopilotProviderExecution
): CopilotProviderModel {
return requireProviderModelSelection(
this.resolveModelRuntimeContext(execution),
cond
);
}
checkParams(input: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions;
withAttachment?: boolean;
execution?: CopilotProviderExecution;
}) {
return checkProviderParams(
this.resolveModelRuntimeContext(input.execution),
input
);
}
getRuntimeHostSeed(): ProviderRuntimeHostSeed {
return {
model: this.resolveModelRuntimeContext(),
resolveExecutionDrivers: () => this.getExecutionDrivers(),
selectModel: this.selectModel.bind(this),
checkParams: this.checkParams.bind(this),
getAttachCapability: this.getAttachCapability.bind(this),
getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this),
getTools: this.getTools.bind(this),
metricLabels: this.metricLabels.bind(this),
};
}
protected getExecutionProfile(execution?: CopilotProviderExecution) {
return execution?.profile?.type === this.type
? execution.profile
: undefined;
}
getActiveProviderMiddleware(
execution?: CopilotProviderExecution
): ProviderMiddlewareConfig {
return (
this.getExecutionProfile(execution)?.middleware ??
resolveProviderMiddleware(this.type)
);
}
metricLabels(
model: string,
labels: Record<string, string | number | boolean | undefined> = {}
labels: Record<string, string | number | boolean | undefined> = {},
execution?: CopilotProviderExecution
) {
const providerId = this.getActiveProviderId();
return { model, providerId, ...labels };
return {
model,
providerId: execution?.providerId ?? `${this.type}-default`,
...labels,
};
}
get config(): C {
const profileId = providerProfileContext.getStore();
if (profileId) {
const profile = this.AFFiNEConfig.copilot.providers.profiles?.find(
profile => profile.id === profileId && profile.type === this.type
);
if (profile) return profile.config as C;
}
protected get config(): C {
return this.AFFiNEConfig.copilot.providers[this.type] as C;
}
@OnEvent('config.init')
async onConfigInit() {
this.setup();
}
@OnEvent('config.changed')
async onConfigChanged(event: Events['config.changed']) {
if ('copilot' in event.updates) {
this.setup();
protected getConfig(execution?: CopilotProviderExecution): C {
const profile = this.getExecutionProfile(execution);
if (profile) {
return profile.config as C;
}
return this.config;
}
protected setup() {
const registry = buildProviderRegistry(this.AFFiNEConfig.copilot.providers);
const providerIds = registry.byType.get(this.type) ?? [];
const nextProviderIds = new Set<string>();
for (const id of providerIds) {
const configured = this.runWithProfile(id, () => this.configured());
if (configured) {
nextProviderIds.add(id);
this.factory.register(id, this);
} else {
this.factory.unregister(id, this);
}
}
for (const providerId of this.#registeredProviderIds) {
if (!nextProviderIds.has(providerId)) {
this.factory.unregister(providerId, this);
}
}
this.#registeredProviderIds.clear();
for (const providerId of nextProviderIds) {
this.#registeredProviderIds.add(providerId);
}
if (env.selfhosted && nextProviderIds.size > 0) {
const [providerId] = Array.from(nextProviderIds);
this.runWithProfile(providerId, () => {
this.refreshOnlineModels().catch(e =>
this.logger.error('Failed to refresh online models', e)
);
});
}
}
async refreshOnlineModels() {}
private unique<T>(values: Iterable<T>) {
return Array.from(new Set(values));
}
private attachmentKindToInputType(
kind: PromptAttachmentKind
): ModelInputType {
switch (kind) {
case 'image':
return ModelInputType.Image;
case 'audio':
return ModelInputType.Audio;
default:
return ModelInputType.File;
}
}
protected async inferModelConditionsFromMessages(
messages?: PromptMessage[],
withAttachment = true
): Promise<Partial<ModelFullConditions>> {
if (!messages?.length || !withAttachment) return {};
const attachmentKinds: PromptAttachmentKind[] = [];
const attachmentSourceKinds: PromptAttachmentSourceKind[] = [];
const inputTypes: ModelInputType[] = [];
let hasRemoteAttachments = false;
for (const message of messages) {
if (!Array.isArray(message.attachments)) continue;
for (const attachment of message.attachments) {
const normalized = await canonicalizePromptAttachment(
attachment,
message
);
attachmentKinds.push(normalized.kind);
inputTypes.push(this.attachmentKindToInputType(normalized.kind));
attachmentSourceKinds.push(normalized.sourceKind);
hasRemoteAttachments = hasRemoteAttachments || normalized.isRemote;
}
}
return {
...(attachmentKinds.length
? { attachmentKinds: this.unique(attachmentKinds) }
: {}),
...(attachmentSourceKinds.length
? { attachmentSourceKinds: this.unique(attachmentSourceKinds) }
: {}),
...(inputTypes.length ? { inputTypes: this.unique(inputTypes) } : {}),
...(hasRemoteAttachments ? { hasRemoteAttachments } : {}),
};
}
private mergeModelConditions(
cond: ModelFullConditions,
inferredCond: Partial<ModelFullConditions>
): ModelFullConditions {
return {
...inferredCond,
...cond,
inputTypes: this.unique([
...(inferredCond.inputTypes ?? []),
...(cond.inputTypes ?? []),
]),
attachmentKinds: this.unique([
...(inferredCond.attachmentKinds ?? []),
...(cond.attachmentKinds ?? []),
]),
attachmentSourceKinds: this.unique([
...(inferredCond.attachmentSourceKinds ?? []),
...(cond.attachmentSourceKinds ?? []),
]),
hasRemoteAttachments:
cond.hasRemoteAttachments ?? inferredCond.hasRemoteAttachments,
};
}
protected getAttachCapability(
getAttachCapability(
model: CopilotProviderModel,
outputType: ModelOutputType
): ModelAttachmentCapability | undefined {
const capability =
model.capabilities.find(cap => cap.output.includes(outputType)) ??
model.capabilities[0];
if (!capability) {
return;
}
return this.resolveAttachmentCapability(capability, outputType);
}
private resolveAttachmentCapability(
cap: ModelCapability,
outputType?: ModelOutputType
): ModelAttachmentCapability | undefined {
if (outputType === ModelOutputType.Structured) {
return cap.structuredAttachments ?? cap.attachments;
}
return cap.attachments;
}
private matchesAttachCapability(
cap: ModelCapability,
cond: ModelFullConditions
) {
const {
attachmentKinds,
attachmentSourceKinds,
hasRemoteAttachments,
outputType,
} = cond;
if (
!attachmentKinds?.length &&
!attachmentSourceKinds?.length &&
!hasRemoteAttachments
) {
return true;
}
const attachmentCapability = this.resolveAttachmentCapability(
cap,
outputType
);
if (!attachmentCapability) {
return !attachmentKinds?.some(
kind => !cap.input.includes(this.attachmentKindToInputType(kind))
);
}
if (
attachmentKinds?.some(kind => !attachmentCapability.kinds.includes(kind))
) {
return false;
}
if (
attachmentSourceKinds?.length &&
attachmentCapability.sourceKinds?.length &&
attachmentSourceKinds.some(
kind => !attachmentCapability.sourceKinds?.includes(kind)
)
) {
return false;
}
if (
hasRemoteAttachments &&
attachmentCapability.allowRemoteUrls === false
) {
return false;
}
return true;
}
private findValidModel(
cond: ModelFullConditions
): CopilotProviderModel | undefined {
const { modelId, outputType, inputTypes } = cond;
const matcher = (cap: ModelCapability) =>
(!outputType || cap.output.includes(outputType)) &&
(!inputTypes?.length ||
inputTypes.every(type => cap.input.includes(type))) &&
this.matchesAttachCapability(cap, cond);
if (modelId) {
const hasOnlineModel = this.onlineModelList.includes(modelId);
const model = this.models.find(
m => m.id === modelId && m.capabilities.some(matcher)
);
if (model) return model;
// allow online model without capabilities check
if (hasOnlineModel) return { id: modelId, capabilities: [] };
return undefined;
}
if (!outputType) return undefined;
return this.models.find(m =>
m.capabilities.some(c => matcher(c) && c.defaultForOutputType)
);
return getAttachCapabilityHelper(model, outputType);
}
// make it async to allow dynamic check available models in some providers
async match(cond: ModelFullConditions = {}): Promise<boolean> {
return this.configured() && !!this.findValidModel(cond);
async match(
cond: ModelFullConditions = {},
execution?: CopilotProviderExecution
): Promise<boolean> {
return (
this.configured(execution) &&
matchProviderModelHelper(this.resolveModelRuntimeContext(execution), cond)
);
}
protected selectModel(cond: ModelFullConditions): CopilotProviderModel {
const model = this.findValidModel(cond);
if (model) return model;
const { modelId, outputType, inputTypes } = cond;
throw new CopilotPromptInvalid(
resolveModel(
modelId: string,
execution?: CopilotProviderExecution
): CopilotProviderModel | undefined {
return resolveProviderModel(
this.resolveModelRuntimeContext(execution),
modelId
? `Model ${modelId} does not support ${outputType ?? '<any>'} output with ${inputTypes ?? '<any>'} input`
: outputType
? `No model supports ${outputType} output with ${inputTypes ?? '<any>'} input for provider ${this.type}`
: 'Output type is required when modelId is not provided'
);
}
@@ -383,302 +220,30 @@ export abstract class CopilotProvider<C = any> {
}
// use for tool use, shared between providers
protected async getTools(
async getTools(
options: CopilotChatOptions,
model: string
): Promise<CopilotToolSet> {
const tools: CopilotToolSet = {};
if (options?.tools?.length) {
this.logger.debug(`getTools: ${JSON.stringify(options.tools)}`);
const ac = this.moduleRef.get(AccessController, { strict: false });
const context = this.moduleRef.get(CopilotContextService, {
strict: false,
});
const docReader = this.moduleRef.get(DocReader, { strict: false });
const docWriter = this.moduleRef.get(DocWriter, { strict: false });
const models = this.moduleRef.get(Models, { strict: false });
const prompt = this.moduleRef.get(PromptService, {
strict: false,
});
for (const tool of options.tools) {
const toolDef = this.getProviderSpecificTools(tool, model);
if (toolDef) {
// allow provider prevent tool creation
if (toolDef[1]) {
tools[toolDef[0]] = toolDef[1];
}
continue;
}
if (
!(env.dev || env.namespaces.canary) &&
['docCreate', 'docUpdate', 'docUpdateMeta'].includes(tool)
) {
continue;
}
switch (tool) {
case 'blobRead': {
const docContext = options.session
? await context.getBySessionId(options.session)
: null;
const getBlobContent = buildBlobContentGetter(ac, docContext);
tools.blob_read = createBlobReadTool(
getBlobContent.bind(null, options)
);
break;
}
case 'codeArtifact': {
tools.code_artifact = createCodeArtifactTool(prompt, this.factory);
break;
}
case 'conversationSummary': {
tools.conversation_summary = createConversationSummaryTool(
options.session,
prompt,
this.factory
);
break;
}
case 'docEdit': {
const getDocContent = buildContentGetter(ac, docReader);
tools.doc_edit = createDocEditTool(
this.factory,
prompt,
getDocContent.bind(null, options)
);
break;
}
case 'docSemanticSearch': {
const docContext = options.session
? await context.getBySessionId(options.session)
: null;
const searchDocs = buildDocSearchGetter(
ac,
context,
docContext,
models
);
tools.doc_semantic_search = createDocSemanticSearchTool(
searchDocs.bind(null, options)
);
break;
}
case 'docKeywordSearch': {
if (this.AFFiNEConfig.indexer.enabled) {
const indexerService = this.moduleRef.get(IndexerService, {
strict: false,
});
const searchDocs = buildDocKeywordSearchGetter(
ac,
indexerService,
models
);
tools.doc_keyword_search = createDocKeywordSearchTool(
searchDocs.bind(null, options)
);
}
break;
}
case 'docRead': {
const getDoc = buildDocContentGetter(ac, docReader, models);
tools.doc_read = createDocReadTool(getDoc.bind(null, options));
break;
}
case 'docCreate': {
const createDoc = buildDocCreateHandler(ac, docWriter);
tools.doc_create = createDocCreateTool(
createDoc.bind(null, options)
);
break;
}
case 'docUpdate': {
const updateDoc = buildDocUpdateHandler(ac, docWriter);
tools.doc_update = createDocUpdateTool(
updateDoc.bind(null, options)
);
break;
}
case 'docUpdateMeta': {
const updateDocMeta = buildDocUpdateMetaHandler(ac, docWriter);
tools.doc_update_meta = createDocUpdateMetaTool(
updateDocMeta.bind(null, options)
);
break;
}
case 'webSearch': {
tools.web_search_exa = createExaSearchTool(this.AFFiNEConfig);
tools.web_crawl_exa = createExaCrawlTool(this.AFFiNEConfig);
break;
}
case 'docCompose': {
tools.doc_compose = createDocComposeTool(prompt, this.factory);
break;
}
case 'sectionEdit': {
tools.section_edit = createSectionEditTool(prompt, this.factory);
break;
}
}
}
return tools;
}
return tools;
}
private handleZodError(ret: z.SafeParseReturnType<any, any>) {
if (ret.success) return;
const issues = ret.error.issues.map(i => {
const path =
'root' +
(i.path.length
? `.${i.path.map(seg => (typeof seg === 'number' ? `[${seg}]` : `.${seg}`)).join('')}`
: '');
return `${i.message}${path}`;
});
throw new CopilotPromptInvalid(issues.join('; '));
}
protected async checkParams({
cond,
messages,
embeddings,
options = {},
withAttachment = true,
}: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?: CopilotChatOptions | CopilotStructuredOptions;
withAttachment?: boolean;
}): Promise<ModelFullConditions> {
if (messages) {
const { requireContent = true, requireAttachment = false } = options;
const MessageSchema = z
.array(
PromptMessageSchema.extend({
content: requireContent
? z.string().trim().min(1)
: z.string().optional().nullable(),
})
.passthrough()
.catchall(z.union([z.string(), z.number(), z.date(), z.null()]))
)
.optional();
this.handleZodError(MessageSchema.safeParse(messages));
const inferredCond = await this.inferModelConditionsFromMessages(
messages,
withAttachment
);
const mergedCond = this.mergeModelConditions(cond, inferredCond);
const model = this.selectModel(mergedCond);
const multimodal = model.capabilities.some(c =>
[ModelInputType.Image, ModelInputType.Audio, ModelInputType.File].some(
t => c.input.includes(t)
)
);
if (
multimodal &&
requireAttachment &&
!messages.some(
message =>
message.role === 'user' &&
Array.isArray(message.attachments) &&
message.attachments.length > 0
)
) {
throw new CopilotPromptInvalid(
'attachments required in multimodal mode'
);
}
if (embeddings) {
this.handleZodError(EmbeddingMessage.safeParse(embeddings));
}
return mergedCond;
}
const inferredCond = await this.inferModelConditionsFromMessages(
messages,
withAttachment
this.logger.debug(`getTools: ${JSON.stringify(options?.tools ?? [])}`);
return await this.toolExecutorHost.getTools(
options,
model,
this.getProviderSpecificTools.bind(this)
);
const mergedCond = this.mergeModelConditions(cond, inferredCond);
if (embeddings) {
this.handleZodError(EmbeddingMessage.safeParse(embeddings));
}
return mergedCond;
}
abstract text(
model: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions
): Promise<string>;
abstract streamText(
model: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions
): AsyncIterable<string>;
streamObject(
_model: ModelConditions,
_messages: PromptMessage[],
_options?: CopilotChatOptions
): AsyncIterable<StreamObject> {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'object',
});
}
structure(
_cond: ModelConditions,
_messages: PromptMessage[],
_options?: CopilotStructuredOptions
): Promise<string> {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'structure',
});
}
streamImages(
_model: ModelConditions,
_messages: PromptMessage[],
_options?: CopilotImageOptions
): AsyncIterable<string> {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'image',
});
}
embedding(
_model: ModelConditions,
_text: string | string[],
_options?: CopilotEmbeddingOptions
): Promise<number[][]> {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'embedding',
});
}
async rerank(
_model: ModelConditions,
_request: CopilotRerankRequest,
_options?: CopilotChatOptions
): Promise<number[]> {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'rerank',
createNativeAdapter(
backend: ToolLoopBackend,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[],
options: {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
} = {}
) {
return this.toolExecutorHost.createNativeAdapter(backend, tools, {
...options,
nodeTextMiddleware: nodeTextMiddleware ?? options.nodeTextMiddleware,
});
}
}
@@ -0,0 +1,28 @@
import { Injectable } from '@nestjs/common';
import { Config } from '../../../base';
import {
buildProviderRegistry,
type CopilotProviderRegistry,
type CopilotProvidersConfigInput,
} from './provider-registry';
@Injectable()
export class CopilotProviderRegistryService {
private lastConfig?: CopilotProvidersConfigInput;
private lastRegistry?: CopilotProviderRegistry;
constructor(private readonly config: Config) {}
getRegistry(): CopilotProviderRegistry {
const providerConfig = this.config.copilot.providers;
if (this.lastConfig === providerConfig && this.lastRegistry) {
return this.lastRegistry;
}
const registry = buildProviderRegistry(providerConfig);
this.lastConfig = providerConfig;
this.lastRegistry = registry;
return registry;
}
}
@@ -2,6 +2,23 @@ import { AiPromptRole } from '@prisma/client';
import { z } from 'zod';
import { JSONSchema } from '../../../base';
import type {
CapabilityAttachmentContract,
CapabilityModelCapability,
ModelConditionsContract,
} from '../../../native';
import type { CopilotModelBackendKind } from '../runtime/contracts';
import {
type StreamObject,
StreamObjectSchema,
} from '../runtime/contracts/runtime-event-contract';
// Owner map:
// - provider/profile/config schemas in this file are backend host ingress.
// - prompt/message/attachment Zod schemas validate Node host ingress and
// persistence surfaces before values cross into native prompt DTOs.
// - model condition/capability types are native-generated facades.
// - StreamObject is app-facing projection, not runtime event truth.
// ========== provider ==========
@@ -163,6 +180,8 @@ const PromptAttachmentSchema = z.discriminatedUnion('kind', [
.object({
kind: z.literal('url'),
url: AttachmentUrlSchema,
data: z.string().optional(),
encoding: z.literal('base64').optional(),
mimeType: z.string().optional(),
fileName: z.string().optional(),
providerHint: AttachmentProviderHintSchema.optional(),
@@ -211,34 +230,14 @@ export const ChatMessageAttachment = z.union([
export const PromptResponseFormatSchema = z
.object({
type: z.literal('json_schema'),
schema: z.any(),
responseSchemaJson: z.record(z.unknown()).optional(),
schemaHash: z.string().optional(),
strict: z.boolean().optional(),
})
.strict();
export const StreamObjectSchema = z.discriminatedUnion('type', [
z.object({
type: z.literal('text-delta'),
textDelta: z.string(),
}),
z.object({
type: z.literal('reasoning'),
textDelta: z.string(),
}),
z.object({
type: z.literal('tool-call'),
toolCallId: z.string(),
toolName: z.string(),
args: z.record(z.any()),
}),
z.object({
type: z.literal('tool-result'),
toolCallId: z.string(),
toolName: z.string(),
args: z.record(z.any()),
result: z.any(),
}),
]);
.strict()
.refine(value => value.responseSchemaJson !== undefined, {
message: 'responseSchemaJson is required',
});
export const PureMessageSchema = z.object({
content: z.string(),
@@ -253,7 +252,8 @@ export const PromptMessageSchema = PureMessageSchema.extend({
}).strict();
export type PromptMessage = z.infer<typeof PromptMessageSchema>;
export type PromptParams = NonNullable<PromptMessage['params']>;
export type StreamObject = z.infer<typeof StreamObjectSchema>;
export { StreamObjectSchema };
export type { StreamObject };
export type PromptAttachment = z.infer<typeof ChatMessageAttachment>;
export type PromptAttachmentSourceKind = z.infer<
typeof PromptAttachmentSourceKindSchema
@@ -286,7 +286,11 @@ export type CopilotChatTools = NonNullable<
export const CopilotStructuredOptionsSchema =
CopilotProviderOptionsSchema.merge(PromptConfigStrictSchema)
.extend({ schema: z.any().optional(), strict: z.boolean().optional() })
.extend({
responseSchemaJson: z.record(z.unknown()).optional(),
schemaHash: z.string().optional(),
strict: z.boolean().optional(),
})
.optional();
export type CopilotStructuredOptions = z.infer<
@@ -299,6 +303,16 @@ export const CopilotImageOptionsSchema = CopilotProviderOptionsSchema.merge(
.extend({
quality: z.string().optional(),
seed: z.number().optional(),
modelName: z.string().nullable().optional(),
loras: z
.array(
z.object({
path: z.string(),
scale: z.number().nullable().optional(),
})
)
.nullable()
.optional(),
})
.optional();
@@ -306,7 +320,7 @@ export type CopilotImageOptions = z.infer<typeof CopilotImageOptionsSchema>;
export const CopilotEmbeddingOptionsSchema =
CopilotProviderOptionsSchema.extend({
dimensions: z.number(),
dimensions: z.number().optional(),
}).optional();
export type CopilotEmbeddingOptions = z.infer<
@@ -324,35 +338,29 @@ export type CopilotRerankRequest = {
topK?: number;
};
export enum ModelInputType {
Text = 'text',
Image = 'image',
Audio = 'audio',
File = 'file',
}
export const ModelInputType = {
Text: 'text',
Image: 'image',
Audio: 'audio',
File: 'file',
} as const;
export enum ModelOutputType {
Text = 'text',
Object = 'object',
Embedding = 'embedding',
Image = 'image',
Rerank = 'rerank',
Structured = 'structured',
}
export type ModelInputType = CapabilityModelCapability['input'][number];
export interface ModelAttachmentCapability {
kinds: PromptAttachmentKind[];
sourceKinds?: PromptAttachmentSourceKind[];
allowRemoteUrls?: boolean;
}
export const ModelOutputType = {
Text: 'text',
Object: 'object',
Embedding: 'embedding',
Image: 'image',
Rerank: 'rerank',
Structured: 'structured',
} as const;
export interface ModelCapability {
input: ModelInputType[];
output: ModelOutputType[];
attachments?: ModelAttachmentCapability;
structuredAttachments?: ModelAttachmentCapability;
defaultForOutputType?: boolean;
}
export type ModelOutputType = CapabilityModelCapability['output'][number];
export type ModelAttachmentCapability = CapabilityAttachmentContract;
export type ModelCapability = CapabilityModelCapability;
export interface CopilotProviderModel {
id: string;
@@ -360,14 +368,8 @@ export interface CopilotProviderModel {
capabilities: ModelCapability[];
}
export type ModelConditions = {
inputTypes?: ModelInputType[];
attachmentKinds?: PromptAttachmentKind[];
attachmentSourceKinds?: PromptAttachmentSourceKind[];
hasRemoteAttachments?: boolean;
modelId?: string;
};
export type { CopilotModelBackendKind };
export type ModelFullConditions = ModelConditions & {
outputType?: ModelOutputType;
};
export type ModelConditions = Omit<ModelConditionsContract, 'outputType'>;
export type ModelFullConditions = ModelConditionsContract;
@@ -1,5 +1,3 @@
import { createHash } from 'node:crypto';
import { BadRequestException, NotFoundException } from '@nestjs/common';
import {
Args,
@@ -24,37 +22,29 @@ import {
CopilotProviderSideError,
CopilotSessionNotFound,
type FileUpload,
ImageFormatNotSupported,
paginate,
Paginated,
PaginationInput,
RequestMutex,
sniffMime,
Throttle,
TooManyRequest,
UserFriendlyError,
} from '../../base';
import { CurrentUser } from '../../core/auth';
import { Admin } from '../../core/common';
import { DocReader } from '../../core/doc';
import {
AccessController,
DocAction,
WorkspacePolicyService,
} from '../../core/permission';
import { AccessController, DocAction } from '../../core/permission';
import { UserType } from '../../core/user';
import type { ListSessionOptions, UpdateChatSession } from '../../models';
import { processImage } from '../../native';
import { CopilotCronJobs } from './cron';
import { CompatHistoryProjector } from './compat/history-projector';
import { ConversationInboxService } from './conversation/inbox';
import { PromptService } from './prompt/service';
import { CopilotProviderFactory } from './providers/factory';
import type { PromptMessage, StreamObject } from './providers/types';
import { ModelOutputType, type StreamObject } from './providers/types';
import { CapabilityRuntime } from './runtime/capability-runtime';
import { ChatSessionService } from './session';
import { CopilotStorage } from './storage';
import { type ChatHistory, type ChatMessage, SubmittedMessage } from './types';
export const COPILOT_LOCKER = 'copilot';
const COPILOT_IMAGE_MAX_EDGE = 1536;
// ================== Input Types ==================
@@ -382,12 +372,13 @@ export class CopilotResolver {
constructor(
private readonly ac: AccessController,
private readonly mutex: RequestMutex,
private readonly policy: WorkspacePolicyService,
private readonly prompt: PromptService,
private readonly chatSession: ChatSessionService,
private readonly storage: CopilotStorage,
private readonly historyProjector: CompatHistoryProjector,
private readonly inbox: ConversationInboxService,
private readonly docReader: DocReader,
private readonly providerFactory: CopilotProviderFactory
private readonly providerFactory: CopilotProviderFactory,
private readonly runtime: CapabilityRuntime
) {}
@ResolveField(() => CopilotQuotaType, {
@@ -436,33 +427,36 @@ export class CopilotResolver {
if (!prompt) {
throw new NotFoundException('Prompt not found');
}
const convertModels = (ids: string[]) => {
return ids
.map(id => ({ id, name: this.modelNames.get(id) }))
.filter(m => !!m.name) as CopilotModelType[];
const convertModels = async (ids: string[]) => {
const models = await Promise.all(
ids.map(async id => {
const cachedName = this.modelNames.get(id);
if (cachedName) return { id, name: cachedName };
const resolved = await this.providerFactory.resolveProvider({
modelId: id,
outputType: ModelOutputType.Text,
});
const name = resolved?.provider.resolveModel(
resolved.modelId ?? id,
resolved.execution
)?.name;
if (name) {
this.modelNames.set(id, name);
return { id, name };
}
return null;
})
);
return models.filter(model => !!model) as CopilotModelType[];
};
const proModels = prompt.config?.proModels || [];
const missing = new Set(
[...prompt.optionalModels, ...proModels].filter(
id => !this.modelNames.has(id)
)
);
if (missing.size) {
for (const model of missing) {
if (this.modelNames.has(model)) continue;
const provider = await this.providerFactory.getProviderByModel(model);
if (provider?.configured()) {
for (const m of provider.models) {
if (m.name) this.modelNames.set(m.id, m.name);
}
}
}
}
return {
defaultModel: prompt.model,
optionalModels: convertModels(prompt.optionalModels),
proModels: convertModels(proModels),
optionalModels: await convertModels(prompt.optionalModels),
proModels: await convertModels(proModels),
};
}
@@ -476,11 +470,20 @@ export class CopilotResolver {
@Args('sessionId') sessionId: string
): Promise<CopilotSessionType> {
await this.assertPermission(user, copilot);
const session = await this.chatSession.getSessionInfo(sessionId);
if (!session) {
const state = await this.chatSession.getMetaState(sessionId);
if (!state) {
throw new NotFoundException('Session not found');
}
return this.transformToSessionType(session);
const projected = this.historyProjector.projectSession(state, {
requestUserId: user.id,
skipVisibilityFilter: true,
});
if (!projected) {
throw new NotFoundException('Session not found');
}
return this.transformToSessionType(projected);
}
@ResolveField(() => [CopilotSessionType], {
@@ -503,12 +506,19 @@ export class CopilotResolver {
Object.assign({}, copilot, { docId: maybeDocId })
);
const sessions = await this.chatSession.list(
Object.assign({}, options, appendOptions),
false
);
const sessions = (
await this.chatSession.listMetaStates(
Object.assign({}, options, appendOptions)
)
)
.map(state =>
this.historyProjector.projectSession(state, {
requestUserId: user.id,
})
)
.filter((history): history is Omit<ChatHistory, 'messages'> => !!history);
if (appendOptions.docId) {
type Session = ChatHistory & { docId: string };
type Session = Omit<ChatHistory, 'messages'> & { docId: string };
const filtered = sessions.filter((s): s is Session => !!s.docId);
const accessible = await this.ac
.user(user.id)
@@ -537,10 +547,20 @@ export class CopilotResolver {
await this.assertPermission(user, { workspaceId, docId }, 'Doc.Read');
}
const histories = await this.chatSession.list(
Object.assign({}, options, { userId: user.id, workspaceId, docId }),
true
);
const histories = (
await this.chatSession.listStates(
Object.assign({}, options, { userId: user.id, workspaceId, docId })
)
)
.map(state =>
this.historyProjector.projectHistory(state, {
requestUserId: user.id,
withMessages: true,
withPrompt: options?.withPrompt,
action: options?.action,
})
)
.filter((history): history is ChatHistory => !!history);
return histories.map(h => ({
...h,
@@ -574,16 +594,31 @@ export class CopilotResolver {
{ skip: pagination.offset, limit: pagination.first }
);
const totalCount = await this.chatSession.count(finalOptions);
const histories = await this.chatSession.list(
finalOptions,
!!options?.withMessages
);
const histories: ChatHistory[] = options?.withMessages
? (await this.chatSession.listStates(finalOptions))
.map(state =>
this.historyProjector.projectHistory(state, {
requestUserId: user.id,
withMessages: true,
withPrompt: options?.withPrompt,
action: options?.action,
})
)
.filter((history): history is ChatHistory => !!history)
: (await this.chatSession.listMetaStates(finalOptions)).flatMap(state => {
const session = this.historyProjector.projectSession(state, {
requestUserId: user.id,
});
return session
? [{ ...session, messages: [] as ChatHistory['messages'] }]
: [];
});
return paginate(
histories.map(h => ({
...h,
// filter out empty messages
messages: h.messages?.filter(
messages: h.messages.filter(
m => m.content || m.attachments?.length
) as ChatMessageType[],
})),
@@ -639,7 +674,16 @@ export class CopilotResolver {
options: CreateChatSessionInput
): Promise<CopilotHistoriesType> {
const sessionId = await this.createCopilotSessionInternal(user, options);
const session = await this.chatSession.getSessionInfo(sessionId);
const state = await this.chatSession.getState(sessionId);
if (!state) {
throw new NotFoundException('Session not found');
}
const session = this.historyProjector.projectHistory(state, {
requestUserId: user.id,
withMessages: true,
withPrompt: false,
action: !!state.prompt.action,
});
if (!session) {
throw new NotFoundException('Session not found');
}
@@ -768,61 +812,8 @@ export class CopilotResolver {
if (!lock) {
throw new TooManyRequest('Server is busy');
}
const session = await this.chatSession.get(options.sessionId);
if (!session || session.config.userId !== user.id) {
throw new BadRequestException('Session not found');
}
const attachments: PromptMessage['attachments'] = options.attachments || [];
if (options.blob || options.blobs) {
const { workspaceId } = session.config;
const blobs = await Promise.all(
options.blob ? [options.blob] : options.blobs || []
);
delete options.blob;
delete options.blobs;
if (blobs.length) {
await this.policy.assertCanUploadBlob(user.id, workspaceId);
}
for (const blob of blobs) {
const uploaded = await this.storage.handleUpload(user.id, blob);
const detectedMime =
sniffMime(uploaded.buffer, blob.mimetype)?.toLowerCase() ||
blob.mimetype;
let attachmentBuffer = uploaded.buffer;
let attachmentMimeType = detectedMime;
if (detectedMime.startsWith('image/')) {
try {
attachmentBuffer = await processImage(
uploaded.buffer,
COPILOT_IMAGE_MAX_EDGE,
true
);
attachmentMimeType = 'image/webp';
} catch {
throw new ImageFormatNotSupported({ format: detectedMime });
}
}
const filename = createHash('sha256')
.update(attachmentBuffer)
.digest('base64url');
const attachment = await this.storage.put(
user.id,
workspaceId,
filename,
attachmentBuffer
);
attachments.push({ attachment, mimeType: attachmentMimeType });
}
}
try {
return await this.chatSession.createMessage({ ...options, attachments });
return await this.inbox.createMessage(user.id, options);
} catch (e: any) {
throw new CopilotFailedToCreateMessage(e.message);
}
@@ -886,15 +877,16 @@ export class CopilotResolver {
const markdown = docContent.markdown.trim();
// Get LLM provider
const provider =
await this.providerFactory.getProviderByModel('morph-v3-large');
if (!provider) {
const resolved = await this.providerFactory.resolveProvider({
modelId: 'morph-v3-large',
outputType: ModelOutputType.Text,
});
if (!resolved) {
throw new BadRequestException('No LLM provider available');
}
try {
return await provider.text(
return await this.runtime.text(
{ modelId: 'morph-v3-large' },
[
{
@@ -909,7 +901,7 @@ export class CopilotResolver {
throw e;
} else {
throw new CopilotProviderSideError({
provider: provider.type,
provider: resolved.provider.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected apply response',
});
@@ -917,9 +909,7 @@ export class CopilotResolver {
}
}
private transformToSessionType(
session: Omit<ChatHistory, 'messages'>
): CopilotSessionType {
private transformToSessionType(session: Omit<ChatHistory, 'messages'>) {
return { id: session.sessionId, ...session };
}
}
@@ -944,25 +934,3 @@ export class UserCopilotResolver {
return { workspaceId: workspaceId || null };
}
}
@Admin()
@Resolver(() => String)
export class PromptsManagementResolver {
constructor(private readonly cron: CopilotCronJobs) {}
@Mutation(() => Boolean, {
description: 'Trigger generate missing titles cron job',
})
async triggerGenerateTitleCron() {
await this.cron.triggerGenerateMissingTitles();
return true;
}
@Mutation(() => Boolean, {
description: 'Trigger cleanup of trashed doc embeddings',
})
async triggerCleanupTrashedDocEmbeddings() {
await this.cron.triggerCleanupTrashedDocEmbeddings();
return true;
}
}
@@ -0,0 +1,192 @@
import type { Turn } from '../core';
import type { ChatSession } from '../session';
import type { ActionRuntimeBridgeEvent } from './action-runtime-bridge';
type ProjectedAssistantTurn = {
content: string;
attachments: string[];
metadata: Record<string, unknown>;
};
type ActionResultProjector = (
result: unknown,
artifacts: unknown[]
) => ProjectedAssistantTurn;
export function summarizeActionResult(result: unknown) {
if (typeof result === 'string') {
return result.slice(0, 500);
}
if (result === undefined || result === null) {
return '';
}
return JSON.stringify(result).slice(0, 500);
}
function stringArray(value: unknown): string[] {
return Array.isArray(value)
? value.filter((item): item is string => typeof item === 'string')
: [];
}
function attachmentUrls(value: unknown): string[] {
if (Array.isArray(value)) {
return value.flatMap(attachmentUrls);
}
if (typeof value === 'string') {
return [value];
}
if (value && typeof value === 'object' && 'url' in value) {
const url = (value as { url?: unknown }).url;
return typeof url === 'string' ? [url] : [];
}
return [];
}
function textResult(result: unknown): string {
if (typeof result === 'string') {
return result;
}
if (result && typeof result === 'object') {
const value = result as {
content?: unknown;
text?: unknown;
result?: unknown;
params?: unknown;
};
if (typeof value.content === 'string') return value.content;
if (typeof value.text === 'string') return value.text;
if (typeof value.result === 'string') return value.result;
}
throw new Error('Action result does not match text output contract');
}
function metadataFromParams(result: unknown): Record<string, unknown> {
if (!result || typeof result !== 'object') return {};
const params = (result as { params?: unknown }).params;
return params && typeof params === 'object' && !Array.isArray(params)
? (params as Record<string, unknown>)
: {};
}
function projectTextResult(result: unknown): ProjectedAssistantTurn {
return {
content: textResult(result),
attachments: stringArray(
result && typeof result === 'object'
? (result as { attachments?: unknown }).attachments
: undefined
),
metadata: metadataFromParams(result),
};
}
function projectImageResult(
result: unknown,
artifacts: unknown[]
): ProjectedAssistantTurn {
const attachments = [
...attachmentUrls(artifacts),
...attachmentUrls(
result && typeof result === 'object'
? (result as { attachments?: unknown }).attachments
: undefined
),
...attachmentUrls(result),
];
if (!attachments.length) {
throw new Error('Action result does not include image attachments');
}
const content = summarizeActionResult(result);
return {
content: typeof content === 'string' ? content : '',
attachments,
metadata: {},
};
}
function isImageAction(actionId: string) {
return actionId.startsWith('image.filter.');
}
function resolveProjector(actionId: string): ActionResultProjector | null {
if (actionId.startsWith('transcript.audio.')) {
return null;
}
if (isImageAction(actionId)) {
return projectImageResult;
}
switch (actionId) {
case 'mindmap.generate':
case 'slides.outline':
return result => projectTextResult(result);
default:
throw new Error(`No action output projector registered for ${actionId}`);
}
}
export function projectActionResultToAssistantTurn(input: {
session: ChatSession;
actionId: string;
result: unknown;
artifacts?: unknown[];
wasAborted: boolean;
}): Turn | null {
const projector = resolveProjector(input.actionId);
if (!projector) {
return null;
}
const projected = input.wasAborted
? { content: '', attachments: [], metadata: {} }
: projector(input.result, input.artifacts ?? []);
return {
conversationId: input.session.config.sessionId,
role: 'assistant',
content: projected.content,
attachments: projected.attachments,
renderTrace: [],
toolEvents: [],
metadata: projected.metadata,
createdAt: new Date(),
};
}
export type ActionChatEvent = {
type: 'event' | 'attachment' | 'message' | 'error';
id?: string;
data: string | object;
};
export function projectActionEventToChatEvent(
messageId: string | undefined,
data: ActionRuntimeBridgeEvent
): ActionChatEvent {
switch (data.type) {
case 'action_done': {
if (
data.status !== 'succeeded' ||
isImageAction(data.actionId) ||
data.result === undefined
) {
return { type: 'event', id: messageId, data };
}
return {
type: 'message',
id: messageId,
data: textResult(data.result),
};
}
case 'attachment':
return {
type: 'attachment',
id: messageId,
data: data.attachment ?? data,
};
case 'error':
return { type: 'error', id: messageId, data };
default:
return { type: 'event', id: messageId, data };
}
}
@@ -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,
});
}
}
}
@@ -0,0 +1,209 @@
import { Injectable } from '@nestjs/common';
import { CopilotPromptInvalid } from '../../../base';
import { ValidatedStructuredValueSchema } from '../core';
import {
type CopilotChatOptions,
type CopilotEmbeddingOptions,
type CopilotImageOptions,
type CopilotProviderType,
type CopilotRerankRequest,
type CopilotStructuredOptions,
type ModelConditions,
type PromptMessage,
type StreamObject,
} from '../providers/types';
import {
type RequiredStructuredOutputContract,
requireStructuredOutputContract,
} from './contracts';
import {
ExecutionPlanBuilder,
type ExecutionPlanForKind,
} from './execution-plan';
import {
NativeExecutionEngine,
type NativeImageArtifact,
} from './native-execution-engine';
type ProviderFilter = {
prefer?: CopilotProviderType;
};
const providerModelId = (modelId?: string) => modelId ?? 'auto';
@Injectable()
export class CapabilityRuntime {
constructor(
private readonly plans: ExecutionPlanBuilder,
private readonly engine: NativeExecutionEngine
) {}
private async executePlan<TPlan, TResult>(
build: () => Promise<TPlan>,
execute: (plan: TPlan) => Promise<TResult>
) {
return await execute(await build());
}
private executeStreamPlan<TPlan, TChunk>(
build: () => Promise<TPlan>,
execute: (plan: TPlan) => AsyncIterableIterator<TChunk>
): AsyncIterableIterator<TChunk> {
return (async function* () {
yield* execute(await build());
})();
}
private hasNativeDispatch(
plan: ExecutionPlanForKind<'embedding'> | ExecutionPlanForKind<'rerank'>,
kind: 'embedding' | 'rerank'
) {
return !!plan.nativeDispatch?.[kind];
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions,
filter?: ProviderFilter
) {
return await this.executePlan(
() => this.plans.buildTextPlan(cond, messages, options, filter),
plan => this.engine.execute(plan)
);
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions,
filter?: ProviderFilter
): AsyncIterableIterator<string> {
yield* this.executeStreamPlan(
() => this.plans.buildStreamTextPlan(cond, messages, options, filter),
plan => this.engine.executeStream(plan)
);
}
async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions,
filter?: ProviderFilter
): AsyncIterableIterator<StreamObject> {
yield* this.executeStreamPlan(
() => this.plans.buildStreamObjectPlan(cond, messages, options, filter),
plan => this.engine.executeStream(plan)
);
}
async generateStructured(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotStructuredOptions,
filter?: ProviderFilter,
responseContract?: RequiredStructuredOutputContract
) {
return await this.executePlan(
() =>
this.plans.buildStructuredPlan(
cond,
messages,
options,
filter,
responseContract
),
plan => this.engine.execute(plan)
);
}
async generateStructuredValue(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotStructuredOptions,
responseContract?: RequiredStructuredOutputContract,
filter?: ProviderFilter
) {
const validatedResponseContract =
requireStructuredOutputContract(responseContract);
if (!options || !validatedResponseContract) {
throw new CopilotPromptInvalid('Structured schema contract is required');
}
const output = await this.generateStructured(
cond,
messages,
options,
filter,
validatedResponseContract
);
const value = JSON.parse(output);
return ValidatedStructuredValueSchema.parse({
value,
schemaHash: validatedResponseContract.schemaHash,
schemaValidationVersion: 'json-schema-v1',
provider: filter?.prefer ?? 'auto',
model: providerModelId(cond.modelId),
});
}
async embeddingConfigured(modelId: string) {
try {
return this.hasNativeDispatch(
await this.plans.buildEmbeddingPlan(modelId, 'ping'),
'embedding'
);
} catch {
return false;
}
}
async embed(
modelId: string,
input: string | string[],
options?: CopilotEmbeddingOptions
) {
return await this.executePlan(
() => this.plans.buildEmbeddingPlan(modelId, input, options),
plan => this.engine.execute(plan)
);
}
async rerankConfigured(modelId: string) {
try {
return this.hasNativeDispatch(
await this.plans.buildRerankPlan(modelId, {
query: 'ping',
candidates: [{ text: 'ping' }],
}),
'rerank'
);
} catch {
return false;
}
}
async rerank(
modelId: string,
request: CopilotRerankRequest,
options?: CopilotChatOptions
) {
return await this.executePlan(
() => this.plans.buildRerankPlan(modelId, request, options),
plan => this.engine.execute(plan)
);
}
async *streamImageArtifacts(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotImageOptions,
filter?: ProviderFilter
): AsyncIterableIterator<NativeImageArtifact> {
yield* this.executeStreamPlan(
() => this.plans.buildImagePlan(cond, messages, options, filter),
plan => this.engine.executeImageArtifacts(plan)
);
}
}
@@ -0,0 +1,104 @@
import {
type LlmBackendConfig,
llmCompileExecutionPlan,
type LlmEmbeddingRequest,
type LlmImageRequest,
type LlmProtocol,
type LlmRequest,
type LlmRerankRequest,
type LlmStructuredRequest,
} from '../../../../native';
import type {
CopilotProviderType,
ModelConditions,
PromptMessage,
} from '../../providers/types';
// Owner: runtime core mirror facade.
// The semantic source of truth is the native/Rust execution-plan contract
// behind llmCompileExecutionPlan(); this file only keeps the TypeScript shape
// needed by Node live-plan assembly until generated/native TS types replace it.
export type ExecutionRequestKind =
| 'text'
| 'streamText'
| 'streamObject'
| 'structured'
| 'embedding'
| 'rerank'
| 'image';
export type ExecutionRoute = {
providerId: string;
protocol: LlmProtocol;
model: string;
backendConfig: LlmBackendConfig;
};
export type ExecutionTransportContract =
| { kind: 'chat'; request: LlmRequest }
| { kind: 'structured'; request: LlmStructuredRequest }
| { kind: 'embedding'; request: LlmEmbeddingRequest }
| { kind: 'rerank'; request: LlmRerankRequest }
| { kind: 'image'; request: LlmImageRequest };
export type SerializableExecutionPlanRequest =
| {
kind: 'text' | 'streamText' | 'streamObject';
cond: ModelConditions;
messages: PromptMessage[];
options?: Record<string, unknown>;
}
| {
kind: 'structured';
cond: ModelConditions;
messages: PromptMessage[];
options?: Record<string, unknown>;
}
| {
kind: 'image';
cond: ModelConditions;
messages: PromptMessage[];
options?: Record<string, unknown>;
}
| {
kind: 'embedding';
cond: ModelConditions;
modelId: string;
input: string | string[];
options?: Record<string, unknown>;
}
| {
kind: 'rerank';
cond: ModelConditions;
modelId: string;
request: {
query: string;
candidates: { id?: string; text: string }[];
topK?: number;
};
options?: Record<string, unknown>;
};
export type SerializableExecutionPlan = {
routes: ExecutionRoute[];
request: SerializableExecutionPlanRequest;
transport?: ExecutionTransportContract;
routePolicy: { fallbackOrder: string[] };
runtimePolicy: {
prefer?: CopilotProviderType;
maxSteps?: number;
};
attachmentPolicy: {
materializeRemoteAttachments: boolean;
};
responsePostprocess: {
mode: ExecutionRequestKind;
};
hostContext?: {
currentMessages?: PromptMessage[];
};
};
export function parseExecutionPlan(value: unknown) {
return llmCompileExecutionPlan<SerializableExecutionPlan>(value);
}
@@ -0,0 +1,7 @@
export * from './execution-plan-contract';
export * from './native-contract';
export * from './prompt-contract';
export * from './runtime-event-contract';
export * from './shared';
export * from './structured-output-contract';
export * from './tool-contract';
@@ -0,0 +1,97 @@
import serverNativeModule, {
type CapabilityMatchRequest,
type CapabilityMatchResponse,
type ModelRegistryMatchRequest,
type ModelRegistryMatchResponse,
type ModelRegistryResolveRequest,
type ModelRegistryResolveResponse,
type ModelRegistryVariantContract,
type ProviderDriverSpec,
type RequestedModelMatchRequest,
type RequestedModelMatchResponse,
} from '@affine/server-native';
// Owner: native/Rust contract facade.
// These types and validators intentionally proxy @affine/server-native and
// must not grow independent runtime semantics in Node.
export type {
CapabilityMatchRequest,
CapabilityMatchResponse,
ProviderDriverSpec,
RequestedModelMatchRequest,
RequestedModelMatchResponse,
};
export type CopilotModelBackendKind = ModelRegistryMatchRequest['backendKind'];
export type ModelRegistryVariant = ModelRegistryVariantContract;
export type ResolveModelRegistryVariantRequest = ModelRegistryResolveRequest;
export type ResolveModelRegistryVariantResponse = ModelRegistryResolveResponse;
export type MatchModelRegistryRequest = ModelRegistryMatchRequest;
export type MatchModelRegistryResponse = ModelRegistryMatchResponse;
function validateNativeContract<T>(name: string, value: unknown): T {
return serverNativeModule.llmValidateContract(name, value) as T;
}
export function parseCapabilityMatchRequest(value: unknown) {
return validateNativeContract<CapabilityMatchRequest>(
'capabilityMatchRequest',
value
);
}
export function parseCapabilityMatchResponse(value: unknown) {
return validateNativeContract<CapabilityMatchResponse>(
'capabilityMatchResponse',
value
);
}
export function parseResolveModelRegistryVariantRequest(value: unknown) {
return validateNativeContract<ResolveModelRegistryVariantRequest>(
'modelRegistryResolveRequest',
value
);
}
export function parseResolveModelRegistryVariantResponse(value: unknown) {
return validateNativeContract<ResolveModelRegistryVariantResponse>(
'modelRegistryResolveResponse',
value
);
}
export function parseMatchModelRegistryRequest(value: unknown) {
return validateNativeContract<MatchModelRegistryRequest>(
'modelRegistryMatchRequest',
value
);
}
export function parseMatchModelRegistryResponse(value: unknown) {
return validateNativeContract<MatchModelRegistryResponse>(
'modelRegistryMatchResponse',
value
);
}
export function parseProviderDriverSpec(value: unknown) {
return validateNativeContract<ProviderDriverSpec>(
'providerDriverSpec',
value
);
}
export function parseRequestedModelMatchRequest(value: unknown) {
return validateNativeContract<RequestedModelMatchRequest>(
'requestedModelMatchRequest',
value
);
}
export function parseRequestedModelMatchResponse(value: unknown) {
return validateNativeContract<RequestedModelMatchResponse>(
'requestedModelMatchResponse',
value
);
}
@@ -0,0 +1,98 @@
import {
llmValidateContract,
type NativePromptCountTokensRequest,
type NativePromptCountTokensResponse,
type NativePromptMetadataRequest,
type NativePromptMetadataResponse,
type NativePromptRenderRequest,
type NativePromptRenderResponse,
type NativePromptSessionRenderRequest,
type NativePromptSessionRenderResponse,
type PromptMessageContract as NativePromptMessageContract,
type PromptStructuredResponseContract as NativePromptStructuredResponseContract,
} from '../../../../native';
import { normalizePromptResponseFormat } from './structured-output-contract';
// Owner: native/Rust prompt contract facade plus Node responseFormat projection.
// Prompt/message/attachment semantics belong to adapter/native contracts; this
// file keeps only TypeScript aliases and host compatibility projection helpers.
export type PromptStructuredResponseContract =
NativePromptStructuredResponseContract;
export type PromptResponseFormat = {
type: 'json_schema';
responseSchemaJson?: Record<string, unknown>;
schemaHash?: string;
strict?: boolean;
};
export type PromptMessageContract = NativePromptMessageContract;
type PromptMessageInput = {
role: PromptMessageContract['role'];
content: string;
attachments?: unknown[] | null;
params?: Record<string, unknown> | null;
responseFormat?: PromptResponseFormat | null;
};
export type PromptRenderContract = NativePromptRenderRequest;
export type PromptRenderResult = NativePromptRenderResponse;
export type PromptTokenCountContract = NativePromptCountTokensRequest;
export type PromptTokenCountResult = NativePromptCountTokensResponse;
export type PromptMetadataContract = NativePromptMetadataRequest;
export type PromptMetadataResult = NativePromptMetadataResponse;
export type PromptSessionContract = NativePromptSessionRenderRequest;
export type PromptSessionResult = NativePromptSessionRenderResponse;
export type NativePromptResponseFormatProjection = {
nativeResponseFormat?: PromptStructuredResponseContract;
};
export type NativePromptMessageProjection = {
message: PromptMessageContract;
nativeResponseFormat?: PromptStructuredResponseContract;
};
export function projectPromptResponseFormatForNative(
responseFormat?: PromptResponseFormat | null
): NativePromptResponseFormatProjection {
const { nativeResponseFormat } =
normalizePromptResponseFormat(responseFormat);
return {
nativeResponseFormat,
};
}
export function projectPromptMessageForNative(
message: PromptMessageInput
): NativePromptMessageProjection {
const { nativeResponseFormat } = projectPromptResponseFormatForNative(
message.responseFormat
);
const nativeMessage: PromptMessageContract = {
role: message.role,
content: message.content,
...(message.attachments
? {
attachments:
message.attachments as PromptMessageContract['attachments'],
}
: {}),
...(message.params
? { params: message.params as PromptMessageContract['params'] }
: {}),
...(nativeResponseFormat ? { responseFormat: nativeResponseFormat } : {}),
};
return { message: nativeMessage, nativeResponseFormat };
}
export function parsePromptRenderContract(value: unknown) {
return llmValidateContract<PromptRenderContract>(
'promptRenderContract',
value
);
}
export function parsePromptSessionContract(value: unknown) {
return llmValidateContract<PromptSessionContract>(
'promptSessionContract',
value
);
}
@@ -0,0 +1,174 @@
import { z } from 'zod';
import { NonEmptyStringSchema, parseToolLoopStreamEvent } from './shared';
// Owner: app-facing stream projection.
// Native/runtime owns incoming tool-loop event validation; this file owns the
// GraphQL/SSE-facing stream object shape and projection helpers only.
export const TextDeltaStreamObjectSchema = z
.object({
type: z.literal('text-delta'),
textDelta: z.string(),
})
.strict();
export const ReasoningStreamObjectSchema = z
.object({
type: z.literal('reasoning'),
textDelta: z.string(),
})
.strict();
export const ToolCallStreamObjectSchema = z
.object({
type: z.literal('tool-call'),
toolCallId: NonEmptyStringSchema,
toolName: NonEmptyStringSchema,
args: z.record(z.unknown()),
rawArgumentsText: z.string().optional(),
argumentParseError: z.string().optional(),
thought: z.string().optional(),
})
.strict();
export const ToolResultStreamObjectSchema = z
.object({
type: z.literal('tool-result'),
toolCallId: NonEmptyStringSchema,
toolName: NonEmptyStringSchema,
args: z.record(z.unknown()),
result: z.unknown(),
rawArgumentsText: z.string().optional(),
argumentParseError: z.string().optional(),
})
.strict();
export const StreamObjectSchema = z.discriminatedUnion('type', [
TextDeltaStreamObjectSchema,
ReasoningStreamObjectSchema,
ToolCallStreamObjectSchema,
ToolResultStreamObjectSchema,
]);
export type StreamObject = z.infer<typeof StreamObjectSchema>;
export const ToolCallEventSchema = z
.object({
type: z.literal('tool_call'),
toolCallId: NonEmptyStringSchema,
toolName: NonEmptyStringSchema,
args: z.record(z.unknown()),
rawArgumentsText: z.string().optional(),
argumentParseError: z.string().optional(),
thought: z.string().optional(),
})
.strict();
export const ToolResultEventSchema = z
.object({
type: z.literal('tool_result'),
toolCallId: NonEmptyStringSchema,
toolName: NonEmptyStringSchema,
args: z.record(z.unknown()),
result: z.unknown(),
rawArgumentsText: z.string().optional(),
argumentParseError: z.string().optional(),
})
.strict();
export const ToolEventSchema = z.discriminatedUnion('type', [
ToolCallEventSchema,
ToolResultEventSchema,
]);
export type ToolEvent = z.infer<typeof ToolEventSchema>;
export function projectRuntimeEventToStreamObject(
value: unknown
): StreamObject | null {
const event = parseToolLoopStreamEvent(value);
switch (event.type) {
case 'text_delta': {
return { type: 'text-delta', textDelta: event.text };
}
case 'reasoning_delta': {
return { type: 'reasoning', textDelta: event.text };
}
case 'tool_call': {
return {
type: 'tool-call',
toolCallId: event.call_id,
toolName: event.name,
args: event.arguments,
rawArgumentsText: event.arguments_text,
argumentParseError: event.arguments_error,
thought: event.thought,
};
}
case 'tool_result': {
return {
type: 'tool-result',
toolCallId: event.call_id,
toolName: event.name,
args: event.arguments,
result: event.output,
rawArgumentsText: event.arguments_text,
argumentParseError: event.arguments_error,
};
}
default:
return null;
}
}
export function streamObjectToToolEvent(
streamObject: StreamObject
): ToolEvent | undefined {
switch (streamObject.type) {
case 'tool-call':
return {
type: 'tool_call',
toolCallId: streamObject.toolCallId,
toolName: streamObject.toolName,
args: streamObject.args,
rawArgumentsText: streamObject.rawArgumentsText,
argumentParseError: streamObject.argumentParseError,
thought: streamObject.thought,
};
case 'tool-result':
return {
type: 'tool_result',
toolCallId: streamObject.toolCallId,
toolName: streamObject.toolName,
args: streamObject.args,
result: streamObject.result,
rawArgumentsText: streamObject.rawArgumentsText,
argumentParseError: streamObject.argumentParseError,
};
default:
return;
}
}
export function toolEventToStreamObject(event: ToolEvent): StreamObject {
return event.type === 'tool_call'
? {
type: 'tool-call',
toolCallId: event.toolCallId,
toolName: event.toolName,
args: event.args,
rawArgumentsText: event.rawArgumentsText,
argumentParseError: event.argumentParseError,
thought: event.thought,
}
: {
type: 'tool-result',
toolCallId: event.toolCallId,
toolName: event.toolName,
args: event.args,
result: event.result,
rawArgumentsText: event.rawArgumentsText,
argumentParseError: event.argumentParseError,
};
}
@@ -0,0 +1,51 @@
import serverNativeModule from '@affine/server-native';
import { z } from 'zod';
import type { LlmToolLoopStreamEvent } from '../../../../native';
// Owner: Node compatibility helpers.
// JsonValue/NonEmptyString support host Zod schemas; ToolLoopStreamEvent is
// validated by the native/runtime contract via llmValidateContract().
const JsonPrimitiveSchema = z.union([
z.string(),
z.number(),
z.boolean(),
z.null(),
]);
export type JsonValue =
| string
| number
| boolean
| null
| JsonValue[]
| { [key: string]: JsonValue };
export const JsonValueSchema: z.ZodType<JsonValue> = z.lazy(() =>
z.union([
JsonPrimitiveSchema,
z.array(JsonValueSchema),
z.record(JsonValueSchema),
])
);
export const JsonObjectSchema = z.record(JsonValueSchema);
export const NonEmptyStringSchema = z.string().trim().min(1);
export const ToolDefinitionBaseSchema = z
.object({
name: NonEmptyStringSchema,
description: z.string().optional(),
parameters: JsonObjectSchema,
})
.strict();
export type ToolLoopStreamEvent = LlmToolLoopStreamEvent;
export function parseToolLoopStreamEvent(value: unknown) {
return serverNativeModule.llmValidateContract(
'toolLoopEvent',
value
) as LlmToolLoopStreamEvent;
}
@@ -0,0 +1,205 @@
import { llmCanonicalJsonSchemaHash } from '../../../../native';
import { toToolJsonSchema } from '../../tools/json-schema';
import type { JsonValue } from './shared';
// Owner: tool-authoring compatibility plus native schema hash facade.
// Zod-to-JSON-Schema conversion stays in Node for tool authoring, but canonical
// schema hashing is delegated to Rust so JSON key order is not semantic.
export type StructuredOutputValidator = {
parse(input: unknown): unknown;
safeParse(input: unknown): unknown;
};
export type StructuredOutputContract = {
responseSchemaJson?: Record<string, unknown>;
schemaHash?: string;
strict?: boolean;
};
export type RequiredStructuredOutputContract = StructuredOutputContract & {
responseSchemaJson: Record<string, unknown>;
schemaHash: string;
};
type StructuredResponseFormatLike = {
type?: string | null;
responseSchemaJson?: Record<string, unknown>;
schemaHash?: string;
strict?: boolean;
} | null;
type StructuredResponseFormatProjection = {
nativeResponseFormat?: {
type: 'json_schema';
responseSchemaJson: Record<string, JsonValue>;
schemaHash: string;
strict?: boolean;
};
hostResponseFormat?: {
type: 'json_schema';
responseSchemaJson?: Record<string, unknown>;
schemaHash?: string;
strict?: boolean;
};
};
export type StructuredOutputContractFields = Pick<
StructuredOutputContract,
'responseSchemaJson' | 'schemaHash' | 'strict'
>;
export function buildStructuredResponseFromSchemaJson(
responseSchemaJson?: Record<string, unknown>
): StructuredOutputContract {
if (!responseSchemaJson) return {};
const schemaHash = ensurePromptResponseSchemaHash(responseSchemaJson);
return { responseSchemaJson, schemaHash };
}
export function ensurePromptResponseSchemaHash(
schemaJson?: Record<string, unknown>,
schemaHash?: string
) {
if (schemaHash || !schemaJson) return schemaHash;
return llmCanonicalJsonSchemaHash(schemaJson);
}
export function isStructuredOutputValidator(
schema: unknown
): schema is StructuredOutputValidator {
return (
!!schema &&
typeof schema === 'object' &&
'parse' in schema &&
typeof schema.parse === 'function' &&
'safeParse' in schema &&
typeof schema.safeParse === 'function'
);
}
export function buildStructuredResponseContract(
schema?: unknown
): StructuredOutputContract {
if (isStructuredOutputValidator(schema)) {
return buildStructuredResponseFromSchemaJson(
toToolJsonSchema(schema) as Record<string, JsonValue>
);
}
if (schema && typeof schema === 'object' && !Array.isArray(schema)) {
return buildStructuredResponseFromSchemaJson(
schema as Record<string, JsonValue>
);
}
return {};
}
export function buildPromptStructuredResponseFromFields(
fields?: StructuredOutputContractFields | null
): StructuredOutputContract | undefined {
if (!fields) {
return;
}
const { responseSchemaJson, schemaHash, strict } = fields;
const normalizedSchemaHash = ensurePromptResponseSchemaHash(
responseSchemaJson,
schemaHash
);
if (!responseSchemaJson) {
return;
}
return {
...(responseSchemaJson ? { responseSchemaJson } : {}),
...(normalizedSchemaHash ? { schemaHash: normalizedSchemaHash } : {}),
...(strict !== undefined ? { strict } : {}),
};
}
export function buildPromptStructuredResponseContractFromResponseFormat(
responseFormat?: StructuredResponseFormatLike
): StructuredOutputContract | undefined {
if (responseFormat?.type !== 'json_schema') {
return;
}
const responseSchemaJson = responseFormat.responseSchemaJson;
const schemaHash = ensurePromptResponseSchemaHash(
responseSchemaJson,
responseFormat.schemaHash
);
if (!responseSchemaJson) {
return;
}
return {
...(responseSchemaJson ? { responseSchemaJson } : {}),
...(schemaHash ? { schemaHash } : {}),
...(responseFormat.strict !== undefined
? { strict: responseFormat.strict }
: {}),
};
}
export function normalizePromptResponseFormat(
responseFormat?: StructuredResponseFormatLike
): StructuredResponseFormatProjection {
if (responseFormat?.type !== 'json_schema') {
return {};
}
const contract =
buildPromptStructuredResponseContractFromResponseFormat(responseFormat);
if (!contract) {
return {};
}
const nextNativeResponseFormat =
contract.responseSchemaJson && contract.schemaHash
? {
type: 'json_schema' as const,
responseSchemaJson: contract.responseSchemaJson as Record<
string,
JsonValue
>,
schemaHash: contract.schemaHash,
...(responseFormat.strict !== undefined
? { strict: responseFormat.strict }
: {}),
}
: undefined;
const nextHostResponseFormat = {
type: 'json_schema' as const,
...(contract.responseSchemaJson
? { responseSchemaJson: contract.responseSchemaJson }
: {}),
...(contract.schemaHash ? { schemaHash: contract.schemaHash } : {}),
...(responseFormat.strict !== undefined
? { strict: responseFormat.strict }
: {}),
};
return {
nativeResponseFormat: nextNativeResponseFormat,
hostResponseFormat: contract.responseSchemaJson
? nextHostResponseFormat
: undefined,
};
}
export function requireStructuredOutputContract(
contract?: StructuredOutputContract
): RequiredStructuredOutputContract | undefined {
if (!contract?.responseSchemaJson || !contract.schemaHash) {
return;
}
return {
responseSchemaJson: contract.responseSchemaJson,
schemaHash: contract.schemaHash,
...(contract.strict !== undefined ? { strict: contract.strict } : {}),
};
}
@@ -0,0 +1,37 @@
import { z } from 'zod';
import type { CopilotToolSet } from '../../tools';
import { ensureToolJsonSchema } from '../../tools/tool';
import { ToolDefinitionBaseSchema } from './shared';
// Owner: tool authoring facade over runtime-owned callback contracts.
// Tool definitions are built from Node-hosted tools; callback request/response
// values are still validated against the runtime/native schema at the boundary.
export const ToolContractSchema = ToolDefinitionBaseSchema;
export type ToolContract = z.infer<typeof ToolContractSchema>;
export interface ToolCallRequest {
callId: string;
name: string;
args: Record<string, unknown>;
rawArgumentsText?: string;
argumentParseError?: string;
}
export interface ToolCallResult extends ToolCallRequest {
output: unknown;
isError?: boolean;
}
export function parseToolContract(value: unknown) {
return ToolContractSchema.parse(value);
}
export function buildToolContracts(toolSet: CopilotToolSet): ToolContract[] {
return Object.entries(toolSet).map(([name, tool]) =>
parseToolContract({
name,
description: tool.description,
parameters: ensureToolJsonSchema(tool, name),
})
);
}
@@ -0,0 +1,65 @@
import { Injectable } from '@nestjs/common';
import { metrics } from '../../../base';
import type { ResolvedCopilotProvider } from '../providers/factory';
import type { CopilotProviderType } from '../providers/types';
import type { ExecutionRequestKind } from './execution-plan';
type ExecutionDispatchPath = 'prepared_routes';
export function summarizePreparedRoutes(
routes: Array<Pick<ResolvedCopilotProvider, 'prepared'>>
) {
const preparedCount = routes.filter(route => !!route.prepared).length;
return {
routeCount: routes.length,
preparedCount,
preparedMode:
preparedCount === 0
? 'none'
: preparedCount === routes.length
? 'all'
: 'partial',
} as const;
}
function planAttrs(
kind: ExecutionRequestKind,
prefer?: CopilotProviderType,
routes?: ResolvedCopilotProvider[]
) {
const summary = summarizePreparedRoutes(routes ?? []);
return {
kind,
prefer: prefer ?? 'auto',
prepared: summary.preparedMode,
route_count: summary.routeCount,
};
}
@Injectable()
export class CopilotExecutionMetrics {
recordPlan(
kind: ExecutionRequestKind,
routes: ResolvedCopilotProvider[],
prefer?: CopilotProviderType
) {
const attrs = planAttrs(kind, prefer, routes);
metrics.ai.counter('execution_plan_total').add(1, attrs);
metrics.ai.histogram('execution_plan_routes').record(attrs.route_count, {
kind: attrs.kind,
prefer: attrs.prefer,
prepared: attrs.prepared,
});
}
recordDispatch(
kind: ExecutionRequestKind,
path: ExecutionDispatchPath,
routeCount: number
) {
const attrs = { kind, path };
metrics.ai.counter('execution_dispatch_total').add(1, attrs);
metrics.ai.histogram('execution_dispatch_routes').record(routeCount, attrs);
}
}

Some files were not shown because too many files have changed in this diff Show More