Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-08-09 05:05:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(server): refactor for byok (#14911)

Browse Source
This commit is contained in:
DarkSky
2026-05-07 04:03:14 +08:00
committed by GitHub
parent 4e169ea5c7
commit eb9cc22502
115 changed files with 10369 additions and 1256 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/src/plugins/copilot/providers/cloudflare.ts
+1 -12
View File
@@ -1,12 +1,11 @@
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import type { CopilotTool } from '../tools';
import { CopilotProvider } from './provider';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { type CopilotChatTools, CopilotProviderType } from './types';
import { CopilotProviderType } from './types';
export type CloudflareWorkersAIConfig = {
apiToken: string;
@@ -25,16 +24,6 @@ export class CloudflareWorkersAIProvider extends CopilotProvider<CloudflareWorke
const config = this.getConfig(execution);
return !!config.apiToken && (!!config.accountId || !!config.baseURL);
}
override getProviderSpecificTools(
toolName: CopilotChatTools,
_model: string
): [string, CopilotTool?] | undefined {
if (toolName === 'docEdit') {
return ['doc_edit', undefined];
}
return;
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
packages/backend/server/src/plugins/copilot/providers/factory.ts
+170 -26
View File
@@ -1,10 +1,14 @@
import { Injectable, Logger } from '@nestjs/common';
import { CopilotQuotaExceeded } from '../../../base';
import { ServerFeature, ServerService } from '../../../core';
import { type CopilotAccessContext, CopilotAccessPolicy } from '../access';
import type { RequiredStructuredOutputContract } from '../runtime/contracts';
import { getProviderRuntimeHost } from '../runtime/provider-runtime-context';
import type { CopilotProvider } from './provider';
import {
buildProviderRegistry,
type CopilotProviderRegistry,
type NormalizedCopilotProviderProfile,
resolveModel,
stripProviderPrefix,
@@ -57,11 +61,18 @@ type RoutePreparationResult = Partial<
>
>;
type EffectiveProviderRegistry = {
byokRegistry: CopilotProviderRegistry;
quotaBackedRegistry: CopilotProviderRegistry;
quotaBackedRoutesAvailable: boolean;
};
@Injectable()
export class CopilotProviderFactory {
constructor(
private readonly server: ServerService,
private readonly registries: CopilotProviderRegistryService
private readonly registries: CopilotProviderRegistryService,
private readonly access: CopilotAccessPolicy
) {}
private readonly logger = new Logger(CopilotProviderFactory.name);
@@ -73,20 +84,84 @@ export class CopilotProviderFactory {
return this.registries.getRegistry();
}
private getPreferredProviderIds(type?: CopilotProviderType) {
private getProviderByProfile(
providerId: string,
profile: NormalizedCopilotProviderProfile
) {
return (
this.#providers.get(providerId) ??
Array.from(this.#providerIdsByType.get(profile.type) ?? [])
.map(id => this.#providers.get(id))
.find((provider): provider is CopilotProvider => !!provider)
);
}
private providerAvailable(
providerId: string,
profile: NormalizedCopilotProviderProfile
) {
return !!this.getProviderByProfile(providerId, profile);
}
private getAvailableProviderIds(registry: CopilotProviderRegistry) {
return Array.from(registry.profiles.entries())
.filter(([providerId, profile]) =>
this.providerAvailable(providerId, profile)
)
.map(([providerId]) => providerId);
}
private getPreferredProviderIds(
registry: CopilotProviderRegistry,
type?: CopilotProviderType
) {
if (!type) return undefined;
return this.#providerIdsByType.get(type);
return registry.byType.get(type)?.filter(providerId => {
const profile = registry.profiles.get(providerId);
return profile ? this.providerAvailable(providerId, profile) : false;
});
}
private normalizeCond(
registry: CopilotProviderRegistry,
providerId: string,
cond: ModelFullConditions
): ModelFullConditions {
const registry = this.getRegistry();
const modelId = stripProviderPrefix(registry, providerId, cond.modelId);
return { ...cond, modelId };
}
private async getEffectiveRegistry(
context: CopilotAccessContext = {}
): Promise<EffectiveProviderRegistry> {
const quotaBackedRegistry = this.getRegistry();
const routeAccess = await this.access.resolveRouteAccess(context);
return {
byokRegistry: buildProviderRegistry({
profiles: routeAccess.byokProfiles,
defaults: {},
}),
quotaBackedRegistry,
quotaBackedRoutesAvailable: routeAccess.quotaBackedRoutesAvailable,
};
}
private getRequestContext(
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions
): CopilotAccessContext {
return {
userId: options?.user,
workspaceId: options?.workspace,
byokLeaseId: options?.byokLeaseId,
featureKind: options?.featureKind,
quotaBackedRoutesAllowed: options?.quotaBackedRoutesAllowed,
};
}
private filterPreparedRoutes(routes: Array<ResolvedCopilotProvider | null>) {
return routes.filter(
(route): route is ResolvedCopilotProvider => route !== null
@@ -113,36 +188,89 @@ export class CopilotProviderFactory {
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
} = {},
context: CopilotAccessContext = {}
): Promise<ResolvedCopilotProvider | null> {
return (await this.resolveRoutes(cond, filter))[0] ?? null;
return (await this.resolveRoutes(cond, filter, context))[0] ?? null;
}
async resolveRoutes(
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
} = {},
context: CopilotAccessContext = {}
): Promise<ResolvedCopilotProvider[]> {
this.logger.debug(
`Resolving copilot provider for output type: ${cond.outputType}`
);
const registry = this.getRegistry();
const { byokRegistry, quotaBackedRegistry, quotaBackedRoutesAvailable } =
await this.getEffectiveRegistry(context);
const byokRoutes = await this.resolveRoutesFromRegistry(
byokRegistry,
cond,
filter
);
const resolved = byokRoutes.length
? byokRoutes
: quotaBackedRoutesAvailable
? await this.resolveRoutesFromRegistry(
quotaBackedRegistry,
cond,
filter
)
: [];
for (const route of resolved) {
this.logger.debug(
`Copilot provider candidate found: ${route.provider.type} (${route.providerId})`
);
}
if (
!resolved.length &&
!quotaBackedRoutesAvailable &&
context.quotaBackedRoutesAllowed !== false
) {
const quotaBackedRoutes = await this.resolveRoutesFromRegistry(
quotaBackedRegistry,
cond,
filter
);
if (quotaBackedRoutes.length) {
throw new CopilotQuotaExceeded();
}
}
return resolved;
}
private async resolveRoutesFromRegistry(
registry: CopilotProviderRegistry,
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const route = resolveModel({
registry,
modelId: cond.modelId,
outputType: cond.outputType,
availableProviderIds: this.#providers.keys(),
preferredProviderIds: this.getPreferredProviderIds(filter.prefer),
availableProviderIds: this.getAvailableProviderIds(registry),
preferredProviderIds: this.getPreferredProviderIds(
registry,
filter.prefer
),
});
const resolved: ResolvedCopilotProvider[] = [];
for (const providerId of route.candidateProviderIds) {
const provider = this.#providers.get(providerId);
const profile = registry.profiles.get(providerId);
const provider = profile
? this.getProviderByProfile(providerId, profile)
: undefined;
if (!provider || !profile) continue;
const normalizedCond = this.normalizeCond(providerId, cond);
const normalizedCond = this.normalizeCond(registry, providerId, cond);
if (
normalizedCond.modelId &&
profile.models?.length &&
@@ -155,9 +283,6 @@ export class CopilotProviderFactory {
const matched = await provider.match(normalizedCond, execution);
if (!matched) continue;
this.logger.debug(
`Copilot provider candidate found: ${provider.type} (${providerId})`
);
resolved.push({
providerId,
provider,
@@ -181,7 +306,11 @@ export class CopilotProviderFactory {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
const routes = await this.resolveRoutes(
cond,
filter,
this.getRequestContext(options)
);
return await this.prepareResolvedRoutes(routes, async route => {
const prepared = await getProviderRuntimeHost(
route.provider
@@ -213,7 +342,11 @@ export class CopilotProviderFactory {
} = {},
responseContract?: RequiredStructuredOutputContract
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
const routes = await this.resolveRoutes(
cond,
filter,
this.getRequestContext(options)
);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedStructured =
(await getProviderRuntimeHost(route.provider).prepare.structured(
@@ -239,10 +372,14 @@ export class CopilotProviderFactory {
input: string | string[],
options: CopilotEmbeddingOptions = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes({
modelId,
outputType: ModelOutputType.Embedding,
});
const routes = await this.resolveRoutes(
{ modelId, outputType: ModelOutputType.Embedding },
{},
{
...this.getRequestContext(options),
featureKind: options?.featureKind ?? 'embedding',
}
);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedEmbedding =
(await getProviderRuntimeHost(route.provider).prepare.embedding(
@@ -267,10 +404,14 @@ export class CopilotProviderFactory {
request: CopilotRerankRequest,
options: CopilotChatOptions = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes({
modelId,
outputType: ModelOutputType.Rerank,
});
const routes = await this.resolveRoutes(
{
modelId,
outputType: ModelOutputType.Rerank,
},
{},
{ ...this.getRequestContext(options), featureKind: 'rerank' }
);
return await this.prepareResolvedRoutes(routes, async route => {
const preparedRerank =
(await getProviderRuntimeHost(route.provider).prepare.rerank(
@@ -298,7 +439,10 @@ export class CopilotProviderFactory {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(cond, filter);
const routes = await this.resolveRoutes(cond, filter, {
...this.getRequestContext(options),
featureKind: options?.featureKind ?? 'image',
});
return await this.prepareResolvedRoutes(routes, async route => {
const preparedImage =
(await getProviderRuntimeHost(route.provider).prepare.image(
packages/backend/server/src/plugins/copilot/providers/index.ts
-1
View File
@@ -8,7 +8,6 @@ export { FalProvider } from './fal';
export { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
export { CopilotProviderLifecycleService } from './lifecycle-service';
export { OpenAIProvider } from './openai';
export { PerplexityProvider } from './perplexity';
export type { CopilotProvider } from './provider';
export { CopilotProviders } from './provider-tokens';
export { CopilotProviderRegistryService } from './registry-service';
packages/backend/server/src/plugins/copilot/providers/morph.ts
-60
View File
@@ -1,60 +0,0 @@
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import { CopilotProvider } from './provider';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType, ModelOutputType } from './types';
export const DEFAULT_DIMENSIONS = 256;
export type MorphConfig = {
apiKey?: string;
};
export class MorphProvider extends CopilotProvider<MorphConfig> {
readonly type = CopilotProviderType.Morph;
protected resolveModelBackendKind() {
return 'morph' as const;
}
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected morph response',
});
}
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
return {
base_url: 'https://api.morphllm.com',
auth_token: this.getConfig(execution).apiKey ?? '',
};
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveOutputType: kind =>
kind === 'streamObject' ? null : ModelOutputType.Text,
},
structured: false,
embedding: false,
rerank: false,
};
}
}
packages/backend/server/src/plugins/copilot/providers/openai.ts
-12
View File
@@ -14,7 +14,6 @@ import {
AttachmentAdmissionHost,
} from '../runtime/hosts/attachment-admission';
import { AttachmentMaterializer } from '../runtime/hosts/attachment-materializer';
import type { CopilotTool } from '../tools';
import { CopilotProvider } from './provider';
import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
import type {
@@ -22,7 +21,6 @@ import type {
ProviderDriverSpec,
} from './provider-runtime-contract';
import {
CopilotChatTools,
CopilotProviderType,
type PromptAttachment,
type PromptMessage,
@@ -64,16 +62,6 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
});
}
override getProviderSpecificTools(
toolName: CopilotChatTools,
_model: string
): [string, CopilotTool?] | undefined {
if (toolName === 'docEdit') {
return ['doc_edit', undefined];
}
return;
}
protected createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
packages/backend/server/src/plugins/copilot/providers/perplexity.ts
-75
View File
@@ -1,75 +0,0 @@
import { CopilotProviderSideError } from '../../../base';
import { type LlmBackendConfig } from '../../../native';
import { CopilotProvider } from './provider';
import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType, ModelOutputType } from './types';
export type PerplexityConfig = {
apiKey: string;
endpoint?: string;
};
export class PerplexityProvider extends CopilotProvider<PerplexityConfig> {
readonly type = CopilotProviderType.Perplexity;
protected resolveModelBackendKind() {
return 'perplexity' as const;
}
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveOutputType: kind =>
kind === 'streamObject' ? null : ModelOutputType.Text,
withAttachment: false,
resolveRequestOptions: async context => ({
withAttachment: !hasProviderModelBehaviorFlag(
context.model,
'no_attachments'
),
include: hasProviderModelBehaviorFlag(
context.model,
'citations_include'
)
? ['citations']
: undefined,
}),
},
structured: false,
embedding: false,
rerank: false,
};
}
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
const config = this.getConfig(execution);
const baseUrl = config.endpoint || 'https://api.perplexity.ai';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: config.apiKey,
};
}
private handleError(e: any) {
if (e instanceof CopilotProviderSideError) {
return e;
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected perplexity response',
});
}
}
packages/backend/server/src/plugins/copilot/providers/provider-middleware.ts
-12
View File
@@ -21,18 +21,6 @@ const DEFAULT_MIDDLEWARE_BY_TYPE: Record<
[CopilotProviderType.AnthropicVertex]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Morph]: {
rust: {
request: ['clamp_max_tokens'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Perplexity]: {
rust: {
request: ['clamp_max_tokens'],
},
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Gemini]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
packages/backend/server/src/plugins/copilot/providers/provider-registry.ts
-2
View File
@@ -15,10 +15,8 @@ const LEGACY_PROVIDER_ORDER: CopilotProviderType[] = [
CopilotProviderType.FAL,
CopilotProviderType.Gemini,
CopilotProviderType.GeminiVertex,
CopilotProviderType.Perplexity,
CopilotProviderType.Anthropic,
CopilotProviderType.AnthropicVertex,
CopilotProviderType.Morph,
];
const LEGACY_PROVIDER_PRIORITY = LEGACY_PROVIDER_ORDER.reduce(
packages/backend/server/src/plugins/copilot/providers/provider-tokens.ts
-4
View File
@@ -5,9 +5,7 @@ import {
import { CloudflareWorkersAIProvider } from './cloudflare';
import { FalProvider } from './fal';
import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
import { MorphProvider } from './morph';
import { OpenAIProvider } from './openai';
import { PerplexityProvider } from './perplexity';
export const CopilotProviders = [
OpenAIProvider,
@@ -15,8 +13,6 @@ export const CopilotProviders = [
FalProvider,
GeminiGenerativeProvider,
GeminiVertexProvider,
PerplexityProvider,
AnthropicOfficialProvider,
AnthropicVertexProvider,
MorphProvider,
];
packages/backend/server/src/plugins/copilot/providers/types.ts
+16 -4
View File
@@ -30,8 +30,6 @@ export enum CopilotProviderType {
Gemini = 'gemini',
GeminiVertex = 'geminiVertex',
OpenAI = 'openai',
Perplexity = 'perplexity',
Morph = 'morph',
}
export const CopilotProviderSchema = z.object({
@@ -80,8 +78,6 @@ export const PromptToolsSchema = z
'blobRead',
'codeArtifact',
'conversationSummary',
// work with morph
'docEdit',
// work with indexer
'docRead',
'docCreate',
@@ -268,6 +264,22 @@ const CopilotProviderOptionsSchema = z.object({
user: z.string().optional(),
session: z.string().optional(),
workspace: z.string().optional(),
byokLeaseId: z.string().optional(),
billingUnitId: z.string().optional(),
taskId: z.string().optional(),
actionId: z.string().optional(),
quotaBackedRoutesAllowed: z.boolean().optional(),
featureKind: z
.enum([
'chat',
'action',
'image',
'embedding',
'workspace_indexing',
'rerank',
'transcript',
])
.optional(),
});
export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
packages/backend/server/src/plugins/copilot/providers/utils.ts
+1 -34
View File
@@ -164,11 +164,6 @@ export function toError(error: unknown): Error {
}
}
type DocEditFootnote = {
intent: string;
result: string;
};
function asRecord(value: unknown): Record<string, unknown> | null {
if (value && typeof value === 'object' && !Array.isArray(value)) {
return value as Record<string, unknown>;
@@ -184,8 +179,6 @@ export class TextStreamParser {
private prefix: string | null = this.CALLOUT_PREFIX;
private readonly docEditFootnotes: DocEditFootnote[] = [];
public parse(chunk: CopilotTextStreamPart) {
let result = '';
switch (chunk.type) {
@@ -233,13 +226,6 @@ export class TextStreamParser {
result += `\nWriting document "${chunk.input.title}"\n`;
break;
}
case 'doc_edit': {
this.docEditFootnotes.push({
intent: String(chunk.input.instructions ?? ''),
result: '',
});
break;
}
}
result = this.markAsCallout(result);
break;
@@ -250,22 +236,6 @@ export class TextStreamParser {
);
result = this.addPrefix(result);
switch (chunk.toolName) {
case 'doc_edit': {
const output = asRecord(chunk.output);
const array = output?.result;
if (Array.isArray(array)) {
result += array
.map(item => {
return `\n${String(asRecord(item)?.changedContent ?? '')}\n`;
})
.join('');
this.docEditFootnotes[this.docEditFootnotes.length - 1].result =
result;
} else {
this.docEditFootnotes.pop();
}
break;
}
case 'doc_semantic_search': {
const output = chunk.output;
if (Array.isArray(output)) {
@@ -319,10 +289,7 @@ export class TextStreamParser {
}
public end() {
const footnotes = this.docEditFootnotes.map((footnote, index) => {
return `[^edit${index + 1}]: ${JSON.stringify({ type: 'doc-edit', ...footnote })}`;
});
return footnotes.join('\n');
return '';
}
private addPrefix(text: string) {