feat: refactor copilot module (#14537)

This commit is contained in:
DarkSky
2026-03-02 13:57:55 +08:00
committed by GitHub
parent 60acd81d4b
commit c5d622531c
92 changed files with 5759 additions and 2170 deletions
@@ -43,7 +43,9 @@ Generated by [AVA](https://avajs.dev).
> Snapshot 5
Buffer @Uint8Array [
66616b65 20696d61 6765
89504e47 0d0a1a0a 0000000d 49484452 00000001 00000001 08040000 00b51c0c
02000000 0b494441 5478da63 fcff1f00 03030200 efa37c9f 00000000 49454e44
ae426082
]
## should preview link
@@ -12,12 +12,12 @@ Generated by [AVA](https://avajs.dev).
{
messages: [
{
content: 'generate text to text',
content: 'generate text to text stream',
role: 'assistant',
},
],
pinned: false,
tokens: 8,
tokens: 10,
},
]
@@ -27,12 +27,12 @@ Generated by [AVA](https://avajs.dev).
{
messages: [
{
content: 'generate text to text',
content: 'generate text to text stream',
role: 'assistant',
},
],
pinned: false,
tokens: 8,
tokens: 10,
},
]
@@ -4,31 +4,31 @@ import type { ExecutionContext, TestFn } from 'ava';
import ava from 'ava';
import { z } from 'zod';
import { ServerFeature, ServerService } from '../core';
import { AuthService } from '../core/auth';
import { QuotaModule } from '../core/quota';
import { Models } from '../models';
import { CopilotModule } from '../plugins/copilot';
import { prompts, PromptService } from '../plugins/copilot/prompt';
import { ServerFeature, ServerService } from '../../core';
import { AuthService } from '../../core/auth';
import { QuotaModule } from '../../core/quota';
import { Models } from '../../models';
import { CopilotModule } from '../../plugins/copilot';
import { prompts, PromptService } from '../../plugins/copilot/prompt';
import {
CopilotProviderFactory,
CopilotProviderType,
StreamObject,
StreamObjectSchema,
} from '../plugins/copilot/providers';
import { TranscriptionResponseSchema } from '../plugins/copilot/transcript/types';
} from '../../plugins/copilot/providers';
import { TranscriptionResponseSchema } from '../../plugins/copilot/transcript/types';
import {
CopilotChatTextExecutor,
CopilotWorkflowService,
GraphExecutorState,
} from '../plugins/copilot/workflow';
} from '../../plugins/copilot/workflow';
import {
CopilotChatImageExecutor,
CopilotCheckHtmlExecutor,
CopilotCheckJsonExecutor,
} from '../plugins/copilot/workflow/executor';
import { createTestingModule, TestingModule } from './utils';
import { TestAssets } from './utils/copilot';
} from '../../plugins/copilot/workflow/executor';
import { createTestingModule, TestingModule } from '../utils';
import { TestAssets } from '../utils/copilot';
type Tester = {
auth: AuthService;
@@ -6,25 +6,25 @@ import type { TestFn } from 'ava';
import ava from 'ava';
import Sinon from 'sinon';
import { AppModule } from '../app.module';
import { JobQueue } from '../base';
import { ConfigModule } from '../base/config';
import { AuthService } from '../core/auth';
import { DocReader } from '../core/doc';
import { CopilotContextService } from '../plugins/copilot/context';
import { AppModule } from '../../app.module';
import { JobQueue } from '../../base';
import { ConfigModule } from '../../base/config';
import { AuthService } from '../../core/auth';
import { DocReader } from '../../core/doc';
import { CopilotContextService } from '../../plugins/copilot/context';
import {
CopilotEmbeddingJob,
MockEmbeddingClient,
} from '../plugins/copilot/embedding';
import { prompts, PromptService } from '../plugins/copilot/prompt';
} from '../../plugins/copilot/embedding';
import { prompts, PromptService } from '../../plugins/copilot/prompt';
import {
CopilotProviderFactory,
CopilotProviderType,
GeminiGenerativeProvider,
OpenAIProvider,
} from '../plugins/copilot/providers';
import { CopilotStorage } from '../plugins/copilot/storage';
import { MockCopilotProvider } from './mocks';
} from '../../plugins/copilot/providers';
import { CopilotStorage } from '../../plugins/copilot/storage';
import { MockCopilotProvider } from '../mocks';
import {
acceptInviteById,
createTestingApp,
@@ -33,7 +33,7 @@ import {
smallestPng,
TestingApp,
TestUser,
} from './utils';
} from '../utils';
import {
addContextDoc,
addContextFile,
@@ -67,7 +67,7 @@ import {
textToEventStream,
unsplashSearch,
updateCopilotSession,
} from './utils/copilot';
} from '../utils/copilot';
const test = ava as TestFn<{
auth: AuthService;
@@ -513,7 +513,11 @@ test('should be able to chat with api', async t => {
);
const messageId = await createCopilotMessage(app, sessionId);
const ret = await chatWithText(app, sessionId, messageId);
t.is(ret, 'generate text to text', 'should be able to chat with text');
t.is(
ret,
'generate text to text stream',
'should be able to chat with text'
);
const ret2 = await chatWithTextStream(app, sessionId, messageId);
t.is(
@@ -657,7 +661,7 @@ test('should be able to retry with api', async t => {
const histories = await getHistories(app, { workspaceId: id, docId });
t.deepEqual(
histories.map(h => h.messages.map(m => m.content)),
[['generate text to text', 'generate text to text']],
[['generate text to text stream', 'generate text to text stream']],
'should be able to list history'
);
}
@@ -794,7 +798,7 @@ test('should be able to list history', async t => {
const histories = await getHistories(app, { workspaceId, docId });
t.deepEqual(
histories.map(h => h.messages.map(m => m.content)),
[['hello', 'generate text to text']],
[['hello', 'generate text to text stream']],
'should be able to list history'
);
}
@@ -807,7 +811,7 @@ test('should be able to list history', async t => {
});
t.deepEqual(
histories.map(h => h.messages.map(m => m.content)),
[['generate text to text', 'hello']],
[['generate text to text stream', 'hello']],
'should be able to list history'
);
}
@@ -858,7 +862,7 @@ test('should reject request that user have not permission', async t => {
const histories = await getHistories(app, { workspaceId, docId });
t.deepEqual(
histories.map(h => h.messages.map(m => m.content)),
[['generate text to text']],
[['generate text to text stream']],
'should able to list history'
);
@@ -8,38 +8,38 @@ import ava from 'ava';
import { nanoid } from 'nanoid';
import Sinon from 'sinon';
import { EventBus, JobQueue } from '../base';
import { ConfigModule } from '../base/config';
import { AuthService } from '../core/auth';
import { QuotaModule } from '../core/quota';
import { StorageModule, WorkspaceBlobStorage } from '../core/storage';
import { EventBus, JobQueue } from '../../base';
import { ConfigModule } from '../../base/config';
import { AuthService } from '../../core/auth';
import { QuotaModule } from '../../core/quota';
import { StorageModule, WorkspaceBlobStorage } from '../../core/storage';
import {
ContextCategories,
CopilotSessionModel,
WorkspaceModel,
} from '../models';
import { CopilotModule } from '../plugins/copilot';
import { CopilotContextService } from '../plugins/copilot/context';
import { CopilotCronJobs } from '../plugins/copilot/cron';
} from '../../models';
import { CopilotModule } from '../../plugins/copilot';
import { CopilotContextService } from '../../plugins/copilot/context';
import { CopilotCronJobs } from '../../plugins/copilot/cron';
import {
CopilotEmbeddingJob,
MockEmbeddingClient,
} from '../plugins/copilot/embedding';
import { prompts, PromptService } from '../plugins/copilot/prompt';
} from '../../plugins/copilot/embedding';
import { prompts, PromptService } from '../../plugins/copilot/prompt';
import {
CopilotProviderFactory,
CopilotProviderType,
ModelInputType,
ModelOutputType,
OpenAIProvider,
} from '../plugins/copilot/providers';
} from '../../plugins/copilot/providers';
import {
CitationParser,
TextStreamParser,
} from '../plugins/copilot/providers/utils';
import { ChatSessionService } from '../plugins/copilot/session';
import { CopilotStorage } from '../plugins/copilot/storage';
import { CopilotTranscriptionService } from '../plugins/copilot/transcript';
} from '../../plugins/copilot/providers/utils';
import { ChatSessionService } from '../../plugins/copilot/session';
import { CopilotStorage } from '../../plugins/copilot/storage';
import { CopilotTranscriptionService } from '../../plugins/copilot/transcript';
import {
CopilotChatTextExecutor,
CopilotWorkflowService,
@@ -48,7 +48,7 @@ import {
WorkflowGraphExecutor,
type WorkflowNodeData,
WorkflowNodeType,
} from '../plugins/copilot/workflow';
} from '../../plugins/copilot/workflow';
import {
CopilotChatImageExecutor,
CopilotCheckHtmlExecutor,
@@ -56,16 +56,16 @@ import {
getWorkflowExecutor,
NodeExecuteState,
NodeExecutorType,
} from '../plugins/copilot/workflow/executor';
import { AutoRegisteredWorkflowExecutor } from '../plugins/copilot/workflow/executor/utils';
import { WorkflowGraphList } from '../plugins/copilot/workflow/graph';
import { CopilotWorkspaceService } from '../plugins/copilot/workspace';
import { PaymentModule } from '../plugins/payment';
import { SubscriptionService } from '../plugins/payment/service';
import { SubscriptionStatus } from '../plugins/payment/types';
import { MockCopilotProvider } from './mocks';
import { createTestingModule, TestingModule } from './utils';
import { WorkflowTestCases } from './utils/copilot';
} from '../../plugins/copilot/workflow/executor';
import { AutoRegisteredWorkflowExecutor } from '../../plugins/copilot/workflow/executor/utils';
import { WorkflowGraphList } from '../../plugins/copilot/workflow/graph';
import { CopilotWorkspaceService } from '../../plugins/copilot/workspace';
import { PaymentModule } from '../../plugins/payment';
import { SubscriptionService } from '../../plugins/payment/service';
import { SubscriptionStatus } from '../../plugins/payment/types';
import { MockCopilotProvider } from '../mocks';
import { createTestingModule, TestingModule } from '../utils';
import { WorkflowTestCases } from '../utils/copilot';
type Context = {
auth: AuthService;
@@ -364,6 +364,21 @@ test('should be able to manage chat session', async t => {
});
t.is(newSessionId, sessionId, 'should get same session id');
}
// should create a fresh session when reuseLatestChat is explicitly disabled
{
const newSessionId = await session.create({
userId,
promptName,
...commonParams,
reuseLatestChat: false,
});
t.not(
newSessionId,
sessionId,
'should create new session id when reuseLatestChat is false'
);
}
});
test('should be able to update chat session prompt', async t => {
@@ -881,6 +896,26 @@ test('should be able to get provider', async t => {
}
});
test('should resolve provider by prefixed model id', async t => {
const { factory } = t.context;
const provider = await factory.getProviderByModel('openai-default/test');
t.truthy(provider, 'should resolve prefixed model id');
t.is(provider?.type, CopilotProviderType.OpenAI);
const result = await provider?.text({ modelId: 'openai-default/test' }, [
{ role: 'user', content: 'hello' },
]);
t.is(result, 'generate text to text');
});
test('should fallback to null when prefixed provider id does not exist', async t => {
const { factory } = t.context;
const provider = await factory.getProviderByModel('unknown/test');
t.is(provider, null);
});
// ==================== workflow ====================
// this test used to preview the final result of the workflow
@@ -2063,25 +2098,23 @@ test('should handle copilot cron jobs correctly', async t => {
});
test('should resolve model correctly based on subscription status and prompt config', async t => {
const { db, session, subscription } = t.context;
const { prompt, session, subscription } = t.context;
// 1) Seed a prompt that has optionalModels and proModels in config
const promptName = 'resolve-model-test';
await db.aiPrompt.create({
data: {
name: promptName,
model: 'gemini-2.5-flash',
messages: {
create: [{ idx: 0, role: 'system', content: 'test' }],
},
config: { proModels: ['gemini-2.5-pro', 'claude-sonnet-4-5@20250929'] },
await prompt.set(
promptName,
'gemini-2.5-flash',
[{ role: 'system', content: 'test' }],
{ proModels: ['gemini-2.5-pro', 'claude-sonnet-4-5@20250929'] },
{
optionalModels: [
'gemini-2.5-flash',
'gemini-2.5-pro',
'claude-sonnet-4-5@20250929',
],
},
});
}
);
// 2) Create a chat session with this prompt
const sessionId = await session.create({
@@ -2106,6 +2139,16 @@ test('should resolve model correctly based on subscription status and prompt con
const model1 = await s.resolveModel(false, 'gemini-2.5-pro');
t.snapshot(model1, 'should honor requested pro model');
const model1WithPrefix = await s.resolveModel(
false,
'openai-default/gemini-2.5-pro'
);
t.is(
model1WithPrefix,
'openai-default/gemini-2.5-pro',
'should honor requested prefixed pro model'
);
const model2 = await s.resolveModel(false, 'not-in-optional');
t.snapshot(model2, 'should fallback to default model');
}
@@ -2119,6 +2162,16 @@ test('should resolve model correctly based on subscription status and prompt con
'should fallback to default model when requesting pro model during trialing'
);
const model3WithPrefix = await s.resolveModel(
true,
'openai-default/gemini-2.5-pro'
);
t.is(
model3WithPrefix,
'gemini-2.5-flash',
'should fallback to default model when requesting prefixed pro model during trialing'
);
const model4 = await s.resolveModel(true, 'gemini-2.5-flash');
t.snapshot(model4, 'should honor requested non-pro model during trialing');
@@ -2141,6 +2194,16 @@ test('should resolve model correctly based on subscription status and prompt con
const model7 = await s.resolveModel(true, 'claude-sonnet-4-5@20250929');
t.snapshot(model7, 'should honor requested pro model during active');
const model7WithPrefix = await s.resolveModel(
true,
'openai-default/claude-sonnet-4-5@20250929'
);
t.is(
model7WithPrefix,
'openai-default/claude-sonnet-4-5@20250929',
'should honor requested prefixed pro model during active'
);
const model8 = await s.resolveModel(true, 'not-in-optional');
t.snapshot(
model8,
@@ -0,0 +1,210 @@
import test from 'ava';
import { z } from 'zod';
import type { NativeLlmRequest, NativeLlmStreamEvent } from '../../native';
import {
buildNativeRequest,
NativeProviderAdapter,
} from '../../plugins/copilot/providers/native';
const mockDispatch = () =>
(async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
yield { type: 'text_delta', text: 'Use [^1] now' };
yield { type: 'citation', index: 1, url: 'https://affine.pro' };
yield { type: 'done', finish_reason: 'stop' };
})();
test('NativeProviderAdapter streamText should append citation footnotes', async t => {
const adapter = new NativeProviderAdapter(mockDispatch, {}, 3);
const chunks: string[] = [];
for await (const chunk of adapter.streamText({
model: 'gpt-4.1',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
})) {
chunks.push(chunk);
}
const text = chunks.join('');
t.true(text.includes('Use [^1] now'));
t.true(
text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}')
);
});
test('NativeProviderAdapter streamObject should append citation footnotes', async t => {
const adapter = new NativeProviderAdapter(mockDispatch, {}, 3);
const chunks = [];
for await (const chunk of adapter.streamObject({
model: 'gpt-4.1',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
})) {
chunks.push(chunk);
}
t.deepEqual(
chunks.map(chunk => chunk.type),
['text-delta', 'text-delta']
);
const text = chunks
.filter(chunk => chunk.type === 'text-delta')
.map(chunk => chunk.textDelta)
.join('');
t.true(text.includes('Use [^1] now'));
t.true(
text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}')
);
});
test('NativeProviderAdapter streamObject should append fallback attachment footnotes', async t => {
const dispatch = () =>
(async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
yield {
type: 'tool_result',
call_id: 'call_1',
name: 'blob_read',
arguments: { blob_id: 'blob_1' },
output: {
blobId: 'blob_1',
fileName: 'a.txt',
fileType: 'text/plain',
content: 'A',
},
};
yield {
type: 'tool_result',
call_id: 'call_2',
name: 'blob_read',
arguments: { blob_id: 'blob_2' },
output: {
blobId: 'blob_2',
fileName: 'b.txt',
fileType: 'text/plain',
content: 'B',
},
};
yield { type: 'text_delta', text: 'Answer from files.' };
yield { type: 'done', finish_reason: 'stop' };
})();
const adapter = new NativeProviderAdapter(dispatch, {}, 3);
const chunks = [];
for await (const chunk of adapter.streamObject({
model: 'gpt-4.1',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
})) {
chunks.push(chunk);
}
const text = chunks
.filter(chunk => chunk.type === 'text-delta')
.map(chunk => chunk.textDelta)
.join('');
t.true(text.includes('Answer from files.'));
t.true(text.includes('[^1][^2]'));
t.true(
text.includes(
'[^1]: {"type":"attachment","blobId":"blob_1","fileName":"a.txt","fileType":"text/plain"}'
)
);
t.true(
text.includes(
'[^2]: {"type":"attachment","blobId":"blob_2","fileName":"b.txt","fileType":"text/plain"}'
)
);
});
test('NativeProviderAdapter streamObject should map tool and text events', async t => {
let round = 0;
const dispatch = (_request: NativeLlmRequest) =>
(async function* (): AsyncIterableIterator<NativeLlmStreamEvent> {
round += 1;
if (round === 1) {
yield {
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: 'ok' };
yield { type: 'done', finish_reason: 'stop' };
})();
const adapter = new NativeProviderAdapter(
dispatch,
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
execute: async () => ({ markdown: '# a1' }),
},
},
4
);
const events = [];
for await (const event of adapter.streamObject({
model: 'gpt-4.1',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'read' }] }],
})) {
events.push(event);
}
t.deepEqual(
events.map(event => event.type),
['tool-call', 'tool-result', 'text-delta']
);
t.deepEqual(events[0], {
type: 'tool-call',
toolCallId: 'call_1',
toolName: 'doc_read',
args: { doc_id: 'a1' },
});
});
test('buildNativeRequest should include rust middleware from profile', async t => {
const { request } = await buildNativeRequest({
model: 'gpt-4.1',
messages: [{ role: 'user', content: 'hello' }],
tools: {},
middleware: {
rust: {
request: ['normalize_messages', 'clamp_max_tokens'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['callout'],
},
},
});
t.deepEqual(request.middleware, {
request: ['normalize_messages', 'clamp_max_tokens'],
stream: ['stream_event_normalize', 'citation_indexing'],
});
});
test('NativeProviderAdapter streamText should skip citation footnotes when disabled', async t => {
const adapter = new NativeProviderAdapter(mockDispatch, {}, 3, {
nodeTextMiddleware: ['callout'],
});
const chunks: string[] = [];
for await (const chunk of adapter.streamText({
model: 'gpt-4.1',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
})) {
chunks.push(chunk);
}
const text = chunks.join('');
t.true(text.includes('Use [^1] now'));
t.false(
text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}')
);
});
@@ -0,0 +1,56 @@
import test from 'ava';
import { resolveProviderMiddleware } from '../../plugins/copilot/providers/provider-middleware';
import { buildProviderRegistry } from '../../plugins/copilot/providers/provider-registry';
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.deepEqual(middleware.node?.text, ['citation_footnote', 'callout']);
});
test('resolveProviderMiddleware should merge defaults and overrides', t => {
const middleware = resolveProviderMiddleware(CopilotProviderType.OpenAI, {
rust: { request: ['clamp_max_tokens'] },
node: { text: ['thinking_format'] },
});
t.deepEqual(middleware.rust?.request, [
'normalize_messages',
'clamp_max_tokens',
]);
t.deepEqual(middleware.node?.text, [
'citation_footnote',
'callout',
'thinking_format',
]);
});
test('buildProviderRegistry should normalize profile middleware defaults', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
],
});
const profile = registry.profiles.get('openai-main');
t.truthy(profile);
t.deepEqual(profile?.middleware.rust?.stream, [
'stream_event_normalize',
'citation_indexing',
]);
t.deepEqual(profile?.middleware.node?.text, ['citation_footnote', 'callout']);
});
@@ -0,0 +1,99 @@
import test from 'ava';
import { ProviderMiddlewareConfig } from '../../plugins/copilot/config';
import { CopilotProvider } from '../../plugins/copilot/providers/provider';
import {
CopilotProviderType,
ModelInputType,
ModelOutputType,
} from '../../plugins/copilot/providers/types';
class TestOpenAIProvider extends CopilotProvider<{ apiKey: string }> {
readonly type = CopilotProviderType.OpenAI;
readonly models = [
{
id: 'gpt-4.1',
capabilities: [
{
input: [ModelInputType.Text],
output: [ModelOutputType.Text],
defaultForOutputType: true,
},
],
},
];
configured() {
return true;
}
async text(_cond: any, _messages: any[], _options?: any) {
return '';
}
async *streamText(_cond: any, _messages: any[], _options?: any) {
yield '';
}
exposeMetricLabels() {
return this.metricLabels('gpt-4.1');
}
exposeMiddleware() {
return this.getActiveProviderMiddleware();
}
}
function createProvider(profileMiddleware?: ProviderMiddlewareConfig) {
const provider = new TestOpenAIProvider();
(provider as any).AFFiNEConfig = {
copilot: {
providers: {
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: 'test' },
middleware: profileMiddleware,
},
],
defaults: {},
openai: { apiKey: 'legacy' },
},
},
};
return provider;
}
test('metricLabels should include active provider id', t => {
const provider = createProvider();
const labels = provider.runWithProfile('openai-main', () =>
provider.exposeMetricLabels()
);
t.is(labels.providerId, 'openai-main');
});
test('getActiveProviderMiddleware should merge defaults with profile override', t => {
const provider = createProvider({
rust: { request: ['clamp_max_tokens'] },
node: { text: ['thinking_format'] },
});
const middleware = provider.runWithProfile('openai-main', () =>
provider.exposeMiddleware()
);
t.deepEqual(middleware.rust?.request, [
'normalize_messages',
'clamp_max_tokens',
]);
t.deepEqual(middleware.rust?.stream, [
'stream_event_normalize',
'citation_indexing',
]);
t.deepEqual(middleware.node?.text, [
'citation_footnote',
'callout',
'thinking_format',
]);
});
@@ -0,0 +1,165 @@
import test from 'ava';
import {
buildProviderRegistry,
resolveModel,
stripProviderPrefix,
} from '../../plugins/copilot/providers/provider-registry';
import {
CopilotProviderType,
ModelOutputType,
} from '../../plugins/copilot/providers/types';
test('buildProviderRegistry should keep explicit profile over legacy compatibility profile', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-default',
type: CopilotProviderType.OpenAI,
priority: 100,
config: { apiKey: 'new' },
},
],
openai: { apiKey: 'legacy' },
});
const profile = registry.profiles.get('openai-default');
t.truthy(profile);
t.deepEqual(profile?.config, { apiKey: 'new' });
});
test('buildProviderRegistry should reject duplicated profile ids', t => {
const error = t.throws(() =>
buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '2' },
},
],
})
) as Error;
t.truthy(error);
t.regex(error.message, /Duplicated copilot provider profile id/);
});
test('buildProviderRegistry should reject defaults that reference unknown providers', t => {
const error = t.throws(() =>
buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
],
defaults: {
fallback: 'unknown-provider',
},
})
) as Error;
t.truthy(error);
t.regex(error.message, /defaults references unknown providerId/);
});
test('resolveModel should support explicit provider prefix and keep slash models untouched', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
{
id: 'fal-main',
type: CopilotProviderType.FAL,
config: { apiKey: '2' },
},
],
});
const prefixed = resolveModel({
registry,
modelId: 'openai-main/gpt-4.1',
});
t.deepEqual(prefixed, {
rawModelId: 'openai-main/gpt-4.1',
modelId: 'gpt-4.1',
explicitProviderId: 'openai-main',
candidateProviderIds: ['openai-main'],
});
const slashModel = resolveModel({
registry,
modelId: 'lora/image-to-image',
});
t.is(slashModel.modelId, 'lora/image-to-image');
t.false(slashModel.candidateProviderIds.includes('lora'));
});
test('resolveModel should follow defaults -> fallback -> order and apply filters', 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: 'openai-main',
},
});
const routed = resolveModel({
registry,
outputType: ModelOutputType.Text,
preferredProviderIds: ['openai-main', 'fal-main'],
});
t.deepEqual(routed.candidateProviderIds, ['openai-main', 'fal-main']);
});
test('stripProviderPrefix should only strip matched provider prefix', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
],
});
t.is(
stripProviderPrefix(registry, 'openai-main', 'openai-main/gpt-4.1'),
'gpt-4.1'
);
t.is(
stripProviderPrefix(registry, 'openai-main', 'another-main/gpt-4.1'),
'another-main/gpt-4.1'
);
t.is(stripProviderPrefix(registry, 'openai-main', 'gpt-4.1'), 'gpt-4.1');
});
@@ -0,0 +1,134 @@
import test from 'ava';
import { z } from 'zod';
import { NativeLlmRequest, NativeLlmStreamEvent } from '../../native';
import {
ToolCallAccumulator,
ToolCallLoop,
ToolSchemaExtractor,
} from '../../plugins/copilot/providers/loop';
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"}',
});
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' },
thought: undefined,
});
});
test('ToolSchemaExtractor should convert zod schema to json schema', t => {
const toolSet = {
doc_read: {
description: 'Read doc',
inputSchema: z.object({
doc_id: z.string(),
limit: z.number().optional(),
}),
execute: async () => ({}),
},
};
const extracted = ToolSchemaExtractor.extract(toolSet);
t.deepEqual(extracted, [
{
name: 'doc_read',
description: 'Read doc',
parameters: {
type: 'object',
properties: {
doc_id: { type: 'string' },
limit: { type: 'number' },
},
additionalProperties: false,
required: ['doc_id'],
},
},
]);
});
test('ToolCallLoop should execute tool call and continue to next round', async t => {
const dispatchRequests: NativeLlmRequest[] = [];
const dispatch = (request: NativeLlmRequest) => {
dispatchRequests.push(request);
const round = dispatchRequests.length;
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 {
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' };
})();
};
let executedArgs: Record<string, unknown> | null = null;
const loop = new ToolCallLoop(
dispatch,
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
execute: async args => {
executedArgs = args;
return { markdown: '# doc' };
},
},
},
4
);
const events: NativeLlmStreamEvent[] = [];
for await (const event of loop.run({
model: 'gpt-4.1',
stream: true,
messages: [{ role: 'user', content: [{ type: 'text', text: 'read doc' }] }],
})) {
events.push(event);
}
t.deepEqual(executedArgs, { doc_id: 'a1' });
t.true(
dispatchRequests[1]?.messages.some(message => message.role === 'tool')
);
t.deepEqual(
events.map(event => event.type),
['tool_call', 'tool_result', 'text_delta', 'done']
);
});
@@ -0,0 +1,116 @@
import test from 'ava';
import { z } from 'zod';
import {
chatToGPTMessage,
CitationFootnoteFormatter,
CitationParser,
StreamPatternParser,
} from '../../plugins/copilot/providers/utils';
test('CitationFootnoteFormatter should format sorted footnotes from citation events', t => {
const formatter = new CitationFootnoteFormatter();
formatter.consume({
type: 'citation',
index: 2,
url: 'https://example.com/b',
});
formatter.consume({
type: 'citation',
index: 1,
url: 'https://example.com/a',
});
t.is(
formatter.end(),
[
'[^1]: {"type":"url","url":"https%3A%2F%2Fexample.com%2Fa"}',
'[^2]: {"type":"url","url":"https%3A%2F%2Fexample.com%2Fb"}',
].join('\n')
);
});
test('CitationFootnoteFormatter should overwrite duplicated index with latest url', t => {
const formatter = new CitationFootnoteFormatter();
formatter.consume({
type: 'citation',
index: 1,
url: 'https://example.com/old',
});
formatter.consume({
type: 'citation',
index: 1,
url: 'https://example.com/new',
});
t.is(
formatter.end(),
'[^1]: {"type":"url","url":"https%3A%2F%2Fexample.com%2Fnew"}'
);
});
test('StreamPatternParser should keep state across chunks', t => {
const parser = new StreamPatternParser(pattern => {
if (pattern.kind === 'wrappedLink') {
return `[^${pattern.url}]`;
}
if (pattern.kind === 'index') {
return `[#${pattern.value}]`;
}
return `[${pattern.text}](${pattern.url})`;
});
const first = parser.write('ref ([AFFiNE](https://affine.pro');
const second = parser.write(')) and [2]');
t.is(first, 'ref ');
t.is(second, '[^https://affine.pro] and [#2]');
t.is(parser.end(), '');
});
test('CitationParser should convert wrapped links to numbered footnotes', t => {
const parser = new CitationParser();
const output = parser.parse('Use ([AFFiNE](https://affine.pro)) now');
t.is(output, 'Use [^1] now');
t.regex(
parser.end(),
/\[\^1\]: \{"type":"url","url":"https%3A%2F%2Faffine.pro"\}/
);
});
test('chatToGPTMessage should not mutate input and should keep system schema', async t => {
const schema = z.object({
query: z.string(),
});
const messages = [
{
role: 'system' as const,
content: 'You are helper',
params: { schema },
},
{
role: 'user' as const,
content: '',
attachments: ['https://example.com/a.png'],
},
];
const firstRef = messages[0];
const secondRef = messages[1];
const [system, normalized, parsedSchema] = await chatToGPTMessage(
messages,
false
);
t.is(system, 'You are helper');
t.is(parsedSchema, schema);
t.is(messages.length, 2);
t.is(messages[0], firstRef);
t.is(messages[1], secondRef);
t.deepEqual(normalized[0], {
role: 'user',
content: [{ type: 'text', text: '[no content]' }],
});
});
@@ -0,0 +1,82 @@
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);
});
@@ -629,14 +629,35 @@ export async function chatWithText(
prefix = '',
retry?: boolean
): Promise<string> {
const endpoint = prefix || '/stream';
const query = messageId
? `?messageId=${messageId}` + (retry ? '&retry=true' : '')
: '';
const res = await app
.GET(`/api/copilot/chat/${sessionId}${prefix}${query}`)
.GET(`/api/copilot/chat/${sessionId}${endpoint}${query}`)
.expect(200);
return res.text;
if (prefix) {
return res.text;
}
const events = sse2array(res.text);
const errorEvent = events.find(event => event.event === 'error');
if (errorEvent?.data) {
let message = errorEvent.data;
try {
const parsed = JSON.parse(errorEvent.data);
message = parsed.message || message;
} catch {
// noop: keep raw error data
}
throw new Error(message);
}
return events
.filter(event => event.event === 'message')
.map(event => event.data ?? '')
.join('');
}
export async function chatWithTextStream(
@@ -38,8 +38,11 @@ test.before(async t => {
t.context.app = app;
});
test.after.always(async t => {
test.afterEach.always(() => {
Sinon.restore();
});
test.after.always(async t => {
__resetDnsLookupForTests();
await t.context.app.close();
});
@@ -80,6 +83,7 @@ const assertAndSnapshotRaw = async (
test('should proxy image', async t => {
const assertAndSnapshot = assertAndSnapshotRaw.bind(null, t);
const imageUrl = `http://example.com/image-${Date.now()}.png`;
await assertAndSnapshot(
'/api/worker/image-proxy',
@@ -105,7 +109,7 @@ test('should proxy image', async t => {
{
await assertAndSnapshot(
'/api/worker/image-proxy?url=http://example.com/image.png',
`/api/worker/image-proxy?url=${imageUrl}`,
'should return 400 if origin and referer are missing',
{ status: 400, origin: null, referer: null }
);
@@ -113,14 +117,17 @@ test('should proxy image', async t => {
{
await assertAndSnapshot(
'/api/worker/image-proxy?url=http://example.com/image.png',
`/api/worker/image-proxy?url=${imageUrl}`,
'should return 400 for invalid origin header',
{ status: 400, origin: 'http://invalid.com' }
);
}
{
const fakeBuffer = Buffer.from('fake image');
const fakeBuffer = Buffer.from(
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jfJ8AAAAASUVORK5CYII=',
'base64'
);
const fakeResponse = new Response(fakeBuffer, {
status: 200,
headers: {
@@ -130,13 +137,14 @@ test('should proxy image', async t => {
});
const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeResponse);
await assertAndSnapshot(
'/api/worker/image-proxy?url=http://example.com/image.png',
'should return image buffer'
);
fetchSpy.restore();
try {
await assertAndSnapshot(
`/api/worker/image-proxy?url=${imageUrl}`,
'should return image buffer'
);
} finally {
fetchSpy.restore();
}
}
});
@@ -200,18 +208,19 @@ test('should preview link', async t => {
});
const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeHTML);
await assertAndSnapshot(
'/api/worker/link-preview',
'should process a valid external URL and return link preview data',
{
status: 200,
method: 'POST',
body: { url: 'http://external.com/page' },
}
);
fetchSpy.restore();
try {
await assertAndSnapshot(
'/api/worker/link-preview',
'should process a valid external URL and return link preview data',
{
status: 200,
method: 'POST',
body: { url: 'http://external.com/page' },
}
);
} finally {
fetchSpy.restore();
}
}
{
@@ -251,18 +260,19 @@ test('should preview link', async t => {
});
const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeHTML);
await assertAndSnapshot(
'/api/worker/link-preview',
'should decode HTML content with charset',
{
status: 200,
method: 'POST',
body: { url: `http://example.com/${charset}` },
}
);
fetchSpy.restore();
try {
await assertAndSnapshot(
'/api/worker/link-preview',
'should decode HTML content with charset',
{
status: 200,
method: 'POST',
body: { url: `http://example.com/${charset}` },
}
);
} finally {
fetchSpy.restore();
}
}
}
});
@@ -42,8 +42,18 @@ export class Ga4Client {
timestamp_micros: event.timestampMicros,
})),
};
await this.post(payload);
try {
await this.post(payload);
} catch {
if (env.DEPLOYMENT_TYPE === 'affine') {
// In production, we want to be resilient to GA4 failures, so we catch and ignore errors.
// In non-production environments, we rethrow to surface issues during development and testing.
console.info(
'Failed to send telemetry event to GA4:',
chunk.map(e => e.eventName).join(', ')
);
}
}
}
}
}
+313
View File
@@ -57,3 +57,316 @@ export const addDocToRootDoc = serverNativeModule.addDocToRootDoc;
export const updateDocTitle = serverNativeModule.updateDocTitle;
export const updateDocProperties = serverNativeModule.updateDocProperties;
export const updateRootDocMetaTitle = serverNativeModule.updateRootDocMetaTitle;
type NativeLlmModule = {
llmDispatch?: (
protocol: string,
backendConfigJson: string,
requestJson: string
) => string | Promise<string>;
llmDispatchStream?: (
protocol: string,
backendConfigJson: string,
requestJson: string,
callback: (error: Error | null, eventJson: string) => void
) => { abort?: () => void } | undefined;
};
const nativeLlmModule = serverNativeModule as typeof serverNativeModule &
NativeLlmModule;
export type NativeLlmProtocol =
| 'openai_chat'
| 'openai_responses'
| 'anthropic';
export type NativeLlmBackendConfig = {
base_url: string;
auth_token: string;
request_layer?: 'anthropic' | 'chat_completions' | 'responses' | 'vertex';
headers?: Record<string, string>;
no_streaming?: boolean;
timeout_ms?: number;
};
export type NativeLlmCoreRole = 'system' | 'user' | 'assistant' | 'tool';
export type NativeLlmCoreContent =
| { type: 'text'; text: string }
| { type: 'reasoning'; text: string; signature?: string }
| {
type: 'tool_call';
call_id: string;
name: string;
arguments: Record<string, unknown>;
thought?: string;
}
| {
type: 'tool_result';
call_id: string;
output: unknown;
is_error?: boolean;
name?: string;
arguments?: Record<string, unknown>;
}
| { type: 'image'; source: Record<string, unknown> | string };
export type NativeLlmCoreMessage = {
role: NativeLlmCoreRole;
content: NativeLlmCoreContent[];
};
export type NativeLlmToolDefinition = {
name: string;
description?: string;
parameters: Record<string, unknown>;
};
export type NativeLlmRequest = {
model: string;
messages: NativeLlmCoreMessage[];
stream?: boolean;
max_tokens?: number;
temperature?: number;
tools?: NativeLlmToolDefinition[];
tool_choice?: 'auto' | 'none' | 'required' | { name: string };
include?: string[];
reasoning?: Record<string, unknown>;
middleware?: {
request?: Array<
'normalize_messages' | 'clamp_max_tokens' | 'tool_schema_rewrite'
>;
stream?: Array<'stream_event_normalize' | 'citation_indexing'>;
config?: {
no_additional_properties?: boolean;
drop_property_format?: boolean;
drop_property_min_length?: boolean;
drop_array_min_items?: boolean;
drop_array_max_items?: boolean;
max_tokens_cap?: number;
};
};
};
export type NativeLlmDispatchResponse = {
id: string;
model: string;
message: NativeLlmCoreMessage;
usage: {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
cached_tokens?: number;
};
finish_reason: string;
reasoning_details?: unknown;
};
export type NativeLlmStreamEvent =
| { type: 'message_start'; id?: string; model?: string }
| { type: 'text_delta'; text: string }
| { type: 'reasoning_delta'; text: string }
| {
type: 'tool_call_delta';
call_id: string;
name?: string;
arguments_delta: string;
}
| {
type: 'tool_call';
call_id: string;
name: string;
arguments: Record<string, unknown>;
thought?: string;
}
| {
type: 'tool_result';
call_id: string;
output: unknown;
is_error?: boolean;
name?: string;
arguments?: Record<string, unknown>;
}
| { type: 'citation'; index: number; url: string }
| {
type: 'usage';
usage: {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
cached_tokens?: number;
};
}
| {
type: 'done';
finish_reason?: string;
usage?: {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
cached_tokens?: number;
};
}
| { type: 'error'; message: string; code?: string; raw?: string };
const LLM_STREAM_END_MARKER = '__AFFINE_LLM_STREAM_END__';
export async function llmDispatch(
protocol: NativeLlmProtocol,
backendConfig: NativeLlmBackendConfig,
request: NativeLlmRequest
): Promise<NativeLlmDispatchResponse> {
if (!nativeLlmModule.llmDispatch) {
throw new Error('native llm dispatch is not available');
}
const response = nativeLlmModule.llmDispatch(
protocol,
JSON.stringify(backendConfig),
JSON.stringify(request)
);
const responseText = await Promise.resolve(response);
return JSON.parse(responseText) as NativeLlmDispatchResponse;
}
export class NativeStreamAdapter<T> implements AsyncIterableIterator<T> {
readonly #queue: T[] = [];
readonly #waiters: ((result: IteratorResult<T>) => void)[] = [];
readonly #handle: { abort?: () => void } | undefined;
readonly #signal?: AbortSignal;
readonly #abortListener?: () => void;
#ended = false;
constructor(
handle: { abort?: () => void } | undefined,
signal?: AbortSignal
) {
this.#handle = handle;
this.#signal = signal;
if (signal?.aborted) {
this.close(true);
return;
}
if (signal) {
this.#abortListener = () => {
this.close(true);
};
signal.addEventListener('abort', this.#abortListener, { once: true });
}
}
private close(abortHandle: boolean) {
if (this.#ended) {
return;
}
this.#ended = true;
if (this.#signal && this.#abortListener) {
this.#signal.removeEventListener('abort', this.#abortListener);
}
if (abortHandle) {
this.#handle?.abort?.();
}
while (this.#waiters.length) {
const waiter = this.#waiters.shift();
waiter?.({ value: undefined as T, done: true });
}
}
push(value: T | null) {
if (this.#ended) {
return;
}
if (value === null) {
this.close(false);
return;
}
const waiter = this.#waiters.shift();
if (waiter) {
waiter({ value, done: false });
return;
}
this.#queue.push(value);
}
[Symbol.asyncIterator]() {
return this;
}
async next(): Promise<IteratorResult<T>> {
if (this.#queue.length > 0) {
const value = this.#queue.shift() as T;
return { value, done: false };
}
if (this.#ended) {
return { value: undefined as T, done: true };
}
return await new Promise(resolve => {
this.#waiters.push(resolve);
});
}
async return(): Promise<IteratorResult<T>> {
this.close(true);
return { value: undefined as T, done: true };
}
}
export function llmDispatchStream(
protocol: NativeLlmProtocol,
backendConfig: NativeLlmBackendConfig,
request: NativeLlmRequest,
signal?: AbortSignal
): AsyncIterableIterator<NativeLlmStreamEvent> {
if (!nativeLlmModule.llmDispatchStream) {
throw new Error('native llm stream dispatch is not available');
}
let adapter: NativeStreamAdapter<NativeLlmStreamEvent> | undefined;
const buffer: (NativeLlmStreamEvent | null)[] = [];
let pushFn = (event: NativeLlmStreamEvent | null) => {
buffer.push(event);
};
const handle = nativeLlmModule.llmDispatchStream(
protocol,
JSON.stringify(backendConfig),
JSON.stringify(request),
(error, eventJson) => {
if (error) {
pushFn({ type: 'error', message: error.message, raw: eventJson });
return;
}
if (eventJson === LLM_STREAM_END_MARKER) {
pushFn(null);
return;
}
try {
pushFn(JSON.parse(eventJson) as NativeLlmStreamEvent);
} catch (error) {
pushFn({
type: 'error',
message:
error instanceof Error
? error.message
: 'failed to parse native stream event',
raw: eventJson,
});
}
}
);
adapter = new NativeStreamAdapter(handle, signal);
pushFn = event => {
adapter.push(event);
};
for (const event of buffer) {
adapter.push(event);
}
return adapter;
}
@@ -1,3 +1,5 @@
import { z } from 'zod';
import {
defineModuleConfig,
StorageJSONSchema,
@@ -13,7 +15,179 @@ import { GeminiGenerativeConfig, GeminiVertexConfig } from './providers/gemini';
import { MorphConfig } from './providers/morph';
import { OpenAIConfig } from './providers/openai';
import { PerplexityConfig } from './providers/perplexity';
import { VertexSchema } from './providers/types';
import {
CopilotProviderType,
ModelOutputType,
VertexSchema,
} from './providers/types';
export type CopilotProviderConfigMap = {
[CopilotProviderType.OpenAI]: OpenAIConfig;
[CopilotProviderType.FAL]: FalConfig;
[CopilotProviderType.Gemini]: GeminiGenerativeConfig;
[CopilotProviderType.GeminiVertex]: GeminiVertexConfig;
[CopilotProviderType.Perplexity]: PerplexityConfig;
[CopilotProviderType.Anthropic]: AnthropicOfficialConfig;
[CopilotProviderType.AnthropicVertex]: AnthropicVertexConfig;
[CopilotProviderType.Morph]: MorphConfig;
};
export type ProviderSpecificConfig =
CopilotProviderConfigMap[keyof CopilotProviderConfigMap];
export const RustRequestMiddlewareValues = [
'normalize_messages',
'clamp_max_tokens',
'tool_schema_rewrite',
] as const;
export type RustRequestMiddleware =
(typeof RustRequestMiddlewareValues)[number];
export const RustStreamMiddlewareValues = [
'stream_event_normalize',
'citation_indexing',
] as const;
export type RustStreamMiddleware = (typeof RustStreamMiddlewareValues)[number];
export const NodeTextMiddlewareValues = [
'citation_footnote',
'callout',
'thinking_format',
] as const;
export type NodeTextMiddleware = (typeof NodeTextMiddlewareValues)[number];
export type ProviderMiddlewareConfig = {
rust?: { request?: RustRequestMiddleware[]; stream?: RustStreamMiddleware[] };
node?: { text?: NodeTextMiddleware[] };
};
type CopilotProviderProfileCommon = {
id: string;
displayName?: string;
priority?: number;
enabled?: boolean;
models?: string[];
middleware?: ProviderMiddlewareConfig;
};
type CopilotProviderProfileVariant<T extends CopilotProviderType> = {
type: T;
config: CopilotProviderConfigMap[T];
};
export type CopilotProviderProfile = CopilotProviderProfileCommon &
{
[Type in CopilotProviderType]: CopilotProviderProfileVariant<Type>;
}[CopilotProviderType];
export type CopilotProviderDefaults = Partial<
Record<ModelOutputType, string>
> & {
fallback?: string;
};
const CopilotProviderProfileBaseShape = z.object({
id: z.string().regex(/^[a-zA-Z0-9-_]+$/),
displayName: z.string().optional(),
priority: z.number().optional(),
enabled: z.boolean().optional(),
models: z.array(z.string()).optional(),
middleware: z
.object({
rust: z
.object({
request: z.array(z.enum(RustRequestMiddlewareValues)).optional(),
stream: z.array(z.enum(RustStreamMiddlewareValues)).optional(),
})
.optional(),
node: z
.object({ text: z.array(z.enum(NodeTextMiddlewareValues)).optional() })
.optional(),
})
.optional(),
});
const OpenAIConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
oldApiStyle: z.boolean().optional(),
});
const FalConfigShape = z.object({
apiKey: z.string(),
});
const GeminiGenerativeConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
});
const VertexProviderConfigShape = z.object({
location: z.string().optional(),
project: z.string().optional(),
baseURL: z.string().optional(),
googleAuthOptions: z.any().optional(),
fetch: z.any().optional(),
});
const PerplexityConfigShape = z.object({
apiKey: z.string(),
endpoint: z.string().optional(),
});
const AnthropicOfficialConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
});
const MorphConfigShape = z.object({
apiKey: z.string().optional(),
});
const CopilotProviderProfileShape = z.discriminatedUnion('type', [
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.OpenAI),
config: OpenAIConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.FAL),
config: FalConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Gemini),
config: GeminiGenerativeConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.GeminiVertex),
config: VertexProviderConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Perplexity),
config: PerplexityConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Anthropic),
config: AnthropicOfficialConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.AnthropicVertex),
config: VertexProviderConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Morph),
config: MorphConfigShape,
}),
]);
const CopilotProviderDefaultsShape = z.object({
[ModelOutputType.Text]: z.string().optional(),
[ModelOutputType.Object]: z.string().optional(),
[ModelOutputType.Embedding]: z.string().optional(),
[ModelOutputType.Image]: z.string().optional(),
[ModelOutputType.Structured]: z.string().optional(),
fallback: z.string().optional(),
});
declare global {
interface AppConfigSchema {
copilot: {
@@ -27,6 +201,8 @@ declare global {
storage: ConfigItem<StorageProviderConfig>;
scenarios: ConfigItem<CopilotPromptScenario>;
providers: {
profiles: ConfigItem<CopilotProviderProfile[]>;
defaults: ConfigItem<CopilotProviderDefaults>;
openai: ConfigItem<OpenAIConfig>;
fal: ConfigItem<FalConfig>;
gemini: ConfigItem<GeminiGenerativeConfig>;
@@ -63,6 +239,16 @@ defineModuleConfig('copilot', {
},
},
},
'providers.profiles': {
desc: 'The profile list for copilot providers.',
default: [],
shape: z.array(CopilotProviderProfileShape),
},
'providers.defaults': {
desc: 'The default provider ids for model output types and global fallback.',
default: {},
shape: CopilotProviderDefaultsShape,
},
'providers.openai': {
desc: 'The config for the openai provider.',
default: {
@@ -36,10 +36,7 @@ import {
BlobNotFound,
CallMetric,
Config,
CopilotFailedToGenerateText,
CopilotSessionNotFound,
InternalServerError,
mapAnyError,
mapSseError,
metrics,
NoCopilotProviderAvailable,
@@ -242,61 +239,6 @@ export class CopilotController implements BeforeApplicationShutdown {
};
}
@Get('/chat/:sessionId')
@CallMetric('ai', 'chat', { timer: true })
async chat(
@CurrentUser() user: CurrentUser,
@Req() req: Request,
@Param('sessionId') sessionId: string,
@Query() query: Record<string, string | string[]>
): Promise<string> {
const info: any = { sessionId, params: query };
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_calls').add(1, { model });
const { reasoning, webSearch, toolsConfig } =
ChatQuerySchema.parse(query);
const content = await provider.text({ modelId: model }, finalMessage, {
...session.config.promptConfig,
signal: getSignal(req).signal,
user: user.id,
session: session.config.sessionId,
workspace: session.config.workspaceId,
reasoning,
webSearch,
tools: getTools(session.config.promptConfig?.tools, toolsConfig),
});
session.push({
role: 'assistant',
content,
createdAt: new Date(),
});
await session.save();
return content;
} catch (e: any) {
metrics.ai.counter('chat_errors').add(1);
let error = mapAnyError(e);
if (error instanceof InternalServerError) {
error = new CopilotFailedToGenerateText(e.message);
}
error.log('CopilotChat', info);
throw error;
}
}
@Sse('/chat/:sessionId/stream')
@CallMetric('ai', 'chat_stream', { timer: true })
async chatStream(
@@ -3,7 +3,7 @@ import { AiPrompt, PrismaClient } from '@prisma/client';
import type { PromptConfig, PromptMessage } from '../providers/types';
type Prompt = Omit<
export type Prompt = Omit<
AiPrompt,
| 'id'
| 'createdAt'
@@ -2095,17 +2095,14 @@ export const prompts: Prompt[] = [
export async function refreshPrompts(db: PrismaClient) {
const needToSkip = await db.aiPrompt
.findMany({
where: { modified: true },
select: { name: true },
})
.findMany({ where: { modified: true }, select: { name: true } })
.then(p => p.map(p => p.name));
for (const prompt of prompts) {
// skip prompt update if already modified by admin panel
if (needToSkip.includes(prompt.name)) {
new Logger('CopilotPrompt').warn(`Skip modified prompt: ${prompt.name}`);
return;
continue;
}
await db.aiPrompt.upsert({
@@ -12,6 +12,7 @@ import {
import { ChatPrompt } from './chat-prompt';
import {
CopilotPromptScenario,
type Prompt,
prompts,
refreshPrompts,
Scenario,
@@ -21,6 +22,7 @@ import {
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,
@@ -28,7 +30,7 @@ export class PromptService implements OnApplicationBootstrap {
) {}
async onApplicationBootstrap() {
this.cache.clear();
this.resetInMemoryPrompts();
await refreshPrompts(this.db);
}
@@ -45,6 +47,7 @@ export class PromptService implements OnApplicationBootstrap {
}
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)) {
@@ -75,25 +78,29 @@ export class PromptService implements OnApplicationBootstrap {
* @returns prompt names
*/
async listNames() {
return this.db.aiPrompt
.findMany({ select: { name: true } })
.then(prompts => Array.from(new Set(prompts.map(p => p.name))));
this.ensureInMemoryPrompts();
return Array.from(this.inMemoryPrompts.keys());
}
async list() {
return this.db.aiPrompt.findMany({
select: {
name: true,
action: true,
model: true,
config: true,
messages: {
select: { role: true, content: true, params: true },
orderBy: { idx: 'asc' },
},
},
orderBy: { action: { sort: 'asc', nulls: 'first' } },
});
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 ?? '');
});
}
/**
@@ -102,40 +109,24 @@ export class PromptService implements OnApplicationBootstrap {
* @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;
}
const prompt = await this.db.aiPrompt.findUnique({
where: {
name,
},
select: {
name: true,
action: true,
model: true,
optionalModels: true,
config: true,
messages: {
select: {
role: true,
content: true,
params: true,
},
orderBy: {
idx: 'asc',
},
},
},
});
const prompt = this.inMemoryPrompts.get(name);
if (!prompt) return null;
const messages = PromptMessageSchema.array().safeParse(prompt?.messages);
const config = PromptConfigSchema.safeParse(prompt?.config);
if (prompt && messages.success && config.success) {
const messages = PromptMessageSchema.array().safeParse(prompt.messages);
const config = PromptConfigSchema.safeParse(prompt.config);
if (messages.success && config.success) {
const chatPrompt = ChatPrompt.createFromPrompt({
...prompt,
...this.clonePrompt(prompt),
action: prompt.action ?? null,
optionalModels: prompt.optionalModels ?? [],
config: config.data,
messages: messages.data,
});
@@ -149,25 +140,69 @@ export class PromptService implements OnApplicationBootstrap {
name: string,
model: string,
messages: PromptMessage[],
config?: PromptConfig | null
config?: PromptConfig | null,
extraConfig?: { optionalModels: string[] }
) {
return await this.db.aiPrompt
.create({
data: {
name,
model,
config: config || undefined,
messages: {
create: messages.map((m, idx) => ({
idx,
...m,
attachments: m.attachments || undefined,
params: m.params || undefined,
})),
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,
})),
},
},
},
})
.then(ret => ret.id);
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;
}
}
@Transactional()
@@ -177,44 +212,123 @@ export class PromptService implements OnApplicationBootstrap {
messages?: PromptMessage[];
model?: string;
modified?: boolean;
config?: PromptConfig;
config?: PromptConfig | null;
},
where?: Prisma.AiPromptWhereInput
) {
this.ensureInMemoryPrompts();
const { config, messages, model, modified } = data;
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: config || undefined,
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,
},
});
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) {
const { id } = await this.db.aiPrompt.delete({ where: { name } });
this.inMemoryPrompts.delete(name);
this.cache.delete(name);
return id;
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[]) {
return messages.map(message => ({
...message,
attachments: message.attachments ? [...message.attachments] : undefined,
params: message.params ? structuredClone(message.params) : undefined,
}));
}
private clonePrompt(prompt: Prompt): Prompt {
return {
...prompt,
optionalModels: prompt.optionalModels
? [...prompt.optionalModels]
: undefined,
config: prompt.config ? structuredClone(prompt.config) : undefined,
messages: this.cloneMessages(prompt.messages),
};
}
private stringifyError(error: unknown) {
return error instanceof Error ? error.message : String(error);
}
}
@@ -1,52 +1,90 @@
import {
type AnthropicProvider as AnthropicSDKProvider,
type AnthropicProviderOptions,
} from '@ai-sdk/anthropic';
import { type GoogleVertexAnthropicProvider } from '@ai-sdk/google-vertex/anthropic';
import { AISDKError, generateText, stepCountIs, streamText } from 'ai';
import type { ToolSet } from 'ai';
import {
CopilotProviderSideError,
metrics,
UserFriendlyError,
} from '../../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
} from '../../../../native';
import type { NodeTextMiddleware } from '../../config';
import { buildNativeRequest, NativeProviderAdapter } from '../native';
import { CopilotProvider } from '../provider';
import type {
CopilotChatOptions,
CopilotProviderModel,
ModelConditions,
PromptMessage,
StreamObject,
} from '../types';
import { ModelOutputType } from '../types';
import {
chatToGPTMessage,
StreamObjectParser,
TextStreamParser,
} from '../utils';
import { CopilotProviderType, ModelOutputType } from '../types';
import { getGoogleAuth, getVertexAnthropicBaseUrl } from '../utils';
export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
protected abstract instance:
| AnthropicSDKProvider
| GoogleVertexAnthropicProvider;
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
} else if (e instanceof AISDKError) {
this.logger.error('Throw error from ai sdk:', e);
return new CopilotProviderSideError({
provider: this.type,
kind: e.name || 'unknown',
message: e.message,
});
} else {
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected anthropic response',
});
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected anthropic response',
});
}
private async createNativeConfig(): Promise<NativeLlmBackendConfig> {
if (this.type === CopilotProviderType.AnthropicVertex) {
const auth = await getGoogleAuth(this.config as any, 'anthropic');
const headers = auth.headers();
const authorization =
headers.Authorization ||
(headers as Record<string, string | undefined>).authorization;
const token =
typeof authorization === 'string'
? authorization.replace(/^Bearer\s+/i, '')
: '';
const baseUrl =
getVertexAnthropicBaseUrl(this.config as any) || auth.baseUrl;
return {
base_url: baseUrl || '',
auth_token: token,
request_layer: 'vertex',
headers,
};
}
const config = this.config as { apiKey: string; baseURL?: string };
const baseUrl = config.baseURL || 'https://api.anthropic.com/v1';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: config.apiKey,
};
}
private createAdapter(
backendConfig: NativeLlmBackendConfig,
tools: ToolSet,
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(
@@ -59,28 +97,29 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_calls').add(1, { model: model.id });
const [system, msgs] = await chatToGPTMessage(messages, true, true);
const modelInstance = this.instance(model.id);
const { text, reasoning } = await generateText({
model: modelInstance,
system,
messages: msgs,
abortSignal: options.signal,
providerOptions: {
anthropic: this.getAnthropicOptions(options, model.id),
},
tools: await this.getTools(options, model.id),
stopWhen: stepCountIs(this.MAX_STEPS),
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 { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
reasoning,
middleware,
});
if (!text) throw new Error('Failed to generate text');
return reasoning ? `${reasoning}\n${text}` : text;
const adapter = this.createAdapter(
backendConfig,
tools,
middleware.node?.text
);
return await adapter.text(request, options.signal);
} catch (e: any) {
metrics.ai.counter('chat_text_errors').add(1, { model: model.id });
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -95,25 +134,32 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_stream_calls').add(1, { model: model.id });
const fullStream = await this.getFullStream(model, messages, options);
const parser = new TextStreamParser();
for await (const chunk of fullStream) {
const result = parser.parse(chunk);
yield result;
if (options.signal?.aborted) {
await fullStream.cancel();
break;
}
}
if (!options.signal?.aborted) {
const footnotes = parser.end();
if (footnotes.length) {
yield `\n\n${footnotes}`;
}
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 { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
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)) {
yield chunk;
}
} catch (e: any) {
metrics.ai.counter('chat_text_stream_errors').add(1, { model: model.id });
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -130,58 +176,34 @@ export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
try {
metrics.ai
.counter('chat_object_stream_calls')
.add(1, { model: model.id });
const fullStream = await this.getFullStream(model, messages, options);
const parser = new StreamObjectParser();
for await (const chunk of fullStream) {
const result = parser.parse(chunk);
if (result) {
yield result;
}
if (options.signal?.aborted) {
await fullStream.cancel();
break;
}
.add(1, this.metricLabels(model.id));
const backendConfig = await this.createNativeConfig();
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
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)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_object_stream_errors')
.add(1, { model: model.id });
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
private async getFullStream(
model: CopilotProviderModel,
messages: PromptMessage[],
options: CopilotChatOptions = {}
) {
const [system, msgs] = await chatToGPTMessage(messages, true, true);
const { fullStream } = streamText({
model: this.instance(model.id),
system,
messages: msgs,
abortSignal: options.signal,
providerOptions: {
anthropic: this.getAnthropicOptions(options, model.id),
},
tools: await this.getTools(options, model.id),
stopWhen: stepCountIs(this.MAX_STEPS),
});
return fullStream;
}
private getAnthropicOptions(options: CopilotChatOptions, model: string) {
const result: AnthropicProviderOptions = {};
if (options?.reasoning && this.isReasoningModel(model)) {
result.thinking = {
type: 'enabled',
budgetTokens: 12000,
};
}
return result;
}
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,3 @@
import {
type AnthropicProvider as AnthropicSDKProvider,
createAnthropic,
} from '@ai-sdk/anthropic';
import z from 'zod';
import { CopilotProviderType, ModelInputType, ModelOutputType } from '../types';
@@ -52,18 +48,12 @@ export class AnthropicOfficialProvider extends AnthropicProvider<AnthropicOffici
},
];
protected instance!: AnthropicSDKProvider;
override configured(): boolean {
return !!this.config.apiKey;
}
override setup() {
super.setup();
this.instance = createAnthropic({
apiKey: this.config.apiKey,
baseURL: this.config.baseURL,
});
}
override async refreshOnlineModels() {
@@ -5,7 +5,11 @@ import {
} from '@ai-sdk/google-vertex/anthropic';
import { CopilotProviderType, ModelInputType, ModelOutputType } from '../types';
import { getGoogleAuth, VertexModelListSchema } from '../utils';
import {
getGoogleAuth,
getVertexAnthropicBaseUrl,
VertexModelListSchema,
} from '../utils';
import { AnthropicProvider } from './anthropic';
export type AnthropicVertexConfig = GoogleVertexAnthropicProviderSettings;
@@ -49,7 +53,8 @@ export class AnthropicVertexProvider extends AnthropicProvider<AnthropicVertexCo
protected instance!: GoogleVertexAnthropicProvider;
override configured(): boolean {
return !!this.config.location && !!this.config.googleAuthOptions;
if (!this.config.location || !this.config.googleAuthOptions) return false;
return !!this.config.project || !!getVertexAnthropicBaseUrl(this.config);
}
override setup() {
@@ -1,16 +1,141 @@
import { Injectable, Logger } from '@nestjs/common';
import { Config } from '../../../base';
import { ServerFeature, ServerService } from '../../../core';
import type { CopilotProvider } from './provider';
import {
buildProviderRegistry,
resolveModel,
stripProviderPrefix,
} from './provider-registry';
import { CopilotProviderType, ModelFullConditions } from './types';
function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
return (
value !== null &&
value !== undefined &&
typeof (value as AsyncIterable<unknown>)[Symbol.asyncIterator] ===
'function'
);
}
@Injectable()
export class CopilotProviderFactory {
constructor(private readonly server: ServerService) {}
constructor(
private readonly server: ServerService,
private readonly config: Config
) {}
private readonly logger = new Logger(CopilotProviderFactory.name);
readonly #providers = new Map<CopilotProviderType, CopilotProvider>();
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);
}
private getPreferredProviderIds(type?: CopilotProviderType) {
if (!type) return undefined;
return this.#providerIdsByType.get(type);
}
private normalizeCond(
providerId: string,
cond: ModelFullConditions
): ModelFullConditions {
const registry = this.getRegistry();
const modelId = stripProviderPrefix(registry, providerId, cond.modelId);
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;
}
const cond = first as Record<string, unknown>;
if (typeof cond.modelId !== 'string') return args;
const registry = this.getRegistry();
const modelId = stripProviderPrefix(registry, providerId, cond.modelId);
return [{ ...cond, modelId }, ...rest];
}
private wrapAsyncIterable<T>(
provider: CopilotProvider,
providerId: string,
iterable: AsyncIterable<T>
): AsyncIterableIterator<T> {
const iterator = iterable[Symbol.asyncIterator]();
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;
},
};
}
private getBoundProvider(providerId: string, provider: CopilotProvider) {
const cached = this.#boundProviders.get(providerId);
if (cached) {
return cached;
}
const wrapped = new Proxy(provider, {
get: (target, prop, receiver) => {
if (prop === 'providerId') {
return providerId;
}
const value = Reflect.get(target, prop, receiver);
if (typeof value !== 'function') {
return value;
}
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;
this.#boundProviders.set(providerId, wrapped);
return wrapped;
}
async getProvider(
cond: ModelFullConditions,
@@ -21,22 +146,41 @@ export class CopilotProviderFactory {
this.logger.debug(
`Resolving copilot provider for output type: ${cond.outputType}`
);
let candidate: CopilotProvider | null = null;
for (const [type, provider] of this.#providers.entries()) {
if (filter.prefer && filter.prefer !== type) {
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 isMatched = await provider.match(cond);
const matched = await provider.runWithProfile(providerId, () =>
provider.match(normalizedCond)
);
if (!matched) continue;
if (isMatched) {
candidate = provider;
this.logger.debug(`Copilot provider candidate found: ${type}`);
break;
}
this.logger.debug(
`Copilot provider candidate found: ${provider.type} (${providerId})`
);
return this.getBoundProvider(providerId, provider);
}
return candidate;
return null;
}
async getProviderByModel(
@@ -46,31 +190,50 @@ export class CopilotProviderFactory {
} = {}
): Promise<CopilotProvider | null> {
this.logger.debug(`Resolving copilot provider for model: ${modelId}`);
return this.getProvider({ modelId }, filter);
}
let candidate: CopilotProvider | null = null;
for (const [type, provider] of this.#providers.entries()) {
if (filter.prefer && filter.prefer !== type) {
continue;
}
if (await provider.match({ modelId })) {
candidate = provider;
this.logger.debug(`Copilot provider candidate found: ${type}`);
register(providerId: string, provider: CopilotProvider) {
const existed = this.#providers.get(providerId);
if (existed?.type && existed.type !== provider.type) {
const ids = this.#providerIdsByType.get(existed.type);
ids?.delete(providerId);
if (!ids?.size) {
this.#providerIdsByType.delete(existed.type);
}
}
return candidate;
}
this.#providers.set(providerId, provider);
this.#boundProviders.delete(providerId);
register(provider: CopilotProvider) {
this.#providers.set(provider.type, provider);
this.logger.log(`Copilot provider [${provider.type}] registered.`);
const ids = this.#providerIdsByType.get(provider.type) ?? new Set<string>();
ids.add(providerId);
this.#providerIdsByType.set(provider.type, ids);
this.logger.log(
`Copilot provider [${provider.type}] registered as [${providerId}].`
);
this.server.enableFeature(ServerFeature.Copilot);
}
unregister(provider: CopilotProvider) {
this.#providers.delete(provider.type);
this.logger.log(`Copilot provider [${provider.type}] unregistered.`);
unregister(providerId: string, provider: CopilotProvider) {
const existed = this.#providers.get(providerId);
if (!existed || existed !== provider) {
return;
}
this.#providers.delete(providerId);
this.#boundProviders.delete(providerId);
const ids = this.#providerIdsByType.get(provider.type);
ids?.delete(providerId);
if (!ids?.size) {
this.#providerIdsByType.delete(provider.type);
}
this.logger.log(
`Copilot provider [${provider.type}] unregistered from [${providerId}].`
);
if (this.#providers.size === 0) {
this.server.disableFeature(ServerFeature.Copilot);
}
@@ -0,0 +1,381 @@
import type { ToolSet } from 'ai';
import { z } from 'zod';
import type {
NativeLlmRequest,
NativeLlmStreamEvent,
NativeLlmToolDefinition,
} from '../../../native';
export type NativeDispatchFn = (
request: NativeLlmRequest,
signal?: AbortSignal
) => AsyncIterableIterator<NativeLlmStreamEvent>;
export type NativeToolCall = {
id: string;
name: string;
args: Record<string, unknown>;
thought?: string;
};
type ToolCallState = {
name?: string;
argumentsText: string;
};
type ToolExecutionResult = {
callId: string;
name: string;
args: Record<string, unknown>;
output: unknown;
isError?: boolean;
};
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);
return {
id: event.call_id,
name: event.name || state?.name || '',
args: this.parseArgs(
event.arguments ?? this.parseJson(state?.argumentsText ?? '{}')
),
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,
args: this.parseArgs(this.parseJson(state.argumentsText)),
});
}
this.#states.clear();
return pending;
}
private parseJson(jsonText: string): unknown {
if (!jsonText.trim()) {
return {};
}
try {
return JSON.parse(jsonText);
} catch {
return {};
}
}
private parseArgs(value: unknown): Record<string, unknown> {
if (value && typeof value === 'object' && !Array.isArray(value)) {
return value as Record<string, unknown>;
}
return {};
}
}
export class ToolSchemaExtractor {
static extract(toolSet: ToolSet): NativeLlmToolDefinition[] {
return Object.entries(toolSet).map(([name, tool]) => {
const unknownTool = tool as Record<string, unknown>;
const inputSchema =
unknownTool.inputSchema ?? unknownTool.parameters ?? z.object({});
return {
name,
description:
typeof unknownTool.description === 'string'
? unknownTool.description
: undefined,
parameters: this.toJsonSchema(inputSchema),
};
});
}
private 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: ToolSet,
private readonly maxSteps = 20
) {}
async *run(
request: NativeLlmRequest,
signal?: AbortSignal
): AsyncIterableIterator<NativeLlmStreamEvent> {
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,
},
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);
messages.push({
role: 'assistant',
content: toolCalls.map(call => ({
type: 'tool_call',
call_id: call.id,
name: call.name,
arguments: call.args,
thought: call.thought,
})),
});
for (const result of toolResults) {
messages.push({
role: 'tool',
content: [
{
type: 'tool_result',
call_id: result.callId,
output: result.output,
is_error: result.isError,
},
],
});
yield {
type: 'tool_result',
call_id: result.callId,
name: result.name,
arguments: result.args,
output: result.output,
is_error: result.isError,
};
}
}
}
private async executeTools(calls: NativeToolCall[]) {
return await Promise.all(calls.map(call => this.executeTool(call)));
}
private async executeTool(
call: NativeToolCall
): Promise<ToolExecutionResult> {
const tool = this.tools[call.name] as
| {
execute?: (args: Record<string, unknown>) => Promise<unknown>;
}
| undefined;
if (!tool?.execute) {
return {
callId: call.id,
name: call.name,
args: call.args,
isError: true,
output: {
message: `Tool not found: ${call.name}`,
},
};
}
try {
const output = await tool.execute(call.args);
return {
callId: call.id,
name: call.name,
args: call.args,
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,
isError: true,
output: {
message: 'Tool execution failed',
},
};
}
}
}
@@ -1,14 +1,17 @@
import {
createOpenAICompatible,
OpenAICompatibleProvider as VercelOpenAICompatibleProvider,
} from '@ai-sdk/openai-compatible';
import { AISDKError, generateText, streamText } from 'ai';
import type { ToolSet } from 'ai';
import {
CopilotProviderSideError,
metrics,
UserFriendlyError,
} from '../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
} from '../../../native';
import type { NodeTextMiddleware } from '../config';
import { buildNativeRequest, NativeProviderAdapter } from './native';
import { CopilotProvider } from './provider';
import type {
CopilotChatOptions,
@@ -16,7 +19,6 @@ import type {
PromptMessage,
} from './types';
import { CopilotProviderType, ModelInputType, ModelOutputType } from './types';
import { chatToGPTMessage, TextStreamParser } from './utils';
export const DEFAULT_DIMENSIONS = 256;
@@ -57,37 +59,48 @@ export class MorphProvider extends CopilotProvider<MorphConfig> {
},
];
#instance!: VercelOpenAICompatibleProvider;
override configured(): boolean {
return !!this.config.apiKey;
}
protected override setup() {
super.setup();
this.#instance = createOpenAICompatible({
name: this.type,
apiKey: this.config.apiKey,
baseURL: 'https://api.morphllm.com/v1',
});
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
} else if (e instanceof AISDKError) {
return new CopilotProviderSideError({
provider: this.type,
kind: e.name || 'unknown',
message: e.message,
});
} else {
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected morph response',
});
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected morph response',
});
}
private createNativeConfig(): NativeLlmBackendConfig {
return {
base_url: 'https://api.morphllm.com',
auth_token: this.config.apiKey ?? '',
};
}
private createNativeAdapter(
tools: ToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream(
'openai_chat',
this.createNativeConfig(),
request,
signal
),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
);
}
async text(
@@ -103,22 +116,22 @@ export class MorphProvider extends CopilotProvider<MorphConfig> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_calls').add(1, { model: model.id });
const [system, msgs] = await chatToGPTMessage(messages);
const modelInstance = this.#instance(model.id);
const { text } = await generateText({
model: modelInstance,
system,
messages: msgs,
abortSignal: options.signal,
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 text.trim();
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
return await adapter.text(request, options.signal);
} catch (e: any) {
metrics.ai.counter('chat_text_errors').add(1, { model: model.id });
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -136,38 +149,26 @@ export class MorphProvider extends CopilotProvider<MorphConfig> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_stream_calls').add(1, { model: model.id });
const [system, msgs] = await chatToGPTMessage(messages);
const modelInstance = this.#instance(model.id);
const { fullStream } = streamText({
model: modelInstance,
system,
messages: msgs,
abortSignal: options.signal,
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 textParser = new TextStreamParser();
for await (const chunk of fullStream) {
switch (chunk.type) {
case 'text-delta': {
let result = textParser.parse(chunk);
yield result;
break;
}
default: {
yield textParser.parse(chunk);
break;
}
}
if (options.signal?.aborted) {
await fullStream.cancel();
break;
}
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
for await (const chunk of adapter.streamText(request, options.signal)) {
yield chunk;
}
} catch (e: any) {
metrics.ai.counter('chat_text_stream_errors').add(1, { model: model.id });
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -0,0 +1,464 @@
import type { ToolSet } from 'ai';
import { ZodType } from 'zod';
import type {
NativeLlmCoreContent,
NativeLlmCoreMessage,
NativeLlmRequest,
NativeLlmStreamEvent,
} from '../../../native';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import { NativeDispatchFn, ToolCallLoop, ToolSchemaExtractor } from './loop';
import type { CopilotChatOptions, PromptMessage, StreamObject } from './types';
import {
CitationFootnoteFormatter,
inferMimeType,
TextStreamParser,
} from './utils';
const SIMPLE_IMAGE_URL_REGEX = /^(https?:\/\/|data:image\/)/;
type BuildNativeRequestOptions = {
model: string;
messages: PromptMessage[];
options?: CopilotChatOptions;
tools?: ToolSet;
withAttachment?: boolean;
include?: string[];
reasoning?: Record<string, unknown>;
middleware?: ProviderMiddlewareConfig;
};
type BuildNativeRequestResult = {
request: NativeLlmRequest;
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';
}
}
async function toCoreContents(
message: PromptMessage,
withAttachment: boolean
): 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) {
let attachmentUrl: string;
let mediaType: string;
if (typeof entry === 'string') {
attachmentUrl = entry;
mediaType =
typeof message.params?.mimetype === 'string'
? message.params.mimetype
: await inferMimeType(entry);
} else {
attachmentUrl = entry.attachment;
mediaType = entry.mimeType;
}
if (!SIMPLE_IMAGE_URL_REGEX.test(attachmentUrl)) continue;
if (!mediaType.startsWith('image/')) continue;
contents.push({ type: 'image', source: { url: attachmentUrl } });
}
return contents;
}
export async function buildNativeRequest({
model,
messages,
options = {},
tools = {},
withAttachment = true,
include,
reasoning,
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 =
systemMessage?.params?.schema instanceof ZodType
? systemMessage.params.schema
: undefined;
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);
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,
middleware: middleware?.rust
? { request: middleware.rust.request, stream: middleware.rust.stream }
: undefined,
},
schema,
};
}
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: ToolSet,
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) {
let output = '';
for await (const chunk of this.streamText(request, signal)) {
output += chunk;
}
return output.trim();
}
async *streamText(
request: NativeLlmRequest,
signal?: AbortSignal
): 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)) {
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
): 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)) {
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;
}
}
}
}
@@ -1,53 +1,35 @@
import {
createOpenAI,
openai,
type OpenAIProvider as VercelOpenAIProvider,
OpenAIResponsesProviderOptions,
} from '@ai-sdk/openai';
import {
createOpenAICompatible,
type OpenAICompatibleProvider as VercelOpenAICompatibleProvider,
} from '@ai-sdk/openai-compatible';
import {
AISDKError,
embedMany,
experimental_generateImage as generateImage,
generateObject,
generateText,
stepCountIs,
streamText,
Tool,
} from 'ai';
import type { Tool, ToolSet } from 'ai';
import { z } from 'zod';
import {
CopilotPromptInvalid,
CopilotProviderNotSupported,
CopilotProviderSideError,
fetchBuffer,
metrics,
OneMB,
readResponseBufferWithLimit,
safeFetch,
UserFriendlyError,
} from '../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
} from '../../../native';
import type { NodeTextMiddleware } from '../config';
import { buildNativeRequest, NativeProviderAdapter } from './native';
import { CopilotProvider } from './provider';
import type {
CopilotChatOptions,
CopilotChatTools,
CopilotEmbeddingOptions,
CopilotImageOptions,
CopilotProviderModel,
CopilotStructuredOptions,
ModelConditions,
PromptMessage,
StreamObject,
} from './types';
import { CopilotProviderType, ModelInputType, ModelOutputType } from './types';
import {
chatToGPTMessage,
CitationParser,
StreamObjectParser,
TextStreamParser,
} from './utils';
import { chatToGPTMessage } from './utils';
export const DEFAULT_DIMENSIONS = 256;
@@ -63,7 +45,12 @@ const ModelListSchema = z.object({
const ImageResponseSchema = z.union([
z.object({
data: z.array(z.object({ b64_json: z.string() })),
data: z.array(
z.object({
b64_json: z.string().optional(),
url: z.string().optional(),
})
),
}),
z.object({
error: z.object({
@@ -87,6 +74,38 @@ const LogProbsSchema = z.array(
})
);
const TRUSTED_ATTACHMENT_HOST_SUFFIXES = ['cdn.affine.pro'];
function normalizeImageFormatToMime(format?: string) {
switch (format?.toLowerCase()) {
case 'jpg':
case 'jpeg':
return 'image/jpeg';
case 'webp':
return 'image/webp';
case 'png':
return 'image/png';
case 'gif':
return 'image/gif';
default:
return 'image/png';
}
}
function normalizeImageResponseData(
data: { b64_json?: string; url?: string }[],
mimeType: string = 'image/png'
) {
return data
.map(image => {
if (image.b64_json) {
return `data:${mimeType};base64,${image.b64_json}`;
}
return image.url;
})
.filter((value): value is string => typeof value === 'string');
}
export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
readonly type = CopilotProviderType.OpenAI;
@@ -319,53 +338,23 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
},
];
#instance!: VercelOpenAIProvider | VercelOpenAICompatibleProvider;
override configured(): boolean {
return !!this.config.apiKey;
}
protected override setup() {
super.setup();
this.#instance =
this.config.oldApiStyle && this.config.baseURL
? createOpenAICompatible({
name: 'openai-compatible-old-style',
apiKey: this.config.apiKey,
baseURL: this.config.baseURL,
})
: createOpenAI({
apiKey: this.config.apiKey,
baseURL: this.config.baseURL,
});
}
private handleError(
e: any,
model: string,
options: CopilotImageOptions = {}
) {
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
} else if (e instanceof AISDKError) {
if (e.message.includes('safety') || e.message.includes('risk')) {
metrics.ai
.counter('chat_text_risk_errors')
.add(1, { model, user: options.user || undefined });
}
return new CopilotProviderSideError({
provider: this.type,
kind: e.name || 'unknown',
message: e.message,
});
} else {
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected openai response',
});
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected openai response',
});
}
override async refreshOnlineModels() {
@@ -389,20 +378,50 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
override getProviderSpecificTools(
toolName: CopilotChatTools,
model: string
_model: string
): [string, Tool?] | undefined {
if (
toolName === 'webSearch' &&
'responses' in this.#instance &&
!this.isReasoningModel(model)
) {
return ['web_search_preview', openai.tools.webSearch({})];
} else if (toolName === 'docEdit') {
if (toolName === 'docEdit') {
return ['doc_edit', undefined];
}
return;
}
private createNativeConfig(): NativeLlmBackendConfig {
const baseUrl = this.config.baseURL || 'https://api.openai.com/v1';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: this.config.apiKey,
};
}
private createNativeAdapter(
tools: ToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream(
this.config.oldApiStyle ? 'openai_chat' : 'openai_responses',
this.createNativeConfig(),
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 { effort: 'medium' };
}
return undefined;
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
@@ -413,33 +432,25 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_calls').add(1, { model: model.id });
const [system, msgs] = await chatToGPTMessage(messages);
const modelInstance =
'responses' in this.#instance
? this.#instance.responses(model.id)
: this.#instance(model.id);
const { text } = await generateText({
model: modelInstance,
system,
messages: msgs,
temperature: options.temperature ?? 0,
maxOutputTokens: options.maxTokens ?? 4096,
providerOptions: {
openai: this.getOpenAIOptions(options, model.id),
},
tools: await this.getTools(options, model.id),
stopWhen: stepCountIs(this.MAX_STEPS),
abortSignal: options.signal,
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,
include: options.webSearch ? ['citations'] : undefined,
reasoning: this.getReasoning(options, model.id),
middleware,
});
return text.trim();
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
return await adapter.text(request, options.signal);
} catch (e: any) {
metrics.ai.counter('chat_text_errors').add(1, { model: model.id });
throw this.handleError(e, model.id, options);
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -456,38 +467,29 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_stream_calls').add(1, { model: model.id });
const fullStream = await this.getFullStream(model, messages, options);
const citationParser = new CitationParser();
const textParser = new TextStreamParser();
for await (const chunk of fullStream) {
switch (chunk.type) {
case 'text-delta': {
let result = textParser.parse(chunk);
result = citationParser.parse(result);
yield result;
break;
}
case 'finish': {
const footnotes = textParser.end();
const result =
citationParser.end() + (footnotes.length ? '\n' + footnotes : '');
yield result;
break;
}
default: {
yield textParser.parse(chunk);
break;
}
}
if (options.signal?.aborted) {
await fullStream.cancel();
break;
}
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,
include: options.webSearch ? ['citations'] : undefined,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
for await (const chunk of adapter.streamText(request, options.signal)) {
yield chunk;
}
} catch (e: any) {
metrics.ai.counter('chat_text_stream_errors').add(1, { model: model.id });
throw this.handleError(e, model.id, options);
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -503,24 +505,27 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
try {
metrics.ai
.counter('chat_object_stream_calls')
.add(1, { model: model.id });
const fullStream = await this.getFullStream(model, messages, options);
const parser = new StreamObjectParser();
for await (const chunk of fullStream) {
const result = parser.parse(chunk);
if (result) {
yield result;
}
if (options.signal?.aborted) {
await fullStream.cancel();
break;
}
.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,
include: options.webSearch ? ['citations'] : undefined,
reasoning: this.getReasoning(options, model.id),
middleware,
});
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
for await (const chunk of adapter.streamObject(request, options.signal)) {
yield chunk;
}
} catch (e: any) {
metrics.ai
.counter('chat_object_stream_errors')
.add(1, { model: model.id });
throw this.handleError(e, model.id, options);
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -535,35 +540,27 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
try {
metrics.ai.counter('chat_text_calls').add(1, { model: model.id });
const [system, msgs, schema] = await chatToGPTMessage(messages);
const tools = await this.getTools(options, model.id);
const middleware = this.getActiveProviderMiddleware();
const { request, schema } = await buildNativeRequest({
model: model.id,
messages,
options,
tools,
reasoning: this.getReasoning(options, model.id),
middleware,
});
if (!schema) {
throw new CopilotPromptInvalid('Schema is required');
}
const modelInstance =
'responses' in this.#instance
? this.#instance.responses(model.id)
: this.#instance(model.id);
const { object } = await generateObject({
model: modelInstance,
system,
messages: msgs,
temperature: options.temperature ?? 0,
maxOutputTokens: options.maxTokens ?? 4096,
maxRetries: options.maxRetries ?? 3,
schema,
providerOptions: {
openai: options.user ? { user: options.user } : {},
},
abortSignal: options.signal,
});
return JSON.stringify(object);
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
const text = await adapter.text(request, options.signal);
const parsed = JSON.parse(text);
const validated = schema.parse(parsed);
return JSON.stringify(validated);
} catch (e: any) {
metrics.ai.counter('chat_text_errors').add(1, { model: model.id });
throw this.handleError(e, model.id, options);
throw this.handleError(e);
}
}
@@ -575,36 +572,32 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
await this.checkParams({ messages: [], cond: fullCond, options });
const model = this.selectModel(fullCond);
// get the log probability of "yes"/"no"
const instance =
'chat' in this.#instance
? this.#instance.chat(model.id)
: this.#instance(model.id);
const scores = await Promise.all(
chunkMessages.map(async messages => {
const [system, msgs] = await chatToGPTMessage(messages);
const result = await generateText({
model: instance,
system,
messages: msgs,
temperature: 0,
maxOutputTokens: 16,
providerOptions: {
openai: {
...this.getOpenAIOptions(options, model.id),
logprobs: 16,
},
const response = await this.requestOpenAIJson(
'/chat/completions',
{
model: model.id,
messages: this.toOpenAIChatMessages(system, msgs),
temperature: 0,
max_tokens: 16,
logprobs: true,
top_logprobs: 16,
},
abortSignal: options.signal,
});
options.signal
);
const topMap: Record<string, number> = LogProbsSchema.parse(
result.providerMetadata?.openai?.logprobs
)[0].top_logprobs.reduce<Record<string, number>>(
const logprobs = response?.choices?.[0]?.logprobs?.content;
if (!Array.isArray(logprobs) || logprobs.length === 0) {
return 0;
}
const parsedLogprobs = LogProbsSchema.parse(logprobs);
const topMap = parsedLogprobs[0].top_logprobs.reduce(
(acc, { token, logprob }) => ({ ...acc, [token]: logprob }),
{}
{} as Record<string, number>
);
const findLogProb = (token: string): number => {
@@ -634,50 +627,212 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
return scores;
}
private async getFullStream(
model: CopilotProviderModel,
messages: PromptMessage[],
options: CopilotChatOptions = {}
) {
const [system, msgs] = await chatToGPTMessage(messages);
const modelInstance =
'responses' in this.#instance
? this.#instance.responses(model.id)
: this.#instance(model.id);
const { fullStream } = streamText({
model: modelInstance,
system,
messages: msgs,
frequencyPenalty: options.frequencyPenalty ?? 0,
presencePenalty: options.presencePenalty ?? 0,
temperature: options.temperature ?? 0,
maxOutputTokens: options.maxTokens ?? 4096,
providerOptions: {
openai: this.getOpenAIOptions(options, model.id),
},
tools: await this.getTools(options, model.id),
stopWhen: stepCountIs(this.MAX_STEPS),
abortSignal: options.signal,
});
return fullStream;
// ====== text to image ======
private buildImageFetchOptions(url: URL) {
const baseOptions = { timeoutMs: 15_000, maxRedirects: 3 } as const;
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 redactUrl(raw: string | URL): string {
try {
const parsed = raw instanceof URL ? raw : new URL(raw);
if (parsed.protocol === 'data:') return 'data:[redacted]';
const segments = parsed.pathname.split('/').filter(Boolean);
const redactedPath =
segments.length <= 2
? parsed.pathname || '/'
: `/${segments[0]}/${segments[1]}/...`;
return `${parsed.origin}${redactedPath}`;
} catch {
return '[invalid-url]';
}
}
private async fetchImage(
url: string,
maxBytes: number,
signal?: AbortSignal
): Promise<{ buffer: Buffer; type: string } | null> {
if (url.startsWith('data:')) {
let response: Response;
try {
response = await fetch(url, { signal });
} catch (error) {
this.logger.warn(
`Skip image attachment data URL due to read failure: ${
error instanceof Error ? error.message : String(error)
}`
);
return null;
}
if (!response.ok) {
this.logger.warn(
`Skip image attachment data URL due to invalid response: ${response.status}`
);
return null;
}
const type =
response.headers.get('content-type') || 'application/octet-stream';
if (!type.startsWith('image/')) {
await response.body?.cancel().catch(() => undefined);
this.logger.warn(
`Skip non-image attachment data URL with content-type ${type}`
);
return null;
}
try {
const buffer = await readResponseBufferWithLimit(response, maxBytes);
return { buffer, type };
} catch (error) {
this.logger.warn(
`Skip image attachment data URL due to read failure/size limit: ${
error instanceof Error ? error.message : String(error)
}`
);
return null;
}
}
let parsed: URL;
try {
parsed = new URL(url);
} catch {
this.logger.warn(
`Skip image attachment with invalid URL: ${this.redactUrl(url)}`
);
return null;
}
const redactedUrl = this.redactUrl(parsed);
if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
this.logger.warn(
`Skip image attachment with unsupported protocol: ${redactedUrl}`
);
return null;
}
let response: Response;
try {
response = await safeFetch(
parsed,
{ method: 'GET', signal },
this.buildImageFetchOptions(parsed)
);
} catch (error) {
this.logger.warn(
`Skip image attachment due to blocked/unreachable URL: ${redactedUrl}, reason: ${
error instanceof Error ? error.message : String(error)
}`
);
return null;
}
if (!response.ok) {
this.logger.warn(
`Skip image attachment fetch failure ${response.status}: ${redactedUrl}`
);
return null;
}
const type =
response.headers.get('content-type') || 'application/octet-stream';
if (!type.startsWith('image/')) {
await response.body?.cancel().catch(() => undefined);
this.logger.warn(
`Skip non-image attachment with content-type ${type}: ${redactedUrl}`
);
return null;
}
const contentLength = Number(response.headers.get('content-length'));
if (Number.isFinite(contentLength) && contentLength > maxBytes) {
await response.body?.cancel().catch(() => undefined);
this.logger.warn(
`Skip oversized image attachment by content-length (${contentLength}): ${redactedUrl}`
);
return null;
}
try {
const buffer = await readResponseBufferWithLimit(response, maxBytes);
return { buffer, type };
} catch (error) {
this.logger.warn(
`Skip image attachment due to read failure/size limit: ${redactedUrl}, reason: ${
error instanceof Error ? error.message : String(error)
}`
);
return null;
}
}
// ====== text to image ======
private async *generateImageWithAttachments(
model: string,
prompt: string,
attachments: NonNullable<PromptMessage['attachments']>
attachments: NonNullable<PromptMessage['attachments']>,
signal?: AbortSignal
): AsyncGenerator<string> {
const form = new FormData();
const outputFormat = 'webp';
const maxBytes = 10 * OneMB;
form.set('model', model);
form.set('prompt', prompt);
form.set('output_format', 'webp');
form.set('output_format', outputFormat);
for (const [idx, entry] of attachments.entries()) {
const url = typeof entry === 'string' ? entry : entry.attachment;
try {
const { buffer, type } = await fetchBuffer(url, 10 * OneMB, 'image/');
const file = new File([buffer], `${idx}.png`, { type });
const attachment = await this.fetchImage(url, maxBytes, signal);
if (!attachment) continue;
const { buffer, type } = attachment;
const extension = type.split(';')[0].split('/')[1] || 'png';
const file = new File([buffer], `${idx}.${extension}`, { type });
form.append('image[]', file);
} catch {
continue;
@@ -703,18 +858,24 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
const json = await res.json();
const imageResponse = ImageResponseSchema.safeParse(json);
if (imageResponse.success) {
const data = imageResponse.data;
if ('error' in data) {
throw new Error(data.error.message);
} else {
for (const image of data.data) {
yield `data:image/webp;base64,${image.b64_json}`;
}
}
} else {
if (!imageResponse.success) {
throw new Error(imageResponse.error.message);
}
const data = imageResponse.data;
if ('error' in data) {
throw new Error(data.error.message);
}
const images = normalizeImageResponseData(
data.data,
normalizeImageFormatToMime(outputFormat)
);
if (!images.length) {
throw new Error('No images returned from OpenAI');
}
for (const image of images) {
yield image;
}
}
override async *streamImages(
@@ -726,13 +887,6 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
await this.checkParams({ messages, cond: fullCond, options });
const model = this.selectModel(fullCond);
if (!('image' in this.#instance)) {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'image',
});
}
metrics.ai
.counter('generate_images_stream_calls')
.add(1, { model: model.id });
@@ -742,22 +896,27 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
try {
if (attachments && attachments.length > 0) {
yield* this.generateImageWithAttachments(model.id, prompt, attachments);
} else {
const modelInstance = this.#instance.image(model.id);
const result = await generateImage({
model: modelInstance,
yield* this.generateImageWithAttachments(
model.id,
prompt,
providerOptions: {
openai: {
quality: options.quality || null,
},
},
});
const imageUrls = result.images.map(
image => `data:image/png;base64,${image.base64}`
attachments,
options.signal
);
} else {
const response = await this.requestOpenAIJson('/images/generations', {
model: model.id,
prompt,
...(options.quality ? { quality: options.quality } : {}),
});
const imageResponse = ImageResponseSchema.parse(response);
if ('error' in imageResponse) {
throw new Error(imageResponse.error.message);
}
const imageUrls = normalizeImageResponseData(imageResponse.data);
if (!imageUrls.length) {
throw new Error('No images returned from OpenAI');
}
for (const imageUrl of imageUrls) {
yield imageUrl;
@@ -769,7 +928,7 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
return;
} catch (e: any) {
metrics.ai.counter('generate_images_errors').add(1, { model: model.id });
throw this.handleError(e, model.id, options);
throw this.handleError(e);
}
}
@@ -783,51 +942,85 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
await this.checkParams({ embeddings: messages, cond: fullCond, options });
const model = this.selectModel(fullCond);
if (!('embedding' in this.#instance)) {
throw new CopilotProviderNotSupported({
provider: this.type,
kind: 'embedding',
});
}
try {
metrics.ai
.counter('generate_embedding_calls')
.add(1, { model: model.id });
const modelInstance = this.#instance.embedding(model.id);
const { embeddings } = await embedMany({
model: modelInstance,
values: messages,
providerOptions: {
openai: {
dimensions: options.dimensions || DEFAULT_DIMENSIONS,
},
},
const response = await this.requestOpenAIJson('/embeddings', {
model: model.id,
input: messages,
dimensions: options.dimensions || DEFAULT_DIMENSIONS,
});
return embeddings.filter(v => v && Array.isArray(v));
const data = Array.isArray(response?.data) ? response.data : [];
return data
.map((item: any) => item?.embedding)
.filter((embedding: unknown) => Array.isArray(embedding)) as number[][];
} catch (e: any) {
metrics.ai
.counter('generate_embedding_errors')
.add(1, { model: model.id });
throw this.handleError(e, model.id, options);
throw this.handleError(e);
}
}
private getOpenAIOptions(options: CopilotChatOptions, model: string) {
const result: OpenAIResponsesProviderOptions = {};
if (options?.reasoning && this.isReasoningModel(model)) {
result.reasoningEffort = 'medium';
result.reasoningSummary = 'detailed';
private toOpenAIChatMessages(
system: string | undefined,
messages: Awaited<ReturnType<typeof chatToGPTMessage>>[1]
) {
const result: Array<{ role: string; content: string }> = [];
if (system) {
result.push({ role: 'system', content: system });
}
if (options?.user) {
result.user = options.user;
for (const message of messages) {
if (typeof message.content === 'string') {
result.push({ role: message.role, content: message.content });
continue;
}
const text = message.content
.filter(
part =>
part &&
typeof part === 'object' &&
'type' in part &&
part.type === 'text' &&
'text' in part
)
.map(part => String((part as { text: string }).text))
.join('\n');
result.push({ role: message.role, content: text || '[no content]' });
}
return result;
}
private async requestOpenAIJson(
path: string,
body: Record<string, unknown>,
signal?: AbortSignal
): Promise<any> {
const baseUrl = this.config.baseURL || 'https://api.openai.com/v1';
const response = await fetch(`${baseUrl}${path}`, {
method: 'POST',
headers: {
Authorization: `Bearer ${this.config.apiKey}`,
'Content-Type': 'application/json',
},
body: JSON.stringify(body),
signal,
});
if (!response.ok) {
throw new Error(
`OpenAI API error ${response.status}: ${await response.text()}`
);
}
return await response.json();
}
private isReasoningModel(model: string) {
// o series reasoning models
return model.startsWith('o') || model.startsWith('gpt-5');
@@ -1,11 +1,13 @@
import {
createPerplexity,
type PerplexityProvider as VercelPerplexityProvider,
} from '@ai-sdk/perplexity';
import { generateText, streamText } from 'ai';
import { z } from 'zod';
import type { ToolSet } from 'ai';
import { CopilotProviderSideError, metrics } from '../../../base';
import {
llmDispatchStream,
type NativeLlmBackendConfig,
type NativeLlmRequest,
} from '../../../native';
import type { NodeTextMiddleware } from '../config';
import { buildNativeRequest, NativeProviderAdapter } from './native';
import { CopilotProvider } from './provider';
import {
CopilotChatOptions,
@@ -15,34 +17,12 @@ import {
ModelOutputType,
PromptMessage,
} from './types';
import { chatToGPTMessage, CitationParser } from './utils';
export type PerplexityConfig = {
apiKey: string;
endpoint?: string;
};
const PerplexityErrorSchema = z.union([
z.object({
detail: z.array(
z.object({
loc: z.array(z.string()),
msg: z.string(),
type: z.string(),
})
),
}),
z.object({
error: z.object({
message: z.string(),
type: z.string(),
code: z.number(),
}),
}),
]);
type PerplexityError = z.infer<typeof PerplexityErrorSchema>;
export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
readonly type = CopilotProviderType.Perplexity;
@@ -90,18 +70,38 @@ export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
},
];
#instance!: VercelPerplexityProvider;
override configured(): boolean {
return !!this.config.apiKey;
}
protected override setup() {
super.setup();
this.#instance = createPerplexity({
apiKey: this.config.apiKey,
baseURL: this.config.endpoint,
});
}
private createNativeConfig(): NativeLlmBackendConfig {
const baseUrl = this.config.endpoint || 'https://api.perplexity.ai';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: this.config.apiKey,
};
}
private createNativeAdapter(
tools: ToolSet,
nodeTextMiddleware?: NodeTextMiddleware[]
) {
return new NativeProviderAdapter(
(request: NativeLlmRequest, signal?: AbortSignal) =>
llmDispatchStream(
'openai_chat',
this.createNativeConfig(),
request,
signal
),
tools,
this.MAX_STEPS,
{ nodeTextMiddleware }
);
}
async text(
@@ -114,32 +114,25 @@ export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_calls').add(1, { model: model.id });
metrics.ai.counter('chat_text_calls').add(1, this.metricLabels(model.id));
const [system, msgs] = await chatToGPTMessage(messages, false);
const modelInstance = this.#instance(model.id);
const { text, sources } = await generateText({
model: modelInstance,
system,
messages: msgs,
temperature: options.temperature ?? 0,
maxOutputTokens: options.maxTokens ?? 4096,
abortSignal: options.signal,
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 parser = new CitationParser();
for (const source of sources.filter(s => s.sourceType === 'url')) {
parser.push(source.url);
}
let result = text.replaceAll(/<\/?think>\n/g, '\n---\n');
result = parser.parse(result);
result += parser.end();
return result;
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
return await adapter.text(request, options.signal);
} catch (e: any) {
metrics.ai.counter('chat_text_errors').add(1, { model: model.id });
metrics.ai
.counter('chat_text_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
@@ -154,79 +147,33 @@ export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
const model = this.selectModel(fullCond);
try {
metrics.ai.counter('chat_text_stream_calls').add(1, { model: model.id });
metrics.ai
.counter('chat_text_stream_calls')
.add(1, this.metricLabels(model.id));
const [system, msgs] = await chatToGPTMessage(messages, false);
const modelInstance = this.#instance(model.id);
const stream = streamText({
model: modelInstance,
system,
messages: msgs,
temperature: options.temperature ?? 0,
maxOutputTokens: options.maxTokens ?? 4096,
abortSignal: options.signal,
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 parser = new CitationParser();
for await (const chunk of stream.fullStream) {
switch (chunk.type) {
case 'source': {
if (chunk.sourceType === 'url') {
parser.push(chunk.url);
}
break;
}
case 'text-delta': {
const text = chunk.text.replaceAll(/<\/?think>\n?/g, '\n---\n');
const result = parser.parse(text);
yield result;
break;
}
case 'finish-step': {
const result = parser.end();
yield result;
break;
}
case 'error': {
const json =
typeof chunk.error === 'string'
? JSON.parse(chunk.error)
: chunk.error;
if (json && typeof json === 'object') {
const data = PerplexityErrorSchema.parse(json);
if ('detail' in data || 'error' in data) {
throw this.convertError(data);
}
}
}
}
const adapter = this.createNativeAdapter(tools, middleware.node?.text);
for await (const chunk of adapter.streamText(request, options.signal)) {
yield chunk;
}
} catch (e) {
metrics.ai.counter('chat_text_stream_errors').add(1, { model: model.id });
throw e;
} catch (e: any) {
metrics.ai
.counter('chat_text_stream_errors')
.add(1, this.metricLabels(model.id));
throw this.handleError(e);
}
}
private convertError(e: PerplexityError) {
function getErrMessage(e: PerplexityError) {
let err = 'Unexpected perplexity response';
if ('detail' in e) {
err = e.detail[0].msg || err;
} else if ('error' in e) {
err = e.error.message || err;
}
return err;
}
throw new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: getErrMessage(e),
});
}
private handleError(e: any) {
if (e instanceof CopilotProviderSideError) {
return e;
@@ -0,0 +1,98 @@
import type { ProviderMiddlewareConfig } from '../config';
import { CopilotProviderType } from './types';
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'],
},
},
[CopilotProviderType.Anthropic]: {
rust: {
request: ['normalize_messages', 'tool_schema_rewrite'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
},
[CopilotProviderType.AnthropicVertex]: {
rust: {
request: ['normalize_messages', 'tool_schema_rewrite'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
},
[CopilotProviderType.Morph]: {
rust: {
request: ['clamp_max_tokens'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
},
[CopilotProviderType.Perplexity]: {
rust: {
request: ['clamp_max_tokens'],
stream: ['stream_event_normalize', 'citation_indexing'],
},
node: {
text: ['citation_footnote', 'callout'],
},
},
[CopilotProviderType.Gemini]: {
node: {
text: ['callout'],
},
},
[CopilotProviderType.GeminiVertex]: {
node: {
text: ['callout'],
},
},
[CopilotProviderType.FAL]: {},
};
function unique<T>(items: T[]) {
return [...new Set(items)];
}
function mergeArray<T>(base: T[] | undefined, override: T[] | undefined) {
if (!base?.length && !override?.length) {
return undefined;
}
return unique([...(base ?? []), ...(override ?? [])]);
}
export function mergeProviderMiddleware(
defaults: ProviderMiddlewareConfig,
override?: ProviderMiddlewareConfig
): ProviderMiddlewareConfig {
return {
rust: {
request: mergeArray(defaults.rust?.request, override?.rust?.request),
stream: mergeArray(defaults.rust?.stream, override?.rust?.stream),
},
node: {
text: mergeArray(defaults.node?.text, override?.node?.text),
},
};
}
export function resolveProviderMiddleware(
type: CopilotProviderType,
override?: ProviderMiddlewareConfig
): ProviderMiddlewareConfig {
const defaults = DEFAULT_MIDDLEWARE_BY_TYPE[type] ?? {};
return mergeProviderMiddleware(defaults, override);
}
@@ -0,0 +1,273 @@
import type {
CopilotProviderConfigMap,
CopilotProviderDefaults,
CopilotProviderProfile,
ProviderMiddlewareConfig,
} from '../config';
import { resolveProviderMiddleware } from './provider-middleware';
import { CopilotProviderType, type ModelOutputType } from './types';
const PROVIDER_ID_PATTERN = /^[a-zA-Z0-9-_]+$/;
const LEGACY_PROVIDER_ORDER: CopilotProviderType[] = [
CopilotProviderType.OpenAI,
CopilotProviderType.FAL,
CopilotProviderType.Gemini,
CopilotProviderType.GeminiVertex,
CopilotProviderType.Perplexity,
CopilotProviderType.Anthropic,
CopilotProviderType.AnthropicVertex,
CopilotProviderType.Morph,
];
const LEGACY_PROVIDER_PRIORITY = LEGACY_PROVIDER_ORDER.reduce(
(acc, type, index) => {
acc[type] = LEGACY_PROVIDER_ORDER.length - index;
return acc;
},
{} as Record<CopilotProviderType, number>
);
type LegacyProvidersConfig = Partial<
Record<CopilotProviderType, CopilotProviderConfigMap[CopilotProviderType]>
>;
export type CopilotProvidersConfigInput = LegacyProvidersConfig & {
profiles?: CopilotProviderProfile[] | null;
defaults?: CopilotProviderDefaults | null;
};
export type NormalizedCopilotProviderProfile = Omit<
CopilotProviderProfile,
'enabled' | 'priority' | 'middleware'
> & {
enabled: boolean;
priority: number;
middleware: ProviderMiddlewareConfig;
};
export type CopilotProviderRegistry = {
profiles: Map<string, NormalizedCopilotProviderProfile>;
defaults: CopilotProviderDefaults;
order: string[];
byType: Map<CopilotProviderType, string[]>;
};
export type ResolveModelResult = {
rawModelId?: string;
modelId?: string;
explicitProviderId?: string;
candidateProviderIds: string[];
};
type ResolveModelOptions = {
registry: CopilotProviderRegistry;
modelId?: string;
outputType?: ModelOutputType;
availableProviderIds?: Iterable<string>;
preferredProviderIds?: Iterable<string>;
};
function unique<T>(list: T[]): T[] {
return [...new Set(list)];
}
function asArray<T>(iter?: Iterable<T>): T[] {
return iter ? Array.from(iter) : [];
}
function parseModelPrefix(
registry: CopilotProviderRegistry,
modelId: string
): { providerId: string; modelId?: string } | null {
const index = modelId.indexOf('/');
if (index <= 0) {
return null;
}
const providerId = modelId.slice(0, index);
if (!registry.profiles.has(providerId)) {
return null;
}
const model = modelId.slice(index + 1);
return { providerId, modelId: model || undefined };
}
function normalizeProfile(
profile: CopilotProviderProfile
): NormalizedCopilotProviderProfile {
return {
...profile,
enabled: profile.enabled !== false,
priority: profile.priority ?? 0,
middleware: resolveProviderMiddleware(profile.type, profile.middleware),
};
}
function toLegacyProfiles(
config: CopilotProvidersConfigInput
): CopilotProviderProfile[] {
const legacyProfiles: CopilotProviderProfile[] = [];
for (const type of LEGACY_PROVIDER_ORDER) {
const legacyConfig = config[type];
if (!legacyConfig) {
continue;
}
legacyProfiles.push({
id: `${type}-default`,
type,
priority: LEGACY_PROVIDER_PRIORITY[type],
config: legacyConfig,
} as CopilotProviderProfile);
}
return legacyProfiles;
}
function mergeProfiles(
explicitProfiles: CopilotProviderProfile[],
legacyProfiles: CopilotProviderProfile[]
): CopilotProviderProfile[] {
const profiles = new Map<string, CopilotProviderProfile>();
for (const profile of explicitProfiles) {
if (!PROVIDER_ID_PATTERN.test(profile.id)) {
throw new Error(`Invalid copilot provider profile id: ${profile.id}`);
}
if (profiles.has(profile.id)) {
throw new Error(`Duplicated copilot provider profile id: ${profile.id}`);
}
profiles.set(profile.id, profile);
}
for (const profile of legacyProfiles) {
if (!profiles.has(profile.id)) {
profiles.set(profile.id, profile);
}
}
return Array.from(profiles.values());
}
function sortProfiles(profiles: NormalizedCopilotProviderProfile[]) {
return profiles.toSorted((a, b) => {
if (a.priority !== b.priority) {
return b.priority - a.priority;
}
return a.id.localeCompare(b.id);
});
}
function assertDefaults(
defaults: CopilotProviderDefaults,
profiles: Map<string, NormalizedCopilotProviderProfile>
) {
for (const providerId of Object.values(defaults)) {
if (!providerId) {
continue;
}
if (!profiles.has(providerId)) {
throw new Error(
`Copilot provider defaults references unknown providerId: ${providerId}`
);
}
}
}
export function buildProviderRegistry(
config: CopilotProvidersConfigInput
): CopilotProviderRegistry {
const explicitProfiles = config.profiles ?? [];
const legacyProfiles = toLegacyProfiles(config);
const mergedProfiles = mergeProfiles(explicitProfiles, legacyProfiles)
.map(normalizeProfile)
.filter(profile => profile.enabled);
const sortedProfiles = sortProfiles(mergedProfiles);
const profiles = new Map(
sortedProfiles.map(profile => [profile.id, profile] as const)
);
const defaults = config.defaults ?? {};
assertDefaults(defaults, profiles);
const order = sortedProfiles.map(profile => profile.id);
const byType = new Map<CopilotProviderType, string[]>();
for (const profile of sortedProfiles) {
const ids = byType.get(profile.type) ?? [];
ids.push(profile.id);
byType.set(profile.type, ids);
}
return { profiles, defaults, order, byType };
}
export function resolveModel({
registry,
modelId,
outputType,
availableProviderIds,
preferredProviderIds,
}: ResolveModelOptions): ResolveModelResult {
const available = new Set(asArray(availableProviderIds));
const preferred = new Set(asArray(preferredProviderIds));
const hasAvailableFilter = available.size > 0;
const hasPreferredFilter = preferred.size > 0;
const isAllowed = (providerId: string) => {
const profile = registry.profiles.get(providerId);
if (!profile?.enabled) {
return false;
}
if (hasAvailableFilter && !available.has(providerId)) {
return false;
}
if (hasPreferredFilter && !preferred.has(providerId)) {
return false;
}
return true;
};
const prefixed = modelId ? parseModelPrefix(registry, modelId) : null;
if (prefixed) {
return {
rawModelId: modelId,
modelId: prefixed.modelId,
explicitProviderId: prefixed.providerId,
candidateProviderIds: isAllowed(prefixed.providerId)
? [prefixed.providerId]
: [],
};
}
const fallbackOrder = [
...(outputType ? [registry.defaults[outputType]] : []),
registry.defaults.fallback,
...registry.order,
].filter((id): id is string => !!id);
return {
rawModelId: modelId,
modelId,
candidateProviderIds: unique(
fallbackOrder.filter(providerId => isAllowed(providerId))
),
};
}
export function stripProviderPrefix(
registry: CopilotProviderRegistry,
providerId: string,
modelId?: string
) {
if (!modelId) {
return modelId;
}
const prefixed = parseModelPrefix(registry, modelId);
if (!prefixed) {
return modelId;
}
if (prefixed.providerId !== providerId) {
return modelId;
}
return prefixed.modelId;
}
@@ -1,3 +1,5 @@
import { AsyncLocalStorage } from 'node:async_hooks';
import { Inject, Injectable, Logger } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import { Tool, ToolSet } from 'ai';
@@ -13,6 +15,7 @@ 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 {
@@ -40,6 +43,8 @@ import {
createSectionEditTool,
} from '../tools';
import { CopilotProviderFactory } from './factory';
import { resolveProviderMiddleware } from './provider-middleware';
import { buildProviderRegistry } from './provider-registry';
import {
type CopilotChatOptions,
CopilotChatTools,
@@ -58,11 +63,14 @@ import {
StreamObject,
} from './types';
const providerProfileContext = new AsyncLocalStorage<string>();
@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;
@@ -70,8 +78,39 @@ export abstract class CopilotProvider<C = any> {
@Inject() protected readonly AFFiNEConfig!: Config;
@Inject() protected readonly factory!: CopilotProviderFactory;
@Inject() protected readonly moduleRef!: ModuleRef;
readonly #registeredProviderIds = new Set<string>();
runWithProfile<T>(providerId: string, callback: () => T): T {
return providerProfileContext.run(providerId, callback);
}
protected getActiveProviderId() {
return providerProfileContext.getStore() ?? `${this.type}-default`;
}
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 metricLabels(
model: string,
labels: Record<string, string | number | boolean | undefined> = {}
) {
const providerId = this.getActiveProviderId();
return { model, providerId, ...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;
}
return this.AFFiNEConfig.copilot.providers[this.type] as C;
}
@@ -88,15 +127,37 @@ export abstract class CopilotProvider<C = any> {
}
protected setup() {
if (this.configured()) {
this.factory.register(this);
if (env.selfhosted) {
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)
);
}
} else {
this.factory.unregister(this);
});
}
}
@@ -91,7 +91,9 @@ export async function chatToGPTMessage(
// so we need to use base64 encoded attachments instead
useBase64Attachment: boolean = false
): Promise<[string | undefined, ChatMessage[], ZodType?]> {
const system = messages[0]?.role === 'system' ? messages.shift() : undefined;
const hasSystem = messages[0]?.role === 'system';
const system = hasSystem ? messages[0] : undefined;
const normalizedMessages = hasSystem ? messages.slice(1) : messages;
const schema =
system?.params?.schema && system.params.schema instanceof ZodType
? system.params.schema
@@ -99,7 +101,7 @@ export async function chatToGPTMessage(
// filter redundant fields
const msgs: ChatMessage[] = [];
for (let { role, content, attachments, params } of messages.filter(
for (let { role, content, attachments, params } of normalizedMessages.filter(
m => m.role !== 'system'
)) {
content = content.trim();
@@ -406,6 +408,34 @@ export class CitationParser {
}
}
export type CitationIndexedEvent = {
type: 'citation';
index: number;
url: string;
};
export class CitationFootnoteFormatter {
private readonly citations = new Map<number, string>();
public consume(event: CitationIndexedEvent) {
if (event.type !== 'citation') {
return '';
}
this.citations.set(event.index, event.url);
return '';
}
public end() {
const footnotes = Array.from(this.citations.entries())
.sort((a, b) => a[0] - b[0])
.map(
([index, citation]) =>
`[^${index}]: {"type":"url","url":"${encodeURIComponent(citation)}"}`
);
return footnotes.join('\n');
}
}
type ChunkType = TextStreamPart<CustomAITools>['type'];
export function toError(error: unknown): Error {
@@ -703,21 +733,39 @@ export const VertexModelListSchema = z.object({
),
});
function normalizeUrl(baseURL?: string) {
if (!baseURL?.trim()) {
return undefined;
}
try {
const url = new URL(baseURL);
const serialized = url.toString();
if (serialized.endsWith('/')) return serialized.slice(0, -1);
return serialized;
} catch {
return undefined;
}
}
export function getVertexAnthropicBaseUrl(
options: GoogleVertexAnthropicProviderSettings
) {
const normalizedBaseUrl = normalizeUrl(options.baseURL);
if (normalizedBaseUrl) return normalizedBaseUrl;
const { location, project } = options;
if (!location || !project) return undefined;
return `https://${location}-aiplatform.googleapis.com/v1/projects/${project}/locations/${location}/publishers/anthropic`;
}
export async function getGoogleAuth(
options: GoogleVertexAnthropicProviderSettings | GoogleVertexProviderSettings,
publisher: 'anthropic' | 'google'
) {
function getBaseUrl() {
const { baseURL, location } = options;
if (baseURL?.trim()) {
try {
const url = new URL(baseURL);
if (url.pathname.endsWith('/')) {
url.pathname = url.pathname.slice(0, -1);
}
return url.toString();
} catch {}
} else if (location) {
const normalizedBaseUrl = normalizeUrl(options.baseURL);
if (normalizedBaseUrl) return normalizedBaseUrl;
const { location } = options;
if (location) {
return `https://${location}-aiplatform.googleapis.com/v1beta1/publishers/${publisher}`;
}
return undefined;
@@ -4,7 +4,6 @@ import { BadRequestException, NotFoundException } from '@nestjs/common';
import {
Args,
Field,
Float,
ID,
InputType,
Mutation,
@@ -15,7 +14,6 @@ import {
ResolveField,
Resolver,
} from '@nestjs/graphql';
import { AiPromptRole } from '@prisma/client';
import { GraphQLJSON, SafeIntResolver } from 'graphql-scalars';
import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
@@ -313,57 +311,6 @@ class CopilotQuotaType {
used!: number;
}
registerEnumType(AiPromptRole, {
name: 'CopilotPromptMessageRole',
});
@InputType('CopilotPromptConfigInput')
@ObjectType()
class CopilotPromptConfigType {
@Field(() => Float, { nullable: true })
frequencyPenalty!: number | null;
@Field(() => Float, { nullable: true })
presencePenalty!: number | null;
@Field(() => Float, { nullable: true })
temperature!: number | null;
@Field(() => Float, { nullable: true })
topP!: number | null;
}
@InputType('CopilotPromptMessageInput')
@ObjectType()
class CopilotPromptMessageType {
@Field(() => AiPromptRole)
role!: AiPromptRole;
@Field(() => String)
content!: string;
@Field(() => GraphQLJSON, { nullable: true })
params!: Record<string, string> | null;
}
@ObjectType()
class CopilotPromptType {
@Field(() => String)
name!: string;
@Field(() => String)
model!: string;
@Field(() => String, { nullable: true })
action!: string | null;
@Field(() => CopilotPromptConfigType, { nullable: true })
config!: CopilotPromptConfigType | null;
@Field(() => [CopilotPromptMessageType])
messages!: CopilotPromptMessageType[];
}
@ObjectType()
class CopilotModelType {
@Field(() => String)
@@ -638,13 +585,8 @@ export class CopilotResolver {
);
}
@Mutation(() => String, {
description: 'Create a chat session',
})
@CallMetric('ai', 'chat_session_create')
async createCopilotSession(
@CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => CreateChatSessionInput })
private async createCopilotSessionInternal(
user: CurrentUser,
options: CreateChatSessionInput
): Promise<string> {
// permission check based on session type
@@ -666,6 +608,42 @@ export class CopilotResolver {
});
}
@Mutation(() => String, {
description: 'Create a chat session',
deprecationReason: 'use `createCopilotSessionWithHistory` instead',
})
@CallMetric('ai', 'chat_session_create')
async createCopilotSession(
@CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => CreateChatSessionInput })
options: CreateChatSessionInput
): Promise<string> {
return await this.createCopilotSessionInternal(user, options);
}
@Mutation(() => CopilotHistoriesType, {
description: 'Create a chat session and return full session payload',
})
@CallMetric('ai', 'chat_session_create_with_history')
async createCopilotSessionWithHistory(
@CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => CreateChatSessionInput })
options: CreateChatSessionInput
): Promise<CopilotHistoriesType> {
const sessionId = await this.createCopilotSessionInternal(user, options);
const session = await this.chatSession.getSessionInfo(sessionId);
if (!session) {
throw new NotFoundException('Session not found');
}
return {
...session,
messages: session.messages.map(message => ({
...message,
id: message.id,
})) as ChatMessageType[],
};
}
@Mutation(() => String, {
description: 'Update a chat session',
})
@@ -939,31 +917,10 @@ export class UserCopilotResolver {
}
}
@InputType()
class CreateCopilotPromptInput {
@Field(() => String)
name!: string;
@Field(() => String)
model!: string;
@Field(() => String, { nullable: true })
action!: string | null;
@Field(() => CopilotPromptConfigType, { nullable: true })
config!: CopilotPromptConfigType | null;
@Field(() => [CopilotPromptMessageType])
messages!: CopilotPromptMessageType[];
}
@Admin()
@Resolver(() => String)
export class PromptsManagementResolver {
constructor(
private readonly cron: CopilotCronJobs,
private readonly promptService: PromptService
) {}
constructor(private readonly cron: CopilotCronJobs) {}
@Mutation(() => Boolean, {
description: 'Trigger generate missing titles cron job',
@@ -980,48 +937,4 @@ export class PromptsManagementResolver {
await this.cron.triggerCleanupTrashedDocEmbeddings();
return true;
}
@Query(() => [CopilotPromptType], {
description: 'List all copilot prompts',
})
async listCopilotPrompts() {
const prompts = await this.promptService.list();
return prompts.filter(
p =>
p.messages.length > 0 &&
// ignore internal prompts
!p.name.startsWith('workflow:') &&
!p.name.startsWith('debug:') &&
!p.name.startsWith('chat:') &&
!p.name.startsWith('action:')
);
}
@Mutation(() => CopilotPromptType, {
description: 'Create a copilot prompt',
})
async createCopilotPrompt(
@Args({ type: () => CreateCopilotPromptInput, name: 'input' })
input: CreateCopilotPromptInput
) {
await this.promptService.set(
input.name,
input.model,
input.messages,
input.config
);
return this.promptService.get(input.name);
}
@Mutation(() => CopilotPromptType, {
description: 'Update a copilot prompt',
})
async updateCopilotPrompt(
@Args('name') name: string,
@Args('messages', { type: () => [CopilotPromptMessageType] })
messages: CopilotPromptMessageType[]
) {
await this.promptService.update(name, { messages, modified: true });
return this.promptService.get(name);
}
}
@@ -7,6 +7,7 @@ import { AiPromptRole } from '@prisma/client';
import { pick } from 'lodash-es';
import {
Config,
CopilotActionTaken,
CopilotMessageNotFound,
CopilotPromptNotFound,
@@ -31,6 +32,7 @@ import { ChatMessageCache } from './message';
import { ChatPrompt } from './prompt/chat-prompt';
import { PromptService } from './prompt/service';
import { CopilotProviderFactory } from './providers/factory';
import { buildProviderRegistry } from './providers/provider-registry';
import {
ModelOutputType,
type PromptMessage,
@@ -105,10 +107,31 @@ export class ChatSession implements AsyncDisposable {
hasPayment: boolean,
requestedModelId?: string
): Promise<string> {
const config = this.moduleRef.get(Config, { strict: false });
const registry = config
? buildProviderRegistry(config.copilot.providers)
: null;
const defaultModel = this.model;
const normalize = (m?: string) =>
!!m && this.optionalModels.includes(m) ? m : defaultModel;
const isPro = (m?: string) => !!m && this.proModels.includes(m);
const normalizeModel = (modelId?: string) => {
if (!modelId) return modelId;
const separatorIndex = modelId.indexOf('/');
if (separatorIndex <= 0) return modelId;
const providerId = modelId.slice(0, separatorIndex);
if (!registry?.profiles.has(providerId)) return modelId;
return modelId.slice(separatorIndex + 1);
};
const inModelList = (models: string[], modelId?: string) => {
if (!modelId) return false;
return (
models.includes(modelId) ||
models.includes(normalizeModel(modelId) ?? '')
);
};
const normalize = (m?: string) => {
if (inModelList(this.optionalModels, m)) return m;
return defaultModel;
};
const isPro = (m?: string) => inModelList(this.proModels, m);
// try resolve payment subscription service lazily
let paymentEnabled = hasPayment;
@@ -132,10 +155,19 @@ export class ChatSession implements AsyncDisposable {
}
if (paymentEnabled && !isUserAIPro && isPro(requestedModelId)) {
if (!defaultModel) {
throw new CopilotSessionInvalidInput(
'Model is required for AI subscription fallback'
);
}
return defaultModel;
}
return normalize(requestedModelId);
const resolvedModel = normalize(requestedModelId);
if (!resolvedModel) {
throw new CopilotSessionInvalidInput('Model is required');
}
return resolvedModel;
}
push(message: ChatMessage) {
@@ -32,16 +32,22 @@ export const buildBlobContentGetter = (
return;
}
const contextFile = context.files.find(
file => file.blobId === blobId || file.id === blobId
);
const canonicalBlobId = contextFile?.blobId ?? blobId;
const targetFileId = contextFile?.id;
const [file, blob] = await Promise.all([
context?.getFileContent(blobId, chunk),
context?.getBlobContent(blobId, chunk),
targetFileId ? context.getFileContent(targetFileId, chunk) : undefined,
context.getBlobContent(canonicalBlobId, chunk),
]);
const content = file?.trim() || blob?.trim();
if (!content) {
return;
}
if (!content) return;
const info = contextFile
? { fileName: contextFile.name, fileType: contextFile.mimeType }
: {};
return { blobId, chunk, content };
return { blobId: canonicalBlobId, chunk, content, ...info };
};
return getBlobContent;
};
@@ -14,6 +14,7 @@ import type {
import { HTMLRewriter } from 'htmlrewriter';
import {
applyAttachHeaders,
BadRequest,
Cache,
readResponseBufferWithLimit,
@@ -127,15 +128,18 @@ export class WorkerController {
if (buffer.length === 0) {
return resp.status(404).header(getCorsHeaders(origin)).send();
}
return resp
.status(200)
.header({
...getCorsHeaders(origin),
...(origin ? { Vary: 'Origin' } : {}),
'Access-Control-Allow-Methods': 'GET',
'Content-Type': 'image/*',
})
.send(buffer);
resp.header({
...getCorsHeaders(origin),
...(origin ? { Vary: 'Origin' } : {}),
'Access-Control-Allow-Methods': 'GET',
});
applyAttachHeaders(resp, { buffer });
const contentType = resp.getHeader('Content-Type') as string | undefined;
if (contentType?.startsWith('image/')) {
return resp.status(200).send(buffer);
} else {
throw new BadRequest('Invalid content type');
}
}
let response: Response;
@@ -171,39 +175,39 @@ export class WorkerController {
throw new BadRequest('Failed to fetch image');
}
if (response.ok) {
const contentType = response.headers.get('Content-Type');
if (contentType?.startsWith('image/')) {
let buffer: Buffer;
try {
buffer = await readResponseBufferWithLimit(
response,
IMAGE_PROXY_MAX_BYTES
);
} catch (error) {
if (error instanceof ResponseTooLargeError) {
this.logger.warn('Image proxy response too large', {
url: imageURL,
limitBytes: error.data?.limitBytes,
receivedBytes: error.data?.receivedBytes,
});
throw new BadRequest('Response too large');
}
throw error;
let buffer: Buffer;
try {
buffer = await readResponseBufferWithLimit(
response,
IMAGE_PROXY_MAX_BYTES
);
} catch (error) {
if (error instanceof ResponseTooLargeError) {
this.logger.warn('Image proxy response too large', {
url: imageURL,
limitBytes: error.data?.limitBytes,
receivedBytes: error.data?.receivedBytes,
});
throw new BadRequest('Response too large');
}
await this.cache.set(cachedUrl, buffer.toString('base64'), {
ttl: CACHE_TTL,
});
const contentDisposition = response.headers.get('Content-Disposition');
return resp
.status(200)
.header({
...getCorsHeaders(origin),
...(origin ? { Vary: 'Origin' } : {}),
'Access-Control-Allow-Methods': 'GET',
'Content-Type': contentType,
'Content-Disposition': contentDisposition,
})
.send(buffer);
throw error;
}
await this.cache.set(cachedUrl, buffer.toString('base64'), {
ttl: CACHE_TTL,
});
const contentDisposition = response.headers.get('Content-Disposition');
resp.header({
...getCorsHeaders(origin),
...(origin ? { Vary: 'Origin' } : {}),
'Access-Control-Allow-Methods': 'GET',
});
if (contentDisposition) {
resp.setHeader('Content-Disposition', contentDisposition);
}
applyAttachHeaders(resp, { buffer });
const contentType = resp.getHeader('Content-Type') as string | undefined;
if (contentType?.startsWith('image/')) {
return resp.status(200).send(buffer);
} else {
throw new BadRequest('Invalid content type');
}
+4 -58
View File
@@ -607,50 +607,10 @@ type CopilotModelsType {
proModels: [CopilotModelType!]!
}
input CopilotPromptConfigInput {
frequencyPenalty: Float
presencePenalty: Float
temperature: Float
topP: Float
}
type CopilotPromptConfigType {
frequencyPenalty: Float
presencePenalty: Float
temperature: Float
topP: Float
}
input CopilotPromptMessageInput {
content: String!
params: JSON
role: CopilotPromptMessageRole!
}
enum CopilotPromptMessageRole {
assistant
system
user
}
type CopilotPromptMessageType {
content: String!
params: JSON
role: CopilotPromptMessageRole!
}
type CopilotPromptNotFoundDataType {
name: String!
}
type CopilotPromptType {
action: String
config: CopilotPromptConfigType
messages: [CopilotPromptMessageType!]!
model: String!
name: String!
}
type CopilotProviderNotSupportedDataType {
kind: String!
provider: String!
@@ -747,14 +707,6 @@ input CreateCheckoutSessionInput {
variant: SubscriptionVariant
}
input CreateCopilotPromptInput {
action: String
config: CopilotPromptConfigInput
messages: [CopilotPromptMessageInput!]!
model: String!
name: String!
}
input CreateUserInput {
email: String!
name: String
@@ -1551,11 +1503,11 @@ type Mutation {
"""Create a chat message"""
createCopilotMessage(options: CreateChatMessageInput!): String!
"""Create a copilot prompt"""
createCopilotPrompt(input: CreateCopilotPromptInput!): CopilotPromptType!
"""Create a chat session"""
createCopilotSession(options: CreateChatSessionInput!): String!
createCopilotSession(options: CreateChatSessionInput!): String! @deprecated(reason: "use `createCopilotSessionWithHistory` instead")
"""Create a chat session and return full session payload"""
createCopilotSessionWithHistory(options: CreateChatSessionInput!): CopilotHistories!
"""Create a stripe customer portal to manage payment methods"""
createCustomerPortal: String!
@@ -1672,9 +1624,6 @@ type Mutation {
"""Update a comment content"""
updateComment(input: CommentUpdateInput!): Boolean!
"""Update a copilot prompt"""
updateCopilotPrompt(messages: [CopilotPromptMessageInput!]!, name: String!): CopilotPromptType!
"""Update a chat session"""
updateCopilotSession(options: UpdateChatSessionInput!): String!
updateDocDefaultRole(input: UpdateDocDefaultRoleInput!): Boolean!
@@ -1923,9 +1872,6 @@ type Query {
"""get workspace invitation info"""
getInviteInfo(inviteId: String!): InvitationType!
"""List all copilot prompts"""
listCopilotPrompts: [CopilotPromptType!]!
prices: [SubscriptionPrice!]!
"""Get public user by id"""