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

feat(server): improve context management (#15448)

Browse Source 
#### PR Dependency Tree


* **PR #15448** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added workspace artifact upload, browsing, removal, deduplication, and
library ownership support.
* Copilot now supports scoped document and artifact search, canvas
reading, live editor context, and frontend tools.
* Added scope and focus selectors with source-resolution receipts in
chat.
* Added embedding health, progress, synchronization, and retrieval
capabilities.
* Added BYOK policy visibility, provider restrictions, endpoint dialect
selection, and validation.
* Added delegated editor interactions and userdata document
authorization.

* **Bug Fixes**
* Improved attachment handling, cancellation, access control, retrieval
fallbacks, workspace synchronization, and configuration validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-08-10 09:27:58 +08:00
committed by GitHub
parent 42322d13fe
commit ee899a267b
311 changed files with 20468 additions and 14806 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/src/plugins/copilot/byok/resolver.ts
+79 -59
View File
@@ -12,7 +12,7 @@ import {
} from '@nestjs/graphql';
import { SafeIntResolver } from 'graphql-scalars';
import { Config, Throttle } from '../../../base';
import { Throttle } from '../../../base';
import { CurrentUser } from '../../../core/auth';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { PermissionAccess } from '../../../core/permission';
@@ -22,27 +22,36 @@ import { llmGetByokCatalog } from '../../../native';
import { CopilotEnabled } from '../feature';
import { ByokEntitlementPolicy } from './policy';
import {
BYOK_ALLOWED_PROVIDERS,
ByokAttachmentKind,
ByokAttachmentSource,
ByokCustomEndpointMode,
ByokEndpointKind,
ByokModelFeature,
ByokModelInput,
ByokModelOutput,
ByokOpenAiDialect,
ByokProbeOperation,
ByokProbeStatusKind,
ByokProvider,
ByokProviderSource,
} from './types';
@ObjectType()
class WorkspaceByokCapabilityType {
@Field(() => [String])
input!: string[];
@Field(() => [ByokModelInput])
input!: ByokModelInput[];
@Field(() => [String])
output!: string[];
@Field(() => [ByokModelOutput])
output!: ByokModelOutput[];
@Field(() => [String])
features!: string[];
@Field(() => [ByokModelFeature])
features!: ByokModelFeature[];
@Field(() => [String])
attachmentKinds!: string[];
@Field(() => [ByokAttachmentKind])
attachmentKinds!: ByokAttachmentKind[];
@Field(() => [String])
attachmentSources!: string[];
@Field(() => [ByokAttachmentSource])
attachmentSources!: ByokAttachmentSource[];
}
@ObjectType()
@@ -59,18 +68,18 @@ class WorkspaceByokModelDeclarationType {
@ObjectType()
class WorkspaceByokEndpointType {
@Field(() => String)
kind!: string;
@Field(() => ByokEndpointKind)
kind!: ByokEndpointKind;
@Field(() => String, { nullable: true })
url!: string | null;
@Field(() => ByokOpenAiDialect, { nullable: true })
dialect!: ByokOpenAiDialect | null;
}
@ObjectType()
class WorkspaceByokProfileDefinitionType {
@Field(() => SafeIntResolver)
version!: number;
@Field(() => WorkspaceByokEndpointType)
endpoint!: WorkspaceByokEndpointType;
@@ -80,8 +89,8 @@ class WorkspaceByokProfileDefinitionType {
@ObjectType()
class WorkspaceByokProbeStatusType {
@Field(() => String)
kind!: string;
@Field(() => ByokProbeStatusKind)
kind!: ByokProbeStatusKind;
@Field(() => Date, { nullable: true })
testedAt!: Date | null;
@@ -92,8 +101,8 @@ class WorkspaceByokProbeStatusType {
@ObjectType()
class WorkspaceByokModelProbeCheckType {
@Field(() => String)
operation!: string;
@Field(() => ByokProbeOperation)
operation!: ByokProbeOperation;
@Field(() => WorkspaceByokProbeStatusType)
status!: WorkspaceByokProbeStatusType;
@@ -204,6 +213,21 @@ class WorkspaceByokCatalogType {
providers!: WorkspaceByokCatalogProviderType[];
}
@ObjectType()
class WorkspaceByokPolicyType {
@Field(() => Boolean)
enabled!: boolean;
@Field(() => [ByokProvider])
allowedProviders!: ByokProvider[];
@Field(() => ByokCustomEndpointMode)
customEndpointMode!: ByokCustomEndpointMode;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
}
@ObjectType()
class WorkspaceByokSettingsType {
@Field(() => String)
@@ -221,14 +245,8 @@ class WorkspaceByokSettingsType {
@Field(() => [WorkspaceByokProfileType])
profiles!: WorkspaceByokProfileType[];
@Field(() => [ByokProvider])
allowedProviders!: ByokProvider[];
@Field(() => Boolean)
customEndpointSupported!: boolean;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
@Field(() => WorkspaceByokPolicyType)
policy!: WorkspaceByokPolicyType;
@Field(() => WorkspaceByokCatalogType)
catalog!: WorkspaceByokCatalogType;
@@ -257,20 +275,20 @@ class CreateWorkspaceByokLocalLeaseResultType {
@InputType()
class WorkspaceByokCapabilityInput {
@Field(() => [String])
input!: string[];
@Field(() => [ByokModelInput])
input!: ByokModelInput[];
@Field(() => [String])
output!: string[];
@Field(() => [ByokModelOutput])
output!: ByokModelOutput[];
@Field(() => [String])
features!: string[];
@Field(() => [ByokModelFeature])
features!: ByokModelFeature[];
@Field(() => [String])
attachmentKinds!: string[];
@Field(() => [ByokAttachmentKind])
attachmentKinds!: ByokAttachmentKind[];
@Field(() => [String])
attachmentSources!: string[];
@Field(() => [ByokAttachmentSource])
attachmentSources!: ByokAttachmentSource[];
}
@InputType()
@@ -287,18 +305,18 @@ class WorkspaceByokModelDeclarationInput {
@InputType()
class WorkspaceByokEndpointInput {
@Field(() => String)
kind!: string;
@Field(() => ByokEndpointKind)
kind!: ByokEndpointKind;
@Field(() => String, { nullable: true })
url!: string | null;
@Field(() => ByokOpenAiDialect, { nullable: true })
dialect!: ByokOpenAiDialect | null;
}
@InputType()
class WorkspaceByokProfileDefinitionInput {
@Field(() => SafeIntResolver)
version!: number;
@Field(() => WorkspaceByokEndpointInput)
endpoint!: WorkspaceByokEndpointInput;
@@ -377,8 +395,8 @@ class WorkspaceByokProbeCheckInput {
@Field(() => String)
modelId!: string;
@Field(() => String)
operation!: string;
@Field(() => ByokProbeOperation)
operation!: ByokProbeOperation;
}
@InputType()
@@ -472,8 +490,7 @@ export class WorkspaceByokResolver {
private readonly ac: PermissionAccess,
private readonly entitlement: ByokEntitlementPolicy,
private readonly runtime: BackendRuntimeProvider,
private readonly models: Models,
private readonly config: Config
private readonly models: Models
) {}
@ResolveField(() => WorkspaceByokSettingsType, {
@@ -491,8 +508,8 @@ export class WorkspaceByokResolver {
const profiles = serverEntitled
? await this.runtime.listByokProfiles(workspace.id)
: [];
const customEndpointSupported =
this.config.copilot.byok.allowCustomEndpoint;
const policy = await this.runtime.getByokPolicy();
const allowedProviders = new Set(policy.allowedProviders);
const catalog = llmGetByokCatalog();
return {
workspaceId: workspace.id,
@@ -500,17 +517,19 @@ export class WorkspaceByokResolver {
serverEntitled,
localEntitled,
profiles: profiles.map(profile => projectProfile(profile)),
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
customEndpointSupported,
privateEndpointSupported:
customEndpointSupported &&
this.config.copilot.byok.allowPrivateEndpoint,
policy: {
...policy,
allowedProviders: policy.allowedProviders as ByokProvider[],
customEndpointMode: policy.customEndpointMode as ByokCustomEndpointMode,
},
catalog: {
...catalog,
providers: catalog.providers.map(provider => ({
...provider,
provider: provider.provider as ByokProvider,
})),
providers: catalog.providers
.filter(provider => allowedProviders.has(provider.provider))
.map(provider => ({
...provider,
provider: provider.provider as ByokProvider,
})),
},
};
}
@@ -711,6 +730,7 @@ function nativeDefinition(input: WorkspaceByokProfileDefinitionInput) {
endpoint: {
...input.endpoint,
url: input.endpoint.url ?? undefined,
dialect: input.endpoint.dialect ?? undefined,
},
};
}
packages/backend/server/src/plugins/copilot/byok/types.ts
+89 -8
View File
@@ -26,6 +26,74 @@ export enum ByokProviderSource {
AffinePlan = 'affine_plan',
}
export enum ByokEndpointKind {
provider_default = 'provider_default',
openai_compatible = 'openai_compatible',
}
export enum ByokOpenAiDialect {
responses = 'responses',
chat_completions = 'chat_completions',
}
export enum ByokModelInput {
text = 'text',
image = 'image',
audio = 'audio',
file = 'file',
}
export enum ByokModelOutput {
text = 'text',
object = 'object',
structured = 'structured',
embedding = 'embedding',
rerank = 'rerank',
image = 'image',
}
export enum ByokModelFeature {
tool_calling = 'tool_calling',
reasoning = 'reasoning',
web_search = 'web_search',
}
export enum ByokAttachmentKind {
image = 'image',
audio = 'audio',
file = 'file',
}
export enum ByokAttachmentSource {
url = 'url',
data = 'data',
bytes = 'bytes',
file_handle = 'file_handle',
}
export enum ByokProbeOperation {
chat = 'chat',
structured = 'structured',
tool_calling = 'tool_calling',
vision = 'vision',
embedding = 'embedding',
rerank = 'rerank',
image = 'image',
transcript = 'transcript',
}
export enum ByokProbeStatusKind {
verified = 'verified',
failed = 'failed',
not_tested = 'not_tested',
}
export enum ByokCustomEndpointMode {
unavailable = 'unavailable',
disabled = 'disabled',
enabled = 'enabled',
}
export type ByokFeatureKind =
| 'chat'
| 'action'
@@ -35,13 +103,6 @@ export type ByokFeatureKind =
| 'transcript'
| 'workspace_indexing';
export const BYOK_ALLOWED_PROVIDERS = [
ByokProvider.openai,
ByokProvider.anthropic,
ByokProvider.gemini,
ByokProvider.fal,
] as const;
export function byokProviderToCopilotType(provider: ByokProvider) {
switch (provider) {
case ByokProvider.openai:
@@ -71,9 +132,29 @@ export function copilotTypeToByokProvider(type: CopilotProviderType) {
}
export function isByokProvider(value: string): value is ByokProvider {
return (BYOK_ALLOWED_PROVIDERS as readonly string[]).includes(value);
switch (value) {
case ByokProvider.openai:
case ByokProvider.anthropic:
case ByokProvider.gemini:
case ByokProvider.fal:
return true;
default:
return false;
}
}
registerEnumType(ByokProvider, { name: 'ByokProvider' });
registerEnumType(ByokKeyStorage, { name: 'ByokKeyStorage' });
registerEnumType(ByokKeyTestStatus, { name: 'ByokKeyTestStatus' });
registerEnumType(ByokEndpointKind, { name: 'ByokEndpointKind' });
registerEnumType(ByokOpenAiDialect, { name: 'ByokOpenAiDialect' });
registerEnumType(ByokModelInput, { name: 'ByokModelInput' });
registerEnumType(ByokModelOutput, { name: 'ByokModelOutput' });
registerEnumType(ByokModelFeature, { name: 'ByokModelFeature' });
registerEnumType(ByokAttachmentKind, { name: 'ByokAttachmentKind' });
registerEnumType(ByokAttachmentSource, { name: 'ByokAttachmentSource' });
registerEnumType(ByokProbeOperation, { name: 'ByokProbeOperation' });
registerEnumType(ByokProbeStatusKind, { name: 'ByokProbeStatusKind' });
registerEnumType(ByokCustomEndpointMode, {
name: 'ByokCustomEndpointMode',
});
packages/backend/server/src/plugins/copilot/config.ts
+45 -98
View File
@@ -1,7 +1,7 @@
import { z } from 'zod';
import serverNativeModule from '@affine/server-native';
import {
defineModuleConfig,
defineNativeModuleConfig,
StorageJSONSchema,
StorageProviderConfig,
} from '../../base';
@@ -9,28 +9,21 @@ import { CopilotProviderType } from './providers/types';
export type ProviderSpecificConfig = Record<string, unknown>;
export const RustRequestMiddlewareValues = [
'normalize_messages',
'clamp_max_tokens',
'tool_schema_rewrite',
'openai_request_compat',
'omit_tool_choice',
] as const;
export type RustRequestMiddleware =
(typeof RustRequestMiddlewareValues)[number];
| 'normalize_messages'
| 'clamp_max_tokens'
| 'tool_schema_rewrite'
| 'openai_request_compat'
| 'omit_tool_choice';
export const RustStreamMiddlewareValues = [
'stream_event_normalize',
'citation_indexing',
] as const;
export type RustStreamMiddleware = (typeof RustStreamMiddlewareValues)[number];
export type RustStreamMiddleware =
| 'stream_event_normalize'
| 'citation_indexing';
export const NodeTextMiddlewareValues = [
'citation_footnote',
'callout',
'thinking_format',
] as const;
export type NodeTextMiddleware = (typeof NodeTextMiddlewareValues)[number];
export type NodeTextMiddleware =
| 'citation_footnote'
| 'callout'
| 'thinking_format';
export type ProviderMiddlewareConfig = {
rust?: { request?: RustRequestMiddleware[]; stream?: RustStreamMiddleware[] };
@@ -51,32 +44,6 @@ export type CopilotProviderProfile = CopilotProviderProfileCommon & {
config: ProviderSpecificConfig;
};
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().min(1)).min(1),
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 CopilotProviderProfileShape = CopilotProviderProfileBaseShape.extend({
type: z.nativeEnum(CopilotProviderType),
config: z.record(z.string(), z.unknown()),
});
declare global {
interface AppConfigSchema {
copilot: {
@@ -103,57 +70,37 @@ declare global {
}
}
defineModuleConfig('copilot', {
enabled: {
desc: 'Enable AI features. Workspace owners configure provider keys in Workspace Settings → Integrations → AI BYOK.',
default: false,
},
'byok.enabled': {
desc: 'Allow workspace owners and admins to configure AI provider keys through AI BYOK.',
default: true,
shape: z.boolean(),
},
'byok.allowedProviders': {
desc: 'AI providers that workspace owners and admins may add through AI BYOK.',
default: ['openai', 'anthropic', 'gemini', 'fal'],
shape: z.array(z.enum(['openai', 'anthropic', 'gemini', 'fal'])),
},
'byok.allowCustomEndpoint': {
desc: 'Allow AI BYOK keys to use a custom provider endpoint.',
default: false,
shape: z.boolean(),
},
'byok.allowPrivateEndpoint': {
desc: 'Whether workspace BYOK custom endpoints may resolve to private network targets. Enabling this allows workspace owners and admins to send provider probe requests to the private network.',
default: false,
shape: z.boolean(),
},
'providers.profiles': {
desc: 'The profile list for copilot providers.',
default: [],
shape: z.array(CopilotProviderProfileShape),
},
unsplash: {
desc: 'The config for the unsplash key.',
default: {
key: '',
defineNativeModuleConfig(
'copilot',
serverNativeModule.appConfigDescriptors('copilot'),
serverNativeModule.validateAppConfigValue,
{
enabled: {
desc: 'Enable AI features. Workspace owners configure provider keys in Workspace Settings → Integrations → AI BYOK.',
default: false,
},
},
exa: {
desc: 'The config for the exa web search key.',
default: {
key: '',
},
},
storage: {
desc: 'The config for the storage provider.',
default: {
provider: 'fs',
bucket: 'copilot',
config: {
path: '~/.affine/storage',
unsplash: {
desc: 'The config for the unsplash key.',
default: {
key: '',
},
},
schema: StorageJSONSchema,
},
});
exa: {
desc: 'The config for the exa web search key.',
default: {
key: '',
},
},
storage: {
desc: 'The config for the storage provider.',
default: {
provider: 'fs',
bucket: 'copilot',
config: {
path: '~/.affine/storage',
},
},
schema: StorageJSONSchema,
},
}
);
packages/backend/server/src/plugins/copilot/context/index.ts
-3
View File
@@ -1,3 +0,0 @@
export { CopilotEmbeddingRealtimeProvider } from './realtime';
export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
export { CopilotContextService } from './service';
packages/backend/server/src/plugins/copilot/context/realtime.ts
-131
View File
@@ -1,131 +0,0 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { Config } from '../../../base/config';
import { OnEvent } from '../../../base/event';
import { PermissionAccess } from '../../../core/permission';
import {
RealtimePublisher,
RealtimeRegistry,
realtimeWorkspaceEmbeddingProgressRoom,
registerRealtimeLiveQuery,
} from '../../../core/realtime';
import { Models } from '../../../models';
import { assertCopilotEnabled } from '../availability';
export function workspaceEmbeddingRoom(workspaceId: string) {
return realtimeWorkspaceEmbeddingProgressRoom(workspaceId);
}
@Injectable()
export class CopilotEmbeddingRealtimeProvider implements OnModuleInit {
constructor(
private readonly ac: PermissionAccess,
private readonly models: Models,
private readonly registry: RealtimeRegistry,
private readonly publisher: RealtimePublisher,
private readonly config: Config
) {}
onModuleInit() {
const input = z.object({ workspaceId: z.string() });
registerRealtimeLiveQuery(this.registry, {
request: {
name: 'workspace.embedding.progress.get',
input,
handle: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
const canEmbedding =
await this.models.copilotWorkspace.checkEmbeddingAvailable();
if (!canEmbedding) {
return { total: 0, embedded: 0 };
}
return await this.models.copilotWorkspace.getEmbeddingStatus(
payload.workspaceId
);
},
},
topic: {
name: 'workspace.embedding.progress.changed',
input,
authorize: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
},
room: (_user, payload) => workspaceEmbeddingRoom(payload.workspaceId),
},
});
}
@OnEvent('workspace.doc.embed.finished', { suppressError: true })
async onDocEmbedFinished(payload: Events['workspace.doc.embed.finished']) {
await this.publishContext(payload.contextId, 'finished');
}
@OnEvent('workspace.doc.embed.failed', { suppressError: true })
async onDocEmbedFailed(payload: Events['workspace.doc.embed.failed']) {
await this.publishContext(payload.contextId, 'failed');
}
@OnEvent('workspace.file.embed.finished', { suppressError: true })
async onFileEmbedFinished(payload: Events['workspace.file.embed.finished']) {
await this.publishEmbeddingProgress(payload, 'finished');
}
@OnEvent('workspace.file.embed.failed', { suppressError: true })
async onFileEmbedFailed(payload: Events['workspace.file.embed.failed']) {
await this.publishEmbeddingProgress(payload, 'failed');
}
@OnEvent('workspace.blob.embed.finished', { suppressError: true })
async onBlobEmbedFinished(payload: Events['workspace.blob.embed.finished']) {
await this.publishContext(payload.contextId, 'finished');
}
@OnEvent('workspace.blob.embed.failed', { suppressError: true })
async onBlobEmbedFailed(payload: Events['workspace.blob.embed.failed']) {
await this.publishContext(payload.contextId, 'failed');
}
private async publishContext(
contextId: string,
reason: 'finished' | 'failed'
) {
if (!this.publisher) return;
const context = await this.models.copilotContext.getConfig(contextId);
if (!context) return;
this.publishWorkspace(context.workspaceId, reason);
}
private async publishEmbeddingProgress(
payload:
| Events['workspace.file.embed.finished']
| Events['workspace.file.embed.failed'],
reason: 'finished' | 'failed'
) {
if (!this.publisher) return;
if (payload.contextId) {
await this.publishContext(payload.contextId, reason);
return;
}
this.publishWorkspace(payload.workspaceId, reason);
}
private publishWorkspace(workspaceId: string, reason: 'finished' | 'failed') {
this.publisher.publish(
'workspace.embedding.progress.changed',
{ workspaceId },
{ reason },
{ room: workspaceEmbeddingRoom(workspaceId) }
);
}
private async assertCopilot(userId: string, workspaceId: string) {
assertCopilotEnabled(this.config);
await this.ac
.user(userId)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Copilot');
}
}
packages/backend/server/src/plugins/copilot/context/resolver.ts
-1063
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/plugins/copilot/context/service.ts
-381
View File
@@ -1,381 +0,0 @@
/* oxlint-disable import/no-cycle -- Context embedding reuses the shared capability runtime. */
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import {
Cache,
CopilotInvalidContext,
NoCopilotProviderAvailable,
OnEvent,
} from '../../../base';
import {
ContextConfig,
ContextConfigSchema,
ContextDoc,
ContextEmbedStatus,
ContextFile,
Models,
} from '../../../models';
import { CopilotEmbeddingClientService } from '../embedding/client';
import type {
EmbeddingCallOptions,
EmbeddingClient,
EmbeddingRouteContext,
} from '../embedding/types';
import { ContextSession } from './session';
const CONTEXT_SESSION_KEY = 'context-session';
@Injectable()
export class CopilotContextService implements OnApplicationBootstrap {
private supportEmbedding = false;
private client: EmbeddingClient | undefined;
constructor(
private readonly embeddingClients: CopilotEmbeddingClientService,
private readonly cache: Cache,
private readonly models: Models
) {}
@OnEvent('config.init')
async onConfigInit() {
await this.setup();
}
@OnEvent('config.changed')
async onConfigChanged() {
await this.setup();
}
private async setup() {
this.client = await this.embeddingClients.refresh();
}
async onApplicationBootstrap() {
const supportEmbedding =
await this.models.copilotContext.checkEmbeddingAvailable();
if (supportEmbedding) {
this.supportEmbedding = true;
}
}
get canEmbedding() {
return this.supportEmbedding;
}
// public this client to allow overriding in tests
get embeddingClient(): EmbeddingClient | undefined {
return this.client ?? this.embeddingClients.getClient();
}
private embeddingOptions(
workspaceId: string,
signal?: AbortSignal,
routeContext: EmbeddingRouteContext = {}
): EmbeddingCallOptions {
return { workspaceId, signal, ...routeContext, featureKind: 'embedding' };
}
private async saveConfig(
contextId: string,
config: ContextConfig,
refreshCache = false
): Promise<void> {
if (!refreshCache) {
await this.models.copilotContext.update(contextId, { config });
}
await this.cache.set(`${CONTEXT_SESSION_KEY}:${contextId}`, config);
}
private async getCachedSession(
contextId: string
): Promise<ContextSession | undefined> {
const cachedSession = await this.cache.get(
`${CONTEXT_SESSION_KEY}:${contextId}`
);
if (cachedSession) {
const config = ContextConfigSchema.safeParse(cachedSession);
if (config.success) {
return new ContextSession(
this.embeddingClient,
contextId,
config.data,
this.models,
this.saveConfig.bind(this, contextId)
);
}
}
return undefined;
}
// NOTE: we only cache config to avoid frequent database queries
// but we do not need to cache session instances because a distributed
// lock is already apply to mutation operation for the same context in
// the resolver, so there will be no simultaneous writing to the config
private async cacheSession(
contextId: string,
config: ContextConfig
): Promise<ContextSession> {
const dispatcher = this.saveConfig.bind(this, contextId);
await dispatcher(config, true);
return new ContextSession(
this.embeddingClient,
contextId,
config,
this.models,
dispatcher
);
}
async create(sessionId: string): Promise<ContextSession> {
// keep the context unique per session
const existsContext = await this.getBySessionId(sessionId);
if (existsContext) return existsContext;
const context = await this.models.copilotContext.create(sessionId);
const config = ContextConfigSchema.parse(context.config);
return await this.cacheSession(context.id, config);
}
async get(id: string): Promise<ContextSession> {
if (!this.embeddingClient) {
throw new NoCopilotProviderAvailable(
{ modelId: 'embedding' },
'embedding client not configured'
);
}
const context = await this.getCachedSession(id);
if (context) return context;
const config = await this.models.copilotContext.getConfig(id);
if (config) {
return this.cacheSession(id, config);
}
throw new CopilotInvalidContext({ contextId: id });
}
async getOwnedContext(
userId: string,
contextId: string,
options: { workspaceId?: string; sessionId?: string } = {}
): Promise<ContextSession> {
const accessInfo =
await this.models.copilotContext.getAccessInfo(contextId);
if (
!accessInfo ||
accessInfo.session.userId !== userId ||
(options.workspaceId &&
accessInfo.session.workspaceId !== options.workspaceId) ||
(options.sessionId && accessInfo.sessionId !== options.sessionId)
) {
throw new CopilotInvalidContext({ contextId });
}
return await this.get(contextId);
}
async getBySessionId(sessionId: string): Promise<ContextSession | null> {
const existsContext =
await this.models.copilotContext.getBySessionId(sessionId);
if (existsContext) return this.get(existsContext.id);
return null;
}
async matchWorkspaceBlobs(
workspaceId: string,
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const blobChunks = await this.models.copilotWorkspace.matchBlobEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
);
if (!blobChunks.length) return [];
return await client.reRank(content, blobChunks, topK, options);
}
async matchWorkspaceFiles(
workspaceId: string,
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const fileChunks = await this.models.copilotWorkspace.matchFileEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
);
if (!fileChunks.length) return [];
return await client.reRank(content, fileChunks, topK, options);
}
async matchWorkspaceDocs(
workspaceId: string,
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const workspaceChunks =
await this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
workspaceId,
topK * 2,
threshold
);
if (!workspaceChunks.length) return [];
return await client.reRank(content, workspaceChunks, topK, options);
}
async matchWorkspaceAll(
workspaceId: string,
content: string,
topK: number,
signal?: AbortSignal,
threshold: number = 0.8,
docIds?: string[],
scopedThreshold: number = 0.85,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const [fileChunks, blobChunks, workspaceChunks, scopedWorkspaceChunks] =
await Promise.all([
this.models.copilotWorkspace.matchFileEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
),
this.models.copilotWorkspace.matchBlobEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
),
this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
workspaceId,
topK * 2,
threshold
),
docIds
? this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
workspaceId,
topK * 2,
scopedThreshold,
docIds
)
: null,
]);
if (
!fileChunks.length &&
!blobChunks.length &&
!workspaceChunks.length &&
!scopedWorkspaceChunks?.length
) {
return [];
}
return await client.reRank(
content,
[
...fileChunks,
...blobChunks,
...workspaceChunks,
...(scopedWorkspaceChunks || []),
],
topK,
options
);
}
@OnEvent('workspace.doc.embed.failed')
async onDocEmbedFailed({
contextId,
docId,
}: Events['workspace.doc.embed.failed']) {
const context = await this.get(contextId);
await context.saveDocRecord(docId, doc => ({
...(doc as ContextDoc),
status: ContextEmbedStatus.failed,
}));
}
@OnEvent('workspace.doc.embed.finished')
async onDocEmbedFinished({
contextId,
docId,
}: Events['workspace.doc.embed.finished']) {
const context = await this.get(contextId);
await context.saveDocRecord(docId, doc => ({
...(doc as ContextDoc),
status: ContextEmbedStatus.finished,
}));
}
@OnEvent('workspace.file.embed.finished')
async onFileEmbedFinish({
contextId,
fileId,
chunkSize,
}: Events['workspace.file.embed.finished']) {
if (!contextId) return;
const context = await this.get(contextId);
await context.saveFileRecord(fileId, file => ({
...(file as ContextFile),
chunkSize,
status: ContextEmbedStatus.finished,
}));
}
@OnEvent('workspace.file.embed.failed')
async onFileEmbedFailed({
contextId,
fileId,
error,
}: Events['workspace.file.embed.failed']) {
if (!contextId) return;
const context = await this.get(contextId);
await context.saveFileRecord(fileId, file => ({
...(file as ContextFile),
error,
status: ContextEmbedStatus.failed,
}));
}
}
packages/backend/server/src/plugins/copilot/context/session.ts
-426
View File
@@ -1,426 +0,0 @@
import { nanoid } from 'nanoid';
import {
ContextBlob,
ContextCategories,
ContextCategory,
ContextConfig,
ContextDoc,
ContextEmbedStatus,
ContextFile,
FileChunkSimilarity,
Models,
} from '../../../models';
import type {
EmbeddingCallOptions,
EmbeddingClient,
EmbeddingRouteContext,
} from '../embedding/types';
export class ContextSession implements AsyncDisposable {
constructor(
private readonly client: EmbeddingClient | undefined,
private readonly contextId: string,
private readonly config: ContextConfig,
private readonly models: Models,
private readonly dispatcher?: (config: ContextConfig) => Promise<void>
) {}
get id() {
return this.contextId;
}
get workspaceId() {
return this.config.workspaceId;
}
get categories(): ContextCategory[] {
return this.config.categories.map(c => ({
...c,
docs: c.docs.map(d => ({ ...d })),
}));
}
get tags() {
const categories = this.config.categories;
return categories.filter(c => c.type === ContextCategories.Tag);
}
get collections() {
const categories = this.config.categories;
return categories.filter(c => c.type === ContextCategories.Collection);
}
get blobs(): ContextBlob[] {
return this.config.blobs.map(d => ({ ...d }));
}
get docs(): ContextDoc[] {
return this.config.docs.map(d => ({ ...d }));
}
get files(): Required<ContextFile>[] {
return this.config.files.map(f => this.fulfillFile(f));
}
get docIds() {
return Array.from(
new Set(
[this.config.docs, this.config.categories.flatMap(c => c.docs)]
.flat()
.map(d => d.id)
)
);
}
private embeddingOptions(
signal?: AbortSignal,
routeContext: EmbeddingRouteContext = {}
): EmbeddingCallOptions {
return {
workspaceId: this.workspaceId,
signal,
...routeContext,
featureKind: 'embedding',
};
}
async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
const category = this.config.categories.find(
c => c.type === type && c.id === id
);
if (category) {
const missingDocs = docs.filter(
docId => !category.docs.some(d => d.id === docId)
);
if (missingDocs.length) {
category.docs.push(
...missingDocs.map(id => ({
id,
createdAt: Date.now(),
status: ContextEmbedStatus.processing,
}))
);
await this.save();
}
return category;
}
const createdAt = Date.now();
const record = {
id,
type,
docs: docs.map(id => ({
id,
createdAt,
status: ContextEmbedStatus.processing,
})),
createdAt,
};
this.config.categories.push(record);
await this.save();
return record;
}
async removeCategoryRecord(type: ContextCategories, id: string) {
const index = this.config.categories.findIndex(
c => c.type === type && c.id === id
);
if (index >= 0) {
this.config.categories.splice(index, 1);
await this.save();
}
return true;
}
async addBlobRecord(blobId: string): Promise<ContextBlob | null> {
const existsBlob = this.config.blobs.find(b => b.id === blobId);
if (existsBlob) {
return existsBlob;
}
const blob = await this.models.blob.get(this.config.workspaceId, blobId);
if (!blob) return null;
const record: ContextBlob = {
id: blobId,
createdAt: Date.now(),
status: ContextEmbedStatus.processing,
};
this.config.blobs.push(record);
await this.save();
return record;
}
async getBlobMetadata() {
const blobIds = this.blobs.map(b => b.id);
const blobs = await this.models.blob.list(this.config.workspaceId, {
where: { key: { in: blobIds } },
select: { key: true, mime: true },
});
const blobChunkSizes = await this.models.copilotWorkspace.getBlobChunkSizes(
this.config.workspaceId,
blobIds
);
return blobs
.filter(b => !!blobChunkSizes.get(b.key))
.map(b => ({
id: b.key,
mimeType: b.mime,
chunkSize: blobChunkSizes.get(b.key),
}));
}
async getBlobContent(
blobId: string,
chunk?: number
): Promise<string | undefined> {
return this.models.copilotWorkspace.getBlobContent(
this.config.workspaceId,
blobId,
chunk
);
}
async removeBlobRecord(blobId: string): Promise<boolean> {
const index = this.config.blobs.findIndex(b => b.id === blobId);
if (index >= 0) {
this.config.blobs.splice(index, 1);
await this.save();
}
return true;
}
async addDocRecord(docId: string): Promise<ContextDoc> {
const doc = this.config.docs.find(f => f.id === docId);
if (doc) {
return doc;
}
const record = { id: docId, createdAt: Date.now() };
this.config.docs.push(record);
await this.save();
return record;
}
async removeDocRecord(docId: string): Promise<boolean> {
const index = this.config.docs.findIndex(f => f.id === docId);
if (index >= 0) {
this.config.docs.splice(index, 1);
await this.save();
}
return true;
}
private fulfillFile(file: ContextFile): Required<ContextFile> {
return {
...file,
mimeType: file.mimeType || 'application/octet-stream',
};
}
async addFile(
blobId: string,
name: string,
mimeType: string
): Promise<Required<ContextFile>> {
let fileId = nanoid();
const existsBlob = this.config.files.find(f => f.blobId === blobId);
if (existsBlob) {
// use exists file id if the blob exists
// we assume that the file content pointed to by the same blobId is consistent.
if (existsBlob.status === ContextEmbedStatus.finished) {
return this.fulfillFile(existsBlob);
}
fileId = existsBlob.id;
} else {
await this.saveFileRecord(fileId, file => ({
...file,
blobId,
chunkSize: 0,
name,
mimeType,
error: null,
createdAt: Date.now(),
}));
}
return this.fulfillFile(this.getFile(fileId) as ContextFile);
}
getFile(fileId: string): ContextFile | undefined {
return this.config.files.find(f => f.id === fileId);
}
async getFileContent(
fileId: string,
chunk?: number
): Promise<string | undefined> {
const file = this.getFile(fileId);
if (!file) return undefined;
return this.models.copilotContext.getFileContent(
this.contextId,
fileId,
chunk
);
}
async removeFile(fileId: string): Promise<boolean> {
await this.models.copilotContext.deleteFileEmbedding(
this.contextId,
fileId
);
this.config.files = this.config.files.filter(f => f.id !== fileId);
await this.save();
return true;
}
/**
* Match the input text with the file chunks
* @param content input text to match
* @param topK number of similar chunks to return, default 5
* @param signal abort signal
* @param threshold relevance threshold for the similarity score, higher threshold means more similar chunks, default 0.7, good enough based on prior experiments
* @returns list of similar chunks
*/
async matchFiles(
content: string,
topK: number = 5,
signal?: AbortSignal,
scopedThreshold: number = 0.85,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
): Promise<FileChunkSimilarity[]> {
if (!this.client) return [];
const options = this.embeddingOptions(signal, routeContext);
const embedding = await this.client.getEmbedding(content, options);
if (!embedding) return [];
const [context, workspace] = await Promise.all([
this.models.copilotContext.matchFileEmbedding(
embedding,
this.id,
topK * 2,
scopedThreshold
),
this.models.copilotWorkspace.matchFileEmbedding(
this.workspaceId,
embedding,
topK * 2,
threshold
),
]);
const files = new Map(this.files.map(f => [f.id, f]));
return this.client.reRank(
content,
[
...context
.filter(f => files.has(f.fileId))
.map(c => {
const { blobId, name, mimeType } = files.get(
c.fileId
) as Required<ContextFile>;
return { ...c, blobId, name, mimeType };
}),
...workspace,
],
topK,
options
);
}
/**
* Match the input text with the workspace chunks
* @param content input text to match
* @param topK number of similar chunks to return, default 5
* @param signal abort signal
* @param threshold relevance threshold for the similarity score, higher threshold means more similar chunks, default 0.7, good enough based on prior experiments
* @returns list of similar chunks
*/
async matchWorkspaceDocs(
content: string,
topK: number = 5,
signal?: AbortSignal,
scopedThreshold: number = 0.85,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
if (!this.client) return [];
const options = this.embeddingOptions(signal, routeContext);
const embedding = await this.client.getEmbedding(content, options);
if (!embedding) return [];
const docIds = this.docIds;
const [inContext, workspace] = await Promise.all([
this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
this.workspaceId,
topK * 2,
scopedThreshold,
docIds
),
this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
this.workspaceId,
topK * 2,
threshold
),
]);
const result = await this.client.reRank(
content,
[...inContext, ...workspace],
topK,
options
);
// sort result, doc recorded in context first
const docIdSet = new Set(docIds);
return result.toSorted(
(a, b) =>
(docIdSet.has(a.docId) ? -1 : 1) - (docIdSet.has(b.docId) ? -1 : 1) ||
(a.distance || Infinity) - (b.distance || Infinity)
);
}
async saveDocRecord(
docId: string,
cb: (
record: Pick<ContextDoc, 'id' | 'status'> &
Partial<Omit<ContextDoc, 'id' | 'status'>>
) => ContextDoc
) {
const docs = [this.config.docs, ...this.config.categories.map(c => c.docs)]
.flat()
.filter(d => d.id === docId);
for (const doc of docs) {
Object.assign(doc, cb({ ...doc }));
}
await this.save();
}
async saveFileRecord(
fileId: string,
cb: (
record: Pick<ContextFile, 'id' | 'status'> &
Partial<Omit<ContextFile, 'id' | 'status'>>
) => ContextFile
) {
const files = this.config.files;
const file = files.find(f => f.id === fileId);
if (file) {
Object.assign(file, cb({ ...file }));
} else {
const file = { id: fileId, status: ContextEmbedStatus.processing };
files.push(cb(file));
}
await this.save();
}
async save() {
await this.dispatcher?.(this.config);
}
async [Symbol.asyncDispose]() {
await this.save();
}
}
packages/backend/server/src/plugins/copilot/conversation/inbox.ts
+28 -1
View File
@@ -8,6 +8,7 @@ import {
sniffMime,
} from '../../../base';
import { PermissionAccess } from '../../../core/permission';
import { Models } from '../../../models';
import { processImage } from '../../../native';
import { CompatSubmissionStore } from '../compat/submission-store';
import type { PromptMessage } from '../providers/types';
@@ -30,6 +31,7 @@ export class ConversationInboxService {
constructor(
private readonly chatSession: ChatSessionService,
private readonly ac: PermissionAccess,
private readonly models: Models,
private readonly storage: CopilotStorage,
private readonly submissions: CompatSubmissionStore
) {}
@@ -48,6 +50,26 @@ export class ConversationInboxService {
options.blob ? [options.blob] : options.blobs || []
);
const focusSelectors = options.params?.focusSelectors;
const hasWorkspaceContext =
attachments.length > 0 ||
blobs.length > 0 ||
(Array.isArray(options.params?.scopeSelectors) &&
options.params.scopeSelectors.length > 0) ||
(Array.isArray(options.params?.preferredSourceIds) &&
options.params.preferredSourceIds.length > 0) ||
(focusSelectors === undefined
? session.config.focus.selectors.length > 0
: Array.isArray(focusSelectors) && focusSelectors.length > 0);
if (
hasWorkspaceContext &&
!(await this.models.workspace.get(session.config.workspaceId))
) {
throw new BadRequestException(
"Local workspaces don't support attachments or references."
);
}
if (blobs.length) {
await this.ac
.user(userId)
@@ -86,7 +108,12 @@ export class ConversationInboxService {
filename,
attachmentBuffer
);
attachments.push({ attachment, mimeType: attachmentMimeType });
attachments.push({
kind: 'url',
url: attachment,
mimeType: attachmentMimeType,
fileName: blob.filename,
});
}
return await this.submissions.create({
packages/backend/server/src/plugins/copilot/conversation/store.ts
+24
View File
@@ -12,6 +12,10 @@ import {
type Turn,
turnFromChatMessage,
} from '../core';
import {
type SessionFocus,
SessionFocusSchema,
} from '../runtime/contracts/shared';
import { type ChatMessage, ChatMessageSchema } from '../types';
type SessionRecord = NonNullable<
@@ -68,6 +72,11 @@ export class ConversationStore {
return parsed.data;
}
private toFocus(focus: unknown): SessionFocus {
const parsed = SessionFocusSchema.safeParse(focus);
return parsed.success ? parsed.data : { selectors: [] };
}
async create(
seed: ConversationSeed,
reuseLatestChat = false
@@ -83,6 +92,7 @@ export class ConversationStore {
conversation: Conversation;
turns: Turn[];
promptName: string;
focus: SessionFocus;
}
| undefined
> {
@@ -95,6 +105,7 @@ export class ConversationStore {
conversation: this.toConversation(session),
turns: this.toTurns(session),
promptName: session.promptName,
focus: this.toFocus(session.focus),
};
}
@@ -102,6 +113,7 @@ export class ConversationStore {
| {
conversation: Conversation;
promptName: string;
focus: SessionFocus;
}
| undefined
> {
@@ -121,6 +133,7 @@ export class ConversationStore {
updatedAt: session.updatedAt,
},
promptName: session.promptName,
focus: this.toFocus(session.focus),
};
}
@@ -142,6 +155,7 @@ export class ConversationStore {
turnFromChatMessage(message, session.id)
),
promptName: session.promptName,
focus: this.toFocus(session.focus),
}));
}
@@ -163,6 +177,7 @@ export class ConversationStore {
updatedAt: session.updatedAt,
} satisfies Conversation,
promptName: session.promptName,
focus: this.toFocus(session.focus),
}));
}
@@ -185,10 +200,19 @@ export class ConversationStore {
userId: string;
turn: Turn;
compatSubmissionId?: string;
focus?: SessionFocus;
artifacts?: Array<{
artifactId: string;
role: string;
displayName?: string;
metadata?: Record<string, unknown>;
}>;
}) {
const message = await this.models.copilotSession.appendMessage({
sessionId: input.sessionId,
userId: input.userId,
focus: input.focus,
artifacts: input.artifacts,
message: (() => {
const { id: _id, ...message } = chatMessageFromTurn(input.turn);
return { ...message, compatSubmissionId: input.compatSubmissionId };
packages/backend/server/src/plugins/copilot/core/adapters.ts
+16 -6
View File
@@ -1,4 +1,5 @@
import type { PromptMessage, StreamObject } from '../providers/types';
import { promptAttachmentMimeType } from '../providers/utils';
import {
streamObjectToToolEvent,
toolEventToStreamObject,
@@ -82,6 +83,7 @@ export const turnFromChatMessage = (
renderTrace: trace.renderTrace,
toolEvents: trace.toolEvents,
metadata: message.params ?? {},
scopeSnapshot: message.scopeSnapshot,
createdAt: message.createdAt,
});
};
@@ -95,14 +97,22 @@ export const chatMessageFromTurn = (turn: Turn): ChatMessage => {
content: turn.content,
attachments: turn.attachments.length ? turn.attachments : undefined,
params: turn.metadata,
scopeSnapshot: turn.scopeSnapshot,
streamObjects: renderTrace.length ? renderTrace : undefined,
createdAt: turn.createdAt,
};
};
export const promptMessageFromTurn = (turn: Turn): PromptMessage => ({
role: turn.role,
content: turn.content,
attachments: turn.attachments.length ? turn.attachments : undefined,
params: Object.keys(turn.metadata).length ? turn.metadata : undefined,
});
export const promptMessageFromTurn = (turn: Turn): PromptMessage => {
const attachments = turn.attachments.filter(attachment => {
const mimeType = promptAttachmentMimeType(attachment);
return !mimeType || mimeType.startsWith('image/');
});
return {
role: turn.role,
content: turn.content,
attachments: attachments.length ? attachments : undefined,
params: Object.keys(turn.metadata).length ? turn.metadata : undefined,
};
};
packages/backend/server/src/plugins/copilot/core/types.ts
+2
View File
@@ -6,6 +6,7 @@ import {
type ToolEvent,
ToolEventSchema,
} from '../runtime/contracts/runtime-event-contract';
import { TurnScopeSnapshotSchema } from '../runtime/contracts/shared';
const CanonicalDateSchema = z.coerce.date();
@@ -35,6 +36,7 @@ export const TurnSchema = z
renderTrace: z.array(StreamObjectSchema).default([]),
toolEvents: z.array(ToolEventSchema).default([]),
metadata: z.record(z.string(), z.any()).default({}),
scopeSnapshot: TurnScopeSnapshotSchema.nullable().optional(),
createdAt: CanonicalDateSchema,
})
.strict();
packages/backend/server/src/plugins/copilot/cron.ts
+1 -37
View File
@@ -1,19 +1,15 @@
import { Injectable, Logger } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule';
import { JOB_SIGNAL, JobQueue, OneDay, OnJob } from '../../base';
import { JobQueue, OneDay, OnJob } from '../../base';
import { Models } from '../../models';
const CLEANUP_EMBEDDING_JOB_BATCH_SIZE = 100;
const BACKGROUND_COPILOT_JOB_PRIORITY = 100;
declare global {
interface Jobs {
'copilot.session.cleanupEmptySessions': {};
'copilot.session.generateMissingTitles': {};
'copilot.workspace.cleanupTrashedDocEmbeddings': {
nextSid?: number;
};
}
}
@@ -39,12 +35,6 @@ export class CopilotCronJobs {
{},
{ jobId: 'daily-copilot-generate-missing-titles' }
);
await this.jobs.add(
'copilot.workspace.cleanupTrashedDocEmbeddings',
{},
{ jobId: 'daily-copilot-cleanup-trashed-doc-embeddings' }
);
}
async triggerGenerateMissingTitles() {
@@ -82,30 +72,4 @@ export class CopilotCronJobs {
`Scheduled title generation for ${sessions.length} sessions`
);
}
@OnJob('copilot.workspace.cleanupTrashedDocEmbeddings')
async cleanupTrashedDocEmbeddings(
params: Jobs['copilot.workspace.cleanupTrashedDocEmbeddings']
) {
const nextSid = params.nextSid ?? 0;
// only consider workspaces that cleared their embeddings more than 24 hours ago
const oneDayAgo = new Date(Date.now() - OneDay);
const workspaces = await this.models.workspace.list(
{ sid: { gt: nextSid }, lastCheckEmbeddings: { lt: oneDayAgo } },
{ id: true, sid: true },
CLEANUP_EMBEDDING_JOB_BATCH_SIZE
);
if (!workspaces.length) {
return JOB_SIGNAL.Done;
}
for (const { id: workspaceId } of workspaces) {
await this.jobs.add(
'copilot.embedding.cleanupTrashedDocEmbeddings',
{ workspaceId },
{ jobId: `cleanup-trashed-doc-embeddings-${workspaceId}` }
);
}
params.nextSid = workspaces[workspaces.length - 1].sid;
return JOB_SIGNAL.Repeat;
}
}
packages/backend/server/src/plugins/copilot/delegated/realtime.ts
+144
View File
@@ -0,0 +1,144 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { EventBus } from '../../../base';
import { RealtimeRegistry, realtimeUserRoom } from '../../../core/realtime';
import { ChatSessionService } from '../session';
import { DelegatedEditorService } from './service';
const identity = {
requestId: z.string().uuid(),
runId: z.string().uuid(),
toolCallId: z.string().min(1).max(256),
sessionId: z.string().min(1),
workspaceId: z.string().min(1),
docId: z.string().min(1),
clientId: z.string().min(1).max(128),
editorStateId: z.string().min(1).max(128),
};
const responseSchema = z
.object({
...identity,
result: z.unknown().optional(),
error: z
.object({
code: z.string().min(1).max(64),
message: z.string().max(500),
retryable: z.boolean(),
})
.strict()
.optional(),
})
.strict()
.refine(
response =>
(response.result !== undefined) !== (response.error !== undefined),
{ message: 'Exactly one of result or error is required.' }
)
.refine(
response =>
Buffer.byteLength(JSON.stringify(response.result ?? null)) <= 512 * 1024,
{ message: 'Delegated tool result is too large.' }
);
@Injectable()
export class DelegatedEditorRealtimeProvider implements OnModuleInit {
constructor(
private readonly registry: RealtimeRegistry,
private readonly event: EventBus,
private readonly sessions: ChatSessionService,
private readonly delegated: DelegatedEditorService
) {}
onModuleInit() {
const leaseInput = z
.object({
clientId: z.string().min(1).max(128),
sessionId: z.string().min(1),
workspaceId: z.string().min(1),
docId: z.string().min(1),
editorStateId: z.string().min(1).max(128),
mode: z.enum(['page', 'edgeless']),
readonly: z.boolean(),
focused: z.boolean(),
capabilities: z
.array(
z.enum([
'frontend_get_editor_state',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
])
)
.max(4),
})
.strict();
this.registry.registerRequest({
name: 'copilot.delegated.editor.upsert',
input: leaseInput,
handle: async (user, input, context) => {
const session = await this.sessions.get(input.sessionId);
if (
!user ||
!context?.connectionId ||
!session ||
session.config.userId !== user.id ||
session.config.workspaceId !== input.workspaceId ||
session.config.docId !== input.docId
) {
throw new Error('INVALID_DELEGATED_EDITOR_SESSION');
}
const lease = this.delegated.upsert(
user.id,
context.connectionId,
input
);
this.event.broadcast('copilot.delegated.editor.upserted', lease);
return { ok: true, expiresAt: lease.expiresAt };
},
});
this.registry.registerRequest({
name: 'copilot.delegated.editor.release',
input: z
.object({
clientId: z.string().min(1).max(128),
editorStateId: z.string().min(1).max(128),
})
.strict(),
handle: async (user, input) => {
if (user) {
this.delegated.release(user.id, input.clientId, input.editorStateId);
this.event.broadcast('copilot.delegated.editor.released', {
userId: user.id,
...input,
});
}
return { ok: true };
},
});
this.registry.registerRequest({
name: 'copilot.delegated.tool.respond',
input: responseSchema,
handle: async (user, response) => {
if (!user) return { accepted: false };
const accepted = this.delegated.receive(user.id, response);
this.event.broadcast('copilot.delegated.tool.responded', {
userId: user.id,
response,
});
return { accepted };
},
});
this.registry.registerTopic({
name: 'copilot.delegated.tool.requested',
input: z.object({ clientId: z.string().min(1).max(128) }).strict(),
authorize: async user => {
if (!user) throw new Error('AUTHENTICATION_REQUIRED');
},
room: (user, input) => {
if (!user) throw new Error('AUTHENTICATION_REQUIRED');
return realtimeUserRoom(user.id, `copilot:${input.clientId}`);
},
});
}
}
packages/backend/server/src/plugins/copilot/delegated/service.ts
+303
View File
@@ -0,0 +1,303 @@
import { randomUUID } from 'node:crypto';
import type {
DelegatedEditorLeaseInput,
DelegatedToolIdentity,
DelegatedToolName,
DelegatedToolResponse,
} from '@affine/realtime';
import { Injectable } from '@nestjs/common';
import { OnEvent } from '../../../base';
import { RealtimePublisher, realtimeUserRoom } from '../../../core/realtime';
import type { CopilotChatOptions } from '../providers/types';
type EditorLease = DelegatedEditorLeaseInput & {
userId: string;
connectionId: string;
expiresAt: number;
};
type PendingRequest = {
identity: DelegatedToolIdentity;
userId: string;
connectionId: string;
resolve: (response: DelegatedToolResponse) => void;
};
declare global {
interface Events {
'copilot.delegated.editor.upserted': EditorLease;
'copilot.delegated.editor.released': {
userId: string;
clientId: string;
editorStateId: string;
};
'copilot.delegated.tool.responded': {
userId: string;
response: DelegatedToolResponse;
};
}
}
const LEASE_TTL_MS = 30_000;
const TOOL_TIMEOUT_MS = 15_000;
@Injectable()
export class DelegatedEditorService {
private readonly leases = new Map<string, EditorLease>();
private readonly pending = new Map<string, PendingRequest>();
constructor(private readonly publisher: RealtimePublisher) {}
leaseKey(userId: string, clientId: string) {
return `${userId}:${clientId}`;
}
upsert(
userId: string,
connectionId: string,
input: DelegatedEditorLeaseInput
) {
const lease = {
...input,
userId,
connectionId,
expiresAt: Date.now() + LEASE_TTL_MS,
};
this.leases.set(this.leaseKey(userId, input.clientId), lease);
return lease;
}
release(userId: string, clientId: string, editorStateId: string) {
const key = this.leaseKey(userId, clientId);
const lease = this.leases.get(key);
if (lease?.editorStateId === editorStateId) {
this.leases.delete(key);
}
}
getLease(options: CopilotChatOptions, tool?: DelegatedToolName) {
if (!options?.user || !options.session || !options.workspace) return null;
const now = Date.now();
let selected: EditorLease | null = null;
for (const [key, lease] of this.leases) {
if (lease.expiresAt <= now) {
this.leases.delete(key);
continue;
}
if (
lease.userId === options.user &&
lease.sessionId === options.session &&
lease.workspaceId === options.workspace &&
lease.focused &&
(!tool || lease.capabilities.includes(tool)) &&
(!selected || lease.expiresAt > selected.expiresAt)
) {
selected = lease;
}
}
return selected;
}
async execute(
options: CopilotChatOptions,
tool: DelegatedToolName,
args: Record<string, unknown>,
signal?: AbortSignal,
execution?: { runId?: string; toolCallId?: string }
) {
const lease = this.getLease(options, tool);
if (!lease) {
return {
error: {
code: 'FRONTEND_UNAVAILABLE',
message: 'No focused editor is available for this session.',
retryable: true,
},
};
}
const identity = {
requestId: randomUUID(),
runId: execution?.runId ?? randomUUID(),
toolCallId: execution?.toolCallId ?? randomUUID(),
sessionId: lease.sessionId,
workspaceId: lease.workspaceId,
docId: lease.docId,
clientId: lease.clientId,
editorStateId: lease.editorStateId,
};
const deadlineAt = Date.now() + TOOL_TIMEOUT_MS;
const response = new Promise<DelegatedToolResponse>(resolve => {
this.pending.set(identity.requestId, {
identity,
userId: lease.userId,
connectionId: lease.connectionId,
resolve,
});
});
this.publisher.publish(
'copilot.delegated.tool.requested',
{ clientId: lease.clientId },
{ type: 'request', ...identity, tool, args, deadlineAt },
{ room: realtimeUserRoom(lease.userId, `copilot:${lease.clientId}`) }
);
let reason: 'aborted' | 'timeout' | undefined;
let timeout: ReturnType<typeof setTimeout> | undefined;
let abort: (() => void) | undefined;
const interrupted = new Promise<DelegatedToolResponse>(resolve => {
timeout = setTimeout(() => {
reason = 'timeout';
resolve({
...identity,
error: {
code: 'FRONTEND_TIMEOUT',
message: 'The focused editor did not respond before the deadline.',
retryable: true,
},
});
}, TOOL_TIMEOUT_MS);
timeout.unref?.();
abort = () => {
reason = 'aborted';
resolve({
...identity,
error: {
code: 'ABORTED',
message: 'The delegated read was cancelled.',
retryable: false,
},
});
};
if (signal?.aborted) {
abort();
} else {
signal?.addEventListener('abort', abort, { once: true });
}
});
const result = await Promise.race([response, interrupted]);
this.pending.delete(identity.requestId);
if (timeout) clearTimeout(timeout);
if (abort) signal?.removeEventListener('abort', abort);
if (reason) {
this.publisher.publish(
'copilot.delegated.tool.requested',
{ clientId: lease.clientId },
{ type: 'cancel', ...identity, reason },
{ room: realtimeUserRoom(lease.userId, `copilot:${lease.clientId}`) }
);
}
if (result.error) return { error: result.error };
if (
tool === 'frontend_get_editor_state' ||
!result.result ||
typeof result.result !== 'object' ||
Array.isArray(result.result)
) {
return result.result;
}
return {
...result.result,
source: {
type: 'document',
workspace_id: lease.workspaceId,
doc_id: lease.docId,
revision: lease.editorStateId,
},
};
}
receive(userId: string, response: DelegatedToolResponse) {
const request = this.pending.get(response.requestId);
if (
!request ||
request.userId !== userId ||
!this.sameIdentity(request.identity, response) ||
!this.validResult(request.identity, response)
) {
return false;
}
this.pending.delete(response.requestId);
request.resolve(response);
return true;
}
@OnEvent('copilot.delegated.editor.upserted', { suppressError: true })
onRemoteUpsert(lease: Events['copilot.delegated.editor.upserted']) {
this.leases.set(this.leaseKey(lease.userId, lease.clientId), lease);
}
@OnEvent('copilot.delegated.editor.released', { suppressError: true })
onRemoteRelease(event: Events['copilot.delegated.editor.released']) {
this.release(event.userId, event.clientId, event.editorStateId);
}
@OnEvent('copilot.delegated.tool.responded', { suppressError: true })
onRemoteResponse(event: Events['copilot.delegated.tool.responded']) {
this.receive(event.userId, event.response);
}
@OnEvent('realtime.connection.disconnected', { suppressError: true })
onDisconnect({ connectionId }: Events['realtime.connection.disconnected']) {
for (const [key, lease] of this.leases) {
if (lease.connectionId === connectionId) {
this.leases.delete(key);
}
}
for (const [requestId, request] of this.pending) {
if (request.connectionId !== connectionId) continue;
this.pending.delete(requestId);
request.resolve({
...request.identity,
error: {
code: 'FRONTEND_DISCONNECTED',
message: 'The focused editor disconnected during the read.',
retryable: true,
},
});
this.publisher.publish(
'copilot.delegated.tool.requested',
{ clientId: request.identity.clientId },
{ type: 'cancel', ...request.identity, reason: 'disconnect' },
{
room: realtimeUserRoom(
request.userId,
`copilot:${request.identity.clientId}`
),
}
);
}
}
private sameIdentity(
expected: DelegatedToolIdentity,
actual: DelegatedToolIdentity
) {
return (
expected.requestId === actual.requestId &&
expected.runId === actual.runId &&
expected.toolCallId === actual.toolCallId &&
expected.sessionId === actual.sessionId &&
expected.workspaceId === actual.workspaceId &&
expected.docId === actual.docId &&
expected.clientId === actual.clientId &&
expected.editorStateId === actual.editorStateId
);
}
private validResult(
identity: DelegatedToolIdentity,
response: DelegatedToolResponse
) {
if (response.error) return true;
return Boolean(
response.result &&
typeof response.result === 'object' &&
'editor_state_id' in response.result &&
response.result.editor_state_id === identity.editorStateId
);
}
}
packages/backend/server/src/plugins/copilot/embedding/client.ts
-237
View File
@@ -1,237 +0,0 @@
/* oxlint-disable import/no-cycle -- Embedding delegates to the shared capability runtime. */
import { createHash } from 'node:crypto';
import { forwardRef, Inject, Injectable, Logger } from '@nestjs/common';
import { CopilotFailedToGenerateEmbedding } from '../../../base/error/errors.gen';
import {
ChunkSimilarity,
Embedding,
EMBEDDING_DIMENSIONS,
} from '../../../models';
import { type CopilotRerankRequest } from '../providers/types';
import { CapabilityRuntime } from '../runtime/capability-runtime';
import {
type EmbeddingCallOptionsInput,
EmbeddingClient,
normalizeEmbeddingCallOptions,
type ReRankResult,
} from './types';
type EmbeddingRuntime = Pick<
CapabilityRuntime,
'embeddingConfigured' | 'embed' | 'rerank'
>;
class ProductionEmbeddingClient extends EmbeddingClient {
private readonly logger = new Logger(ProductionEmbeddingClient.name);
constructor(private readonly runtime: EmbeddingRuntime) {
super();
}
override async configured(): Promise<boolean> {
const result = await this.runtime.embeddingConfigured('route-selected');
if (!result) {
this.logger.warn(
'Copilot embedding client is not configured properly, please check your configuration.'
);
}
return result;
}
async getEmbeddings(
input: string[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const modelId = 'route-selected';
const embeddings = await this.runtime.embed(modelId, input, {
dimensions: EMBEDDING_DIMENSIONS,
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: normalizedOptions.featureKind ?? 'embedding',
});
if (embeddings.length !== input.length) {
throw new CopilotFailedToGenerateEmbedding({
provider: modelId,
message: `Expected ${input.length} embeddings, got ${embeddings.length}`,
});
}
return Array.from(embeddings.entries()).map(([index, embedding]) => ({
index,
embedding,
content: input[index],
}));
}
private getTargetId<T extends ChunkSimilarity>(embedding: T) {
return 'docId' in embedding && typeof embedding.docId === 'string'
? embedding.docId
: 'fileId' in embedding && typeof embedding.fileId === 'string'
? embedding.fileId
: '';
}
private async getEmbeddingRelevance<
Chunk extends ChunkSimilarity = ChunkSimilarity,
>(
query: string,
embeddings: Chunk[],
options?: EmbeddingCallOptionsInput
): Promise<ReRankResult> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
if (!embeddings.length) return [];
const rerankRequest: CopilotRerankRequest = {
query,
candidates: embeddings.map((embedding, index) => ({
id: String(index),
text: embedding.content,
})),
};
const ranks = await this.runtime.rerank('route-selected', rerankRequest, {
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: 'rerank',
});
try {
return ranks.map((score, i) => {
const chunk = embeddings[i];
return {
chunk: chunk.chunk,
targetId: this.getTargetId(chunk),
score: Math.max(score, 1 - (chunk.distance || -Infinity)),
};
});
} catch (error) {
this.logger.error('Failed to parse rerank results', error);
// silent error, will fallback to default sorting in parent method
return [];
}
}
override async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
query: string,
embeddings: Chunk[],
topK: number,
options?: EmbeddingCallOptionsInput
): Promise<Chunk[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
// search in context and workspace may find same chunks, de-duplicate them
const { deduped: dedupedEmbeddings } = embeddings.reduce(
(acc, e) => {
const key = `${this.getTargetId(e)}:${e.chunk}`;
if (!acc.seen.has(key)) {
acc.seen.add(key);
acc.deduped.push(e);
}
return acc;
},
{ deduped: [] as Chunk[], seen: new Set<string>() }
);
const sortedEmbeddings = dedupedEmbeddings.toSorted(
(a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity)
);
const chunks = sortedEmbeddings.reduce(
(acc, e) => {
const targetId = this.getTargetId(e);
const key = `${targetId}:${e.chunk}`;
acc[key] = e;
return acc;
},
{} as Record<string, Chunk>
);
try {
// The rerank prompt is expected to handle the full deduped candidate list.
const ranks = await this.getEmbeddingRelevance(
query,
sortedEmbeddings,
normalizedOptions
);
if (sortedEmbeddings.length !== ranks.length) {
// llm return wrong result, fallback to default sorting
this.logger.warn(
`Batch size mismatch: expected ${sortedEmbeddings.length}, got ${ranks.length}`
);
return await super.reRank(
query,
dedupedEmbeddings,
topK,
normalizedOptions
);
}
const highConfidenceChunks = ranks
.flat()
.toSorted((a, b) => b.score - a.score)
.filter(r => r.score > 0.5)
.map(r => chunks[`${r.targetId}:${r.chunk}`])
.filter(Boolean);
this.logger.verbose(
`ReRank completed: ${highConfidenceChunks.length} high-confidence results found, total ${sortedEmbeddings.length} embeddings`,
highConfidenceChunks.length !== sortedEmbeddings.length
? JSON.stringify(ranks)
: undefined
);
return highConfidenceChunks.slice(0, topK);
} catch (error) {
this.logger.warn('ReRank failed, falling back to default sorting', error);
return await super.reRank(
query,
dedupedEmbeddings,
topK,
normalizedOptions
);
}
}
}
@Injectable()
export class CopilotEmbeddingClientService {
private client: EmbeddingClient | undefined;
constructor(
@Inject(forwardRef(() => CapabilityRuntime))
private readonly runtime: EmbeddingRuntime
) {}
async refresh() {
const client = new ProductionEmbeddingClient(this.runtime);
await client.configured();
this.client = client;
return this.client;
}
getClient() {
return this.client;
}
}
export class MockEmbeddingClient extends EmbeddingClient {
private embed(content: string) {
const seed = createHash('sha256').update(content).digest();
return Array.from({ length: EMBEDDING_DIMENSIONS }, (_, index) => {
const byte = seed[index % seed.length];
return byte / 255;
});
}
async getEmbeddings(input: string[]): Promise<Embedding[]> {
return input.map((content, i) => ({
index: i,
content,
embedding: this.embed(content),
}));
}
}
packages/backend/server/src/plugins/copilot/embedding/index.ts
+6 -4
View File
@@ -1,4 +1,6 @@
export { CopilotEmbeddingClientService, MockEmbeddingClient } from './client';
export { CopilotEmbeddingJob } from './job';
export type { Chunk, DocFragment } from './types';
export { EmbeddingClient } from './types';
export { NativeEmbeddingService } from './native';
export { CopilotRerankService } from './rerank';
export {
EMBEDDING_RERANK_RUNTIME,
type EmbeddingRerankRuntime,
} from './route-context';
packages/backend/server/src/plugins/copilot/embedding/job.ts
-674
View File
@@ -1,674 +0,0 @@
import { Injectable, Logger } from '@nestjs/common';
import {
BlobNotFound,
CallMetric,
CopilotContextFileNotSupported,
EventBus,
JobQueue,
mapAnyError,
OneDay,
OnEvent,
OnJob,
} from '../../../base';
import { DocReader } from '../../../core/doc';
import { WorkspaceBlobStorage } from '../../../core/storage';
import { readAllDocIdsFromWorkspaceSnapshot } from '../../../core/utils/blocksuite';
import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
import { CopilotEmbeddingClientService } from './client';
import type { Chunk, DocFragment, EmbeddingCallOptions } from './types';
import { EmbeddingClient } from './types';
@Injectable()
export class CopilotEmbeddingJob {
private readonly logger = new Logger(CopilotEmbeddingJob.name);
private readonly workspaceJobAbortController: Map<string, AbortController> =
new Map();
private supportEmbedding = false;
private client: EmbeddingClient | undefined;
constructor(
private readonly embeddingClients: CopilotEmbeddingClientService,
private readonly doc: DocReader,
private readonly event: EventBus,
private readonly models: Models,
private readonly queue: JobQueue,
private readonly storage: CopilotStorage,
private readonly workspaceStorage: WorkspaceBlobStorage
) {}
@OnEvent('config.init')
async onConfigInit() {
await this.setup();
}
@OnEvent('config.changed')
async onConfigChanged() {
await this.setup();
}
private async setup() {
this.supportEmbedding =
await this.models.copilotContext.checkEmbeddingAvailable();
if (this.supportEmbedding) {
this.client = await this.embeddingClients.refresh();
}
}
// public this client to allow overriding in tests
get embeddingClient() {
return this.client as EmbeddingClient;
}
@CallMetric('ai', 'addFileEmbeddingQueue')
async addFileEmbeddingQueue(
file: Jobs['copilot.embedding.files'],
options?: { priority?: number }
) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.files', file, {
priority: options?.priority,
});
}
@CallMetric('ai', 'addBlobEmbeddingQueue')
async addBlobEmbeddingQueue(blob: Jobs['copilot.embedding.blobs']) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.blobs', blob);
}
@OnEvent('workspace.doc.embedding')
async addDocEmbeddingQueue(
docs: Events['workspace.doc.embedding'],
options?: { contextId: string; priority: number }
) {
if (!this.supportEmbedding) return;
for (const { workspaceId, docId } of docs) {
const jobId = `workspace:embedding:${workspaceId}:${docId}`;
const job = await this.queue.get(jobId, 'copilot.embedding.docs');
// if the job exists and is older than 5 minute, remove it
if (job && job.timestamp + 5 * 60 * 1000 < Date.now()) {
this.logger.verbose(`Removing old embedding job ${jobId}`);
await this.queue.remove(jobId, 'copilot.embedding.docs');
}
await this.queue.add(
'copilot.embedding.docs',
{
contextId: options?.contextId,
workspaceId,
docId,
},
{
jobId: `workspace:embedding:${workspaceId}:${docId}`,
priority: options?.priority ?? 1,
timestamp: Date.now(),
}
);
}
}
@OnEvent('workspace.updated')
async onWorkspaceConfigUpdate({
id,
enableDocEmbedding,
}: Events['workspace.updated']) {
// trigger workspace embedding
this.event.emit('workspace.embedding', {
workspaceId: id,
enableDocEmbedding,
});
}
@OnEvent('workspace.embedding')
async addWorkspaceEmbeddingQueue({
workspaceId,
enableDocEmbedding,
}: Events['workspace.embedding']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
if (enableDocEmbedding === undefined) {
enableDocEmbedding =
await this.models.workspace.allowEmbedding(workspaceId);
}
if (enableDocEmbedding) {
const toBeEmbedDocIds =
await this.models.copilotWorkspace.findDocsToEmbed(workspaceId);
if (!toBeEmbedDocIds.length) {
return;
}
// filter out trashed docs
const rootSnapshot = await this.models.doc.getSnapshot(
workspaceId,
workspaceId
);
if (!rootSnapshot) {
this.logger.warn(
`Root snapshot for workspace ${workspaceId} not found, skipping embedding.`
);
return;
}
const allDocIds = new Set(
readAllDocIdsFromWorkspaceSnapshot(rootSnapshot.blob)
);
this.logger.log(
`Trigger embedding for ${toBeEmbedDocIds.length} docs in workspace ${workspaceId}`
);
const finalToBeEmbedDocIds = toBeEmbedDocIds.filter(docId =>
allDocIds.has(docId)
);
for (const docId of finalToBeEmbedDocIds) {
await this.queue.add(
'copilot.embedding.docs',
{
workspaceId,
docId,
},
{
jobId: `workspace:embedding:${workspaceId}:${docId}`,
priority: 1,
}
);
}
} else {
const controller = this.workspaceJobAbortController.get(workspaceId);
if (controller) {
controller.abort();
this.workspaceJobAbortController.delete(workspaceId);
}
}
}
@OnJob('copilot.embedding.updateDoc')
async addDocEmbeddingQueueFromEvent(
doc: Jobs['copilot.embedding.updateDoc']
) {
if (!this.supportEmbedding || !this.embeddingClient) return;
await this.queue.add(
'copilot.embedding.docs',
{
workspaceId: doc.workspaceId,
docId: doc.docId,
},
{
jobId: `workspace:embedding:${doc.workspaceId}:${doc.docId}`,
priority: 2,
}
);
}
private async deleteDocEmbedding(doc: {
workspaceId: string;
docId: string;
}) {
await this.queue.remove(
`workspace:embedding:${doc.workspaceId}:${doc.docId}`,
'copilot.embedding.docs'
);
await this.models.copilotContext.purgeWorkspaceEmbedding(
doc.workspaceId,
doc.docId
);
}
@OnJob('copilot.embedding.reconcileDocumentCleanup')
async reconcileDocumentCleanup({
workspaceId,
docId,
cleanupVersion,
}: Jobs['copilot.embedding.reconcileDocumentCleanup']) {
const root = await this.doc.getDoc(workspaceId, workspaceId);
if (!root) {
throw new Error(`workspace root ${workspaceId} not found`);
}
const live = readAllDocIdsFromWorkspaceSnapshot(root.bin, true).includes(
docId
);
if (live) {
if (!(await this.doc.getDoc(workspaceId, docId))) {
throw new Error(`restored document ${workspaceId}/${docId} not found`);
}
await this.addDocEmbeddingQueueFromEvent({ workspaceId, docId });
} else {
await this.deleteDocEmbedding({ workspaceId, docId });
}
await this.queue.add('backendRuntime.ackDocumentCleanupEffect', {
workspaceId,
docId,
cleanupVersion,
effect: 'copilot',
});
}
private async readCopilotBlob(
userId: string,
workspaceId: string,
blobId: string,
fileName: string
) {
const { body } = await this.storage.get(userId, workspaceId, blobId);
if (!body) throw new BlobNotFound({ spaceId: workspaceId, blobId });
const buffer = await readStream(body);
return new File([buffer], fileName);
}
private async readWorkspaceBlob(
workspaceId: string,
blobId: string,
fileName: string
) {
const { body } = await this.workspaceStorage.get(workspaceId, blobId);
if (!body) throw new BlobNotFound({ spaceId: workspaceId, blobId });
const buffer = await readStream(body);
return new File([buffer], fileName);
}
private workspaceIndexingOptions(
workspaceId: string,
signal?: AbortSignal,
userId?: string
): EmbeddingCallOptions {
return {
workspaceId,
userId,
signal,
featureKind: 'workspace_indexing',
};
}
@OnJob('copilot.embedding.files')
async embedPendingFile({
userId,
workspaceId,
contextId,
blobId,
fileId,
fileName,
}: Jobs['copilot.embedding.files']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
try {
const file = await this.readCopilotBlob(
userId,
workspaceId,
blobId,
fileName
);
// no need to check if embeddings is empty, will throw internally
const chunks = await this.embeddingClient.getFileChunks(file);
const total = chunks.reduce((acc, c) => acc + c.length, 0);
for (const chunk of chunks) {
const embeddings = await this.embeddingClient.generateEmbeddings(
chunk,
this.workspaceIndexingOptions(workspaceId, undefined, userId)
);
if (contextId) {
// for context files
await this.models.copilotContext.insertFileEmbedding(
contextId,
fileId,
embeddings
);
} else {
// for workspace files
await this.models.copilotWorkspace.insertFileEmbeddings(
workspaceId,
fileId,
embeddings
);
}
}
this.event.emit('workspace.file.embed.finished', {
contextId,
workspaceId,
fileId,
chunkSize: total,
});
} catch (error: any) {
this.event.emit('workspace.file.embed.failed', {
contextId,
workspaceId,
fileId,
error: mapAnyError(error).message,
});
// passthrough error to job queue
throw error;
}
}
@OnJob('copilot.embedding.blobs')
async embedPendingBlob({
workspaceId,
contextId,
blobId,
}: Jobs['copilot.embedding.blobs']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
try {
const file = await this.readWorkspaceBlob(workspaceId, blobId, 'blob');
const chunks = await this.embeddingClient.getFileChunks(file);
const total = chunks.reduce((acc, c) => acc + c.length, 0);
for (const chunk of chunks) {
const embeddings = await this.embeddingClient.generateEmbeddings(
chunk,
this.workspaceIndexingOptions(workspaceId)
);
await this.models.copilotWorkspace.insertBlobEmbeddings(
workspaceId,
blobId,
embeddings
);
}
if (contextId) {
this.event.emit('workspace.blob.embed.finished', {
contextId,
blobId,
chunkSize: total,
});
}
} catch (error: any) {
if (contextId) {
this.event.emit('workspace.blob.embed.failed', {
contextId,
blobId,
error: mapAnyError(error).message,
});
}
throw error;
}
}
private async getDocFragment(
workspaceId: string,
docId: string
): Promise<DocFragment | null> {
const docContent = await this.doc.getFullDocContent(workspaceId, docId);
const authors = await this.models.doc.getAuthors(workspaceId, docId);
if (docContent && authors) {
const { title, summary } = docContent;
const { createdAt, updatedAt, createdByUser, updatedByUser } = authors;
return {
title: title || 'Untitled',
summary,
createdAt: createdAt.toDateString(),
updatedAt: updatedAt.toDateString(),
createdBy: createdByUser?.name,
updatedBy: updatedByUser?.name,
};
}
return null;
}
private formatDocChunks(chunks: Chunk[], fragment: DocFragment): Chunk[] {
return chunks.map(chunk => ({
index: chunk.index,
content: [
`Title: ${fragment.title}`,
`Created at: ${fragment.createdAt}`,
`Updated at: ${fragment.updatedAt}`,
fragment.createdBy ? `Created by: ${fragment.createdBy}` : undefined,
fragment.updatedBy ? `Updated by: ${fragment.updatedBy}` : undefined,
chunk.content,
]
.filter(Boolean)
.join('\n'),
}));
}
private getWorkspaceSignal(workspaceId: string) {
let controller = this.workspaceJobAbortController.get(workspaceId);
if (!controller) {
controller = new AbortController();
this.workspaceJobAbortController.set(workspaceId, controller);
}
return controller.signal;
}
private normalize(s: string) {
return s.replaceAll(/[\p{White_Space}]+/gu, '');
}
@OnJob('copilot.embedding.docs')
async embedPendingDocs({
contextId,
workspaceId,
docId,
}: Jobs['copilot.embedding.docs']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
if (workspaceId === docId || docId.includes('$')) return;
const signal = this.getWorkspaceSignal(workspaceId);
try {
const hasNewDoc = await this.models.doc.exists(
workspaceId,
docId.split(':space:')[1] || ''
);
const needEmbedding =
await this.models.copilotWorkspace.checkDocNeedEmbedded(
workspaceId,
docId
);
this.logger.debug(
`Check if doc ${docId} in workspace ${workspaceId} needs embedding: ${needEmbedding}`
);
if (needEmbedding) {
if (signal.aborted) {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} is aborted, skipping embedding.`
);
return;
}
// if doc id deprecated, skip embedding and fulfill empty embedding
const fragment = !hasNewDoc
? await this.getDocFragment(workspaceId, docId)
: undefined;
if (!hasNewDoc && fragment) {
// fast fall for empty doc, journal is easily to create a empty doc
if (fragment.summary.trim()) {
const existsContent =
await this.models.copilotContext.getWorkspaceContent(
workspaceId,
docId
);
if (
existsContent &&
this.normalize(existsContent) === this.normalize(fragment.summary)
) {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no content change, skipping embedding.`
);
if (contextId) {
this.event.emit('workspace.doc.embed.finished', {
contextId,
docId,
});
}
return;
}
const embeddings = await this.embeddingClient.getFileEmbeddings(
new File(
[fragment.summary],
`${fragment.title || 'Untitled'}.md`
),
chunks => this.formatDocChunks(chunks, fragment),
this.workspaceIndexingOptions(workspaceId, signal)
);
for (const chunks of embeddings) {
await this.models.copilotContext.insertWorkspaceEmbedding(
workspaceId,
docId,
chunks
);
}
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has summary, embedding done.`
);
} else {
// for empty doc, insert empty embedding
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no summary, fulfilling empty embedding.`
);
await this.models.copilotContext.fulfillEmptyEmbedding(
workspaceId,
docId
);
}
} else {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no fragment, fulfilling empty embedding.`
);
await this.models.copilotContext.fulfillEmptyEmbedding(
workspaceId,
docId
);
}
}
if (contextId) {
this.event.emit('workspace.doc.embed.finished', {
contextId,
docId,
});
}
} catch (error: any) {
if (contextId) {
this.event.emit('workspace.doc.embed.failed', {
contextId,
docId,
});
}
if (
error instanceof CopilotContextFileNotSupported &&
error.message.includes('no content found')
) {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no content, fulfilling empty embedding.`
);
// if the doc is empty, we still need to fulfill the embedding
await this.models.copilotContext.fulfillEmptyEmbedding(
workspaceId,
docId
);
return;
}
// log error and skip the job
this.logger.error(
`Error embedding doc ${docId} in workspace ${workspaceId}`,
error
);
}
}
@OnJob('copilot.embedding.cleanupTrashedDocEmbeddings')
async cleanupTrashedDocEmbeddings({
workspaceId,
}: Jobs['copilot.embedding.cleanupTrashedDocEmbeddings']) {
const workspace = await this.models.workspace.get(workspaceId);
if (!workspace) {
this.logger.warn(`workspace ${workspaceId} not found`);
return;
}
const oneMonthAgo = new Date(Date.now() - OneDay * 30);
const snapshot = await this.models.doc.getSnapshot(
workspaceId,
workspaceId
);
if (!snapshot) {
// maybe local workspace or empty workspace
this.logger.verbose(`workspace root snapshot ${workspaceId} not found`);
// mark last check time to avoid repeated checking
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
return;
} else if (
// always check if never cleared
workspace.lastCheckEmbeddings > new Date(0) &&
snapshot.updatedAt < oneMonthAgo
) {
this.logger.verbose(
`workspace ${workspaceId} is too old, skipping embeddings cleanup`
);
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
return;
}
const [docIdsInEmbedding, docIdsInSnapshots] = await Promise.all([
this.models.copilotContext.listWorkspaceDocEmbedding(workspaceId),
this.models.copilotWorkspace.listEmbeddableDocIds(workspaceId),
]);
if (!docIdsInEmbedding.length && !docIdsInSnapshots.length) {
this.logger.verbose(
`No doc embeddings and snapshots found in workspace ${workspaceId}, skipping cleanup`
);
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
return;
}
const docIdsInWorkspace = readAllDocIdsFromWorkspaceSnapshot(snapshot.blob);
const docIdsInWorkspaceSet = new Set(docIdsInWorkspace);
const deletedDocIds = new Set(
[...docIdsInEmbedding, ...docIdsInSnapshots].filter(
docId => !docIdsInWorkspaceSet.has(docId)
)
);
for (const docId of deletedDocIds) {
const isPlaceholder = await this.models.copilotWorkspace.hasPlaceholder(
workspaceId,
docId
);
if (isPlaceholder) continue;
await this.models.copilotContext.deleteWorkspaceEmbedding(
workspaceId,
docId
);
}
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
}
@OnEvent('workspace.updated')
async onWorkspaceUpdated({ id }: Events['workspace.updated']) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.cleanupTrashedDocEmbeddings', {
workspaceId: id,
});
}
}
packages/backend/server/src/plugins/copilot/embedding/native.ts
+201
View File
@@ -0,0 +1,201 @@
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import { nanoid } from 'nanoid';
import { metrics } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import type { DocChunkSimilarity } from '../../../models';
import type {
RuntimeEmbeddingCandidate,
RuntimeRetrievalScope,
} from '../../../native';
import { CopilotRerankService } from './rerank';
import type { EmbeddingRouteContext } from './route-context';
@Injectable()
export class NativeEmbeddingService implements OnApplicationBootstrap {
private supportEmbedding = false;
constructor(
private readonly runtime: BackendRuntimeProvider,
private readonly rerank: CopilotRerankService
) {}
async onApplicationBootstrap() {
this.supportEmbedding = (await this.health()).enabled;
}
get canEmbedding() {
return this.supportEmbedding;
}
async health() {
const health = await this.runtime.embeddingHealth();
metrics.ai.counter('embedding_capability_check').add(1, {
state: health.state,
enabled: health.enabled,
reason: health.reason ?? 'none',
schema: String(health.schemaVersion ?? 0),
worker: health.workerRunning ? 'running' : 'stopped',
});
return health;
}
async progress(workspaceId: string) {
return await this.runtime.embeddingWorkspaceProgress(workspaceId);
}
async readSourceContent(
workspaceId: string,
sourceKind: 'document' | 'artifact',
sourceKey: string,
retrieval: RuntimeRetrievalScope,
maxChars?: number,
cursor?: string
) {
return await this.runtime.readEmbeddingSourceContent({
workspaceId,
sourceKind,
sourceKey,
retrieval,
maxChars,
cursor,
});
}
async match(
workspaceId: string,
query: string,
sourceKind: 'document' | 'artifact',
retrieval: RuntimeRetrievalScope,
limit: number,
signal?: AbortSignal
): Promise<RuntimeEmbeddingCandidate[]> {
const startedAt = performance.now();
signal?.throwIfAborted();
const requestId = nanoid();
const abort = () => {
void this.runtime
.cancelEmbeddingCandidateRequest(requestId)
.catch(() => {});
};
signal?.addEventListener('abort', abort, { once: true });
try {
const candidates = await this.runtime.matchEmbeddingCandidates({
requestId,
workspaceId,
query,
sourceKind,
retrieval,
limit,
});
signal?.throwIfAborted();
metrics.ai
.histogram('embedding_candidate_latency_ms')
.record(performance.now() - startedAt, {
corpus: sourceKind,
mode: retrieval.mode,
outcome: 'success',
});
return candidates;
} catch (error) {
metrics.ai.counter('embedding_operation_failure').add(1, {
operation: 'match',
kind: sourceKind,
code: embeddingErrorCode(error),
});
throw error;
} finally {
signal?.removeEventListener('abort', abort);
}
}
async matchWorkspaceDocCandidates(
workspaceId: string,
content: string,
topK = 5,
docIds?: string[]
): Promise<DocChunkSimilarity[]> {
const retrieval: RuntimeRetrievalScope = {
mode: docIds ? 'required' : 'workspace',
requiredDocIds: docIds ?? [],
requiredArtifactIds: [],
preferredSourceIds: [],
};
return (
await this.match(workspaceId, content, 'document', retrieval, topK * 2)
)
.filter(candidate => candidate.docId)
.map(candidate => ({
docId: candidate.docId as string,
chunk: candidate.chunk,
content: candidate.content,
distance: candidate.distance,
unitId: candidate.unitId ?? '',
visibility: (candidate.visibility ?? 'page') as
| 'page'
| 'edgeless'
| 'both',
blockId: candidate.blockId ?? undefined,
elementId: candidate.elementId ?? undefined,
frameId: candidate.frameId ?? undefined,
}));
}
async rerankWorkspaceDocs(
workspaceId: string,
content: string,
candidates: DocChunkSimilarity[],
topK = 5,
routeContext?: EmbeddingRouteContext
) {
if (!candidates.length) return [];
return await this.rerank.rerank(
content,
candidates,
topK,
workspaceId,
routeContext
);
}
async recordQueueCounts() {
const counts = await this.runtime.embeddingQueueCounts();
for (const status of [
'pending',
'running',
'retryWait',
'ready',
'failed',
] as const) {
metrics.ai
.gauge('embedding_queue_status')
.record(Number(counts[status]), { status });
}
metrics.ai
.gauge('embedding_vector_rows')
.record(Number(counts.activeVectorRows), { state: 'active' });
metrics.ai
.gauge('embedding_vector_rows')
.record(Number(counts.inactiveVectorRows), { state: 'inactive' });
metrics.ai
.gauge('embedding_index_size_bytes')
.record(Number(counts.indexBytes));
metrics.ai
.gauge('embedding_index_retry')
.record(Number(counts.retryingIndexes), { measure: 'indexes' });
metrics.ai
.gauge('embedding_index_retry')
.record(Number(counts.maxIndexRetrySeconds), {
measure: 'max_delay_seconds',
});
}
}
function embeddingErrorCode(error: unknown) {
if (!(error instanceof Error)) return 'unknown';
if (error.message.includes('resource_exceeded')) return 'resource_exceeded';
if (error.message.includes('embedding_unavailable')) return 'unavailable';
if (error.message.includes('not_found')) return 'not_found';
if (error.message.includes('disabled')) return 'disabled';
return 'failed';
}
packages/backend/server/src/plugins/copilot/embedding/realtime.ts
+57
View File
@@ -0,0 +1,57 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { Config } from '../../../base/config';
import { PermissionAccess } from '../../../core/permission';
import {
RealtimeRegistry,
realtimeWorkspaceEmbeddingProgressRoom,
registerRealtimeLiveQuery,
} from '../../../core/realtime';
import { assertCopilotEnabled } from '../availability';
import { NativeEmbeddingService } from './native';
@Injectable()
export class CopilotEmbeddingRealtimeProvider implements OnModuleInit {
constructor(
private readonly ac: PermissionAccess,
private readonly embedding: NativeEmbeddingService,
private readonly registry: RealtimeRegistry,
private readonly config: Config
) {}
onModuleInit() {
const input = z.object({ workspaceId: z.string() });
registerRealtimeLiveQuery(this.registry, {
request: {
name: 'workspace.embedding.progress.get',
input,
handle: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
const health = await this.embedding.health();
return health.enabled
? await this.embedding.progress(payload.workspaceId)
: { total: 0, embedded: 0 };
},
},
topic: {
name: 'workspace.embedding.progress.changed',
input,
authorize: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
},
room: (_user, payload) =>
realtimeWorkspaceEmbeddingProgressRoom(payload.workspaceId),
},
});
}
private async assertCopilot(userId: string, workspaceId: string) {
assertCopilotEnabled(this.config);
await this.ac
.user(userId)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Copilot');
}
}
packages/backend/server/src/plugins/copilot/embedding/rerank.ts
+69
View File
@@ -0,0 +1,69 @@
import { Inject, Injectable } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import type { ChunkSimilarity } from '../../../models';
import {
EMBEDDING_RERANK_RUNTIME,
type EmbeddingRerankRuntime,
type EmbeddingRouteContext,
} from './route-context';
@Injectable()
export class CopilotRerankService {
constructor(
@Inject(ModuleRef)
private readonly moduleRef: ModuleRef
) {}
async rerank<T extends ChunkSimilarity>(
query: string,
candidates: T[],
topK: number,
workspaceId: string,
routeContext: EmbeddingRouteContext = {},
signal?: AbortSignal
): Promise<T[]> {
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
if (!candidates.length) return [];
try {
const runtime = this.moduleRef.get<EmbeddingRerankRuntime>(
EMBEDDING_RERANK_RUNTIME,
{ strict: false }
);
const scores = await runtime.rerank(
'route-selected',
{
query,
candidates: candidates.map((candidate, index) => ({
id: String(index),
text: candidate.content,
})),
},
{
workspace: workspaceId,
byokLeaseId: routeContext.byokLeaseId,
featureKind: 'rerank',
signal,
}
);
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
if (scores.length !== candidates.length) {
return candidates
.toSorted(
(a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity)
)
.slice(0, topK);
}
return candidates
.map((candidate, index) => ({ candidate, score: scores[index] }))
.toSorted((a, b) => b.score - a.score)
.slice(0, topK)
.map(item => item.candidate);
} catch (error) {
if (signal?.aborted) throw new Error('SEARCH_ABORTED', { cause: error });
return candidates
.toSorted((a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity))
.slice(0, topK);
}
}
}
packages/backend/server/src/plugins/copilot/embedding/route-context.ts
+18
View File
@@ -0,0 +1,18 @@
export type EmbeddingRouteContext = {
byokLeaseId?: string;
};
export const EMBEDDING_RERANK_RUNTIME = Symbol('EMBEDDING_RERANK_RUNTIME');
export interface EmbeddingRerankRuntime {
rerank(
modelId: string,
request: { query: string; candidates: { id: string; text: string }[] },
options: {
workspace: string;
byokLeaseId?: string;
featureKind: 'rerank';
signal?: AbortSignal;
}
): Promise<number[]>;
}
packages/backend/server/src/plugins/copilot/embedding/types.ts
-252
View File
@@ -1,252 +0,0 @@
import { File } from 'node:buffer';
import { z } from 'zod';
import { CopilotContextFileNotSupported } from '../../../base';
import type { PageDocContent } from '../../../core/utils/blocksuite';
import { ChunkSimilarity, Embedding } from '../../../models';
import { parseDoc } from '../../../native';
import type { ByokFeatureKind } from '../byok/types';
declare global {
interface Events {
'workspace.embedding': {
workspaceId: string;
enableDocEmbedding?: boolean;
};
'workspace.blob.embed.finished': {
contextId: string;
blobId: string;
chunkSize: number;
};
'workspace.blob.embed.failed': {
contextId: string;
blobId: string;
error: string;
};
'workspace.doc.embedding': Array<{
workspaceId: string;
docId: string;
}>;
'workspace.doc.embed.failed': {
contextId: string;
docId: string;
};
'workspace.doc.embed.finished': {
contextId: string;
docId: string;
};
'workspace.file.embed.finished': {
contextId?: string;
workspaceId: string;
fileId: string;
chunkSize: number;
};
'workspace.file.embed.failed': {
contextId?: string;
workspaceId: string;
fileId: string;
error: string;
};
}
interface Jobs {
'copilot.embedding.docs': {
contextId?: string;
workspaceId: string;
docId: string;
};
'copilot.embedding.updateDoc': {
workspaceId: string;
docId: string;
};
'copilot.embedding.reconcileDocumentCleanup': {
workspaceId: string;
docId: string;
cleanupVersion: string;
};
'copilot.embedding.files': {
contextId?: string;
userId: string;
workspaceId: string;
blobId: string;
fileId: string;
fileName: string;
};
'copilot.embedding.blobs': {
contextId?: string;
workspaceId: string;
blobId: string;
};
'copilot.embedding.cleanupTrashedDocEmbeddings': {
workspaceId: string;
};
}
}
export type DocFragment = PageDocContent & {
createdAt: string;
createdBy?: string;
updatedAt: string;
updatedBy?: string;
};
export type Chunk = {
index: number;
content: string;
};
export type EmbeddingCallOptions = {
signal?: AbortSignal;
userId?: string;
workspaceId?: string;
byokLeaseId?: string;
featureKind?: Extract<
ByokFeatureKind,
'embedding' | 'workspace_indexing' | 'rerank'
>;
};
export type EmbeddingCallOptionsInput = AbortSignal | EmbeddingCallOptions;
export type EmbeddingRouteContext = Pick<
EmbeddingCallOptions,
'userId' | 'byokLeaseId'
>;
export function normalizeEmbeddingCallOptions(
options?: EmbeddingCallOptionsInput
): EmbeddingCallOptions {
if (!options) {
return {};
}
if ('aborted' in options && 'addEventListener' in options) {
return { signal: options };
}
return options;
}
export abstract class EmbeddingClient {
async configured() {
return true;
}
async getFileEmbeddings(
file: File,
chunkMapper: (chunk: Chunk[]) => Chunk[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[][]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const chunks = await this.getFileChunks(file, normalizedOptions.signal);
const chunkedEmbeddings = await Promise.all(
chunks.map(chunk =>
this.generateEmbeddings(chunkMapper(chunk), normalizedOptions)
)
);
return chunkedEmbeddings;
}
async getFileChunks(file: File, signal?: AbortSignal): Promise<Chunk[][]> {
const buffer = Buffer.from(await file.arrayBuffer());
let doc;
try {
doc = await parseDoc(file.name, buffer);
} catch (e: any) {
throw new CopilotContextFileNotSupported({
fileName: file.name,
message: e?.message || e?.toString?.() || 'format not supported',
});
}
if (doc && !signal?.aborted) {
if (!doc.chunks.length) {
throw new CopilotContextFileNotSupported({
fileName: file.name,
message: 'no content found',
});
}
const input = doc.chunks.toSorted((a, b) => a.index - b.index);
// chunk input into 128 every array
const chunks: Chunk[][] = [];
for (let i = 0; i < input.length; i += 128) {
chunks.push(input.slice(i, i + 128));
}
return chunks;
}
throw new CopilotContextFileNotSupported({
fileName: file.name,
message: 'failed to parse file',
});
}
async generateEmbeddings(
chunks: Chunk[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const retry = 3;
let embeddings: Embedding[] = [];
let error = null;
for (let i = 0; i < retry; i++) {
try {
embeddings = await this.getEmbeddings(
chunks.map(c => c.content),
normalizedOptions
);
break;
} catch (e) {
error = e;
}
}
if (error) throw error;
// fix the index of the embeddings
return embeddings.map(e => ({ ...e, index: chunks[e.index].index }));
}
async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
_query: string,
embeddings: Chunk[],
topK: number,
_options?: EmbeddingCallOptionsInput
): Promise<Chunk[]> {
// sort by distance with ascending order
return embeddings
.toSorted((a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity))
.slice(0, topK);
}
async getEmbedding(query: string, options?: EmbeddingCallOptionsInput) {
const embedding = await this.getEmbeddings([query], options);
return embedding?.[0]?.embedding;
}
abstract getEmbeddings(
input: string[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]>;
}
const ReRankItemSchema = z.object({
chunk: z.number().describe('The chunk index of the search result.'),
targetId: z.string().describe('The id of the target.'),
score: z
.number()
.min(0)
.max(10)
.describe(
'The relevance score of the results should be 0-10, with 0 being the least relevant and 10 being the most relevant.'
),
});
export type ReRankResult = z.infer<typeof ReRankItemSchema>[];
packages/backend/server/src/plugins/copilot/index.ts
+1 -8
View File
@@ -17,7 +17,6 @@ import { McpCredentialService } from './mcp/credential';
import { McpCredentialResolver } from './mcp/resolver';
import {
COPILOT_API_PROVIDERS,
COPILOT_CONTEXT_REALTIME_PROVIDERS,
COPILOT_FEATURE_PROVIDERS,
COPILOT_KERNEL_PROVIDERS,
COPILOT_TRANSCRIPT_REALTIME_PROVIDERS,
@@ -49,17 +48,11 @@ export class CopilotAvailabilityModule {}
export class CopilotKernelModule {}
@Module({
imports: [PermissionModule, CopilotAvailabilityModule],
imports: [PermissionModule, CopilotAvailabilityModule, CopilotKernelModule],
providers: [...COPILOT_TRANSCRIPT_REALTIME_PROVIDERS],
})
export class CopilotRealtimeModule {}
@Module({
imports: [PermissionModule, CopilotAvailabilityModule],
providers: [...COPILOT_CONTEXT_REALTIME_PROVIDERS],
})
export class CopilotEmbeddingRealtimeModule {}
@Module({
imports: [...COPILOT_SHARED_IMPORTS, CopilotKernelModule],
providers: [...COPILOT_FEATURE_PROVIDERS],
packages/backend/server/src/plugins/copilot/mcp/provider.ts
+39 -91
View File
@@ -1,13 +1,10 @@
import { Injectable } from '@nestjs/common';
import { McpAccessMode } from '@prisma/client';
import { pick } from 'lodash-es';
import z from 'zod/v3';
import { DocReader, DocWriter } from '../../../core/doc';
import { PermissionAccess } from '../../../core/permission';
import { clearEmbeddingChunk } from '../../../models';
import { IndexerService } from '../../indexer';
import { CopilotContextService } from '../context/service';
import { DocumentRetrievalService } from '../retrieval/document';
type McpTextContent = {
type: 'text';
@@ -103,8 +100,7 @@ export class WorkspaceMcpProvider {
private readonly ac: PermissionAccess,
private readonly reader: DocReader,
private readonly writer: DocWriter,
private readonly context: CopilotContextService,
private readonly indexer: IndexerService
private readonly retrieval: DocumentRetrievalService
) {}
async for(
@@ -154,105 +150,57 @@ export class WorkspaceMcpProvider {
},
});
const semanticSearch = defineTool({
name: 'semantic_search',
title: 'Semantic Search',
const docSearch = defineTool({
name: 'doc_search',
title: 'Document Search',
description:
'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts, recent documents).',
parser: z.object({ query: z.string() }),
'Search persisted workspace documents and return bounded passages with Page or canvas locators. Retrieval strategy is selected by the server and never includes files, blobs, attachments, or the web.',
parser: z.object({
query: z.string().trim().min(1).max(2000),
doc_ids: z.array(z.string().min(1).max(128)).max(50).optional(),
limit: z.number().int().min(1).max(20).optional(),
}),
inputSchema: {
type: 'object',
properties: {
query: { type: 'string' },
doc_ids: {
type: 'array',
items: { type: 'string' },
maxItems: 50,
},
limit: { type: 'integer', minimum: 1, maximum: 20 },
},
required: ['query'],
additionalProperties: false,
},
execute: async ({ query }, options) => {
const trimmed = query.trim();
if (!trimmed) {
return toolError('Query is required for semantic search.');
}
const chunks = await this.context.matchWorkspaceDocs(
workspaceId,
trimmed,
5,
execute: async ({ query, doc_ids, limit }, options) => {
const result = await this.retrieval.search(
{ user: userId, workspace: workspaceId },
query,
doc_ids,
limit ?? 10,
options.signal
);
const abortedAfterMatch = abortIfNeeded(options.signal);
if (abortedAfterMatch) return abortedAfterMatch;
const docs = await this.ac
.user(userId)
.workspace(workspaceId)
.docs(
chunks.filter(chunk => 'docId' in chunk),
'Doc.Read'
);
const abortedAfterDocs = abortIfNeeded(options.signal);
if (abortedAfterDocs) return abortedAfterDocs;
if (!docs || docs.length === 0) {
return toolText('No matching documents found.');
}
return {
content: docs.map(doc => ({
type: 'text',
text: clearEmbeddingChunk(doc).content,
})),
};
return toolText(
JSON.stringify({
retrieval_mode: result.retrievalMode,
degraded_reason: result.degradedReason,
hits: result.hits.map(hit => ({
doc_id: hit.docId,
title: hit.title,
excerpt: hit.excerpt,
visibility: hit.visibility,
block_id: hit.blockId,
element_id: hit.elementId,
frame_id: hit.frameId,
})),
})
);
},
});
const keywordSearch = defineTool({
name: 'keyword_search',
title: 'Keyword Search',
description:
'Fuzzy search all workspace documents for the exact keyword or phrase supplied and return passages ranked by textual match. Use this tool by default whenever a straightforward term-based or keyword-base lookup is sufficient.',
parser: z.object({ query: z.string() }),
inputSchema: {
type: 'object',
properties: {
query: { type: 'string' },
},
required: ['query'],
additionalProperties: false,
},
execute: async ({ query }, options) => {
const trimmed = query.trim();
if (!trimmed) return toolError('Query is required for keyword search.');
let docs = await this.indexer.searchDocsByKeyword(workspaceId, trimmed);
const abortedAfterSearch = abortIfNeeded(options.signal);
if (abortedAfterSearch) return abortedAfterSearch;
docs = await this.ac
.user(userId)
.workspace(workspaceId)
.docs(docs, 'Doc.Read');
const abortedAfterDocs = abortIfNeeded(options.signal);
if (abortedAfterDocs) return abortedAfterDocs;
if (!docs || docs.length === 0) {
return toolText('No matching documents found.');
}
return {
content: docs.map(doc => ({
type: 'text',
text: JSON.stringify(pick(doc, 'docId', 'title', 'createdAt')),
})),
};
},
});
const tools = [readDocument, semanticSearch, keywordSearch];
const tools = [readDocument, docSearch];
if (
accessMode === McpAccessMode.READ_WRITE &&
packages/backend/server/src/plugins/copilot/module-providers.ts
+21 -22
View File
@@ -4,23 +4,26 @@ import { CompatHistoryProjector } from './compat/history-projector';
import { HistoryPromptPreloadProjector } from './compat/history-prompt-preload-projector';
import { HistoryVisibilityPolicy } from './compat/history-visibility-policy';
import { CompatSubmissionStore } from './compat/submission-store';
import {
CopilotContextResolver,
CopilotContextRootResolver,
CopilotContextService,
CopilotEmbeddingRealtimeProvider,
} from './context';
import { ConversationInboxService } from './conversation/inbox';
import { ConversationPolicy } from './conversation/policy';
import { ConversationStore } from './conversation/store';
import { CopilotCronJobs } from './cron';
import { DelegatedEditorRealtimeProvider } from './delegated/realtime';
import { DelegatedEditorService } from './delegated/service';
import {
CopilotEmbeddingClientService,
CopilotEmbeddingJob,
CopilotRerankService,
EMBEDDING_RERANK_RUNTIME,
NativeEmbeddingService,
} from './embedding';
import { CopilotEmbeddingRealtimeProvider } from './embedding/realtime';
import { WorkspaceMcpProvider } from './mcp/provider';
import { PromptService } from './prompt';
import { CopilotResolver, UserCopilotResolver } from './resolver';
import { ArtifactRetrievalService } from './retrieval/artifact';
import {
DOCUMENT_VECTOR_SEARCH,
DocumentRetrievalService,
} from './retrieval/document';
import { ActionRuntimeBridge } from './runtime/action-runtime-bridge';
import { CapabilityRuntime } from './runtime/capability-runtime';
import { CopilotRuntimeEventConsumer } from './runtime/copilot-runtime-event-consumer';
@@ -61,14 +64,19 @@ export const COPILOT_RUNTIME_PROVIDERS = [
HistoryPromptPreloadProjector,
CompatSubmissionStore,
HistoryVisibilityPolicy,
CopilotContextService,
CopilotEmbeddingClientService,
NativeEmbeddingService,
CopilotRerankService,
PromptService,
{ provide: DOCUMENT_VECTOR_SEARCH, useExisting: NativeEmbeddingService },
DocumentRetrievalService,
ArtifactRetrievalService,
DelegatedEditorService,
ActionRuntimeBridge,
CopilotRuntimeEventConsumer,
PromptRuntime,
ConversationHost,
CapabilityRuntime,
{ provide: EMBEDDING_RERANK_RUNTIME, useExisting: CapabilityRuntime },
ToolRuntime,
AttachmentMaterializer,
AttachmentAdmissionHost,
@@ -79,18 +87,11 @@ export const COPILOT_RUNTIME_PROVIDERS = [
TurnPersistence,
];
export const COPILOT_CONTEXT_REALTIME_PROVIDERS = [
CopilotEmbeddingRealtimeProvider,
];
export const COPILOT_CONTEXT_PROVIDERS = [
CopilotContextResolver,
...COPILOT_CONTEXT_REALTIME_PROVIDERS,
];
export const COPILOT_TRANSCRIPT_REALTIME_PROVIDERS = [
CopilotTranscriptionReader,
CopilotTranscriptRealtimeProvider,
CopilotEmbeddingRealtimeProvider,
DelegatedEditorRealtimeProvider,
];
export const COPILOT_TRANSCRIPT_PROVIDERS = [
@@ -108,11 +109,10 @@ export const COPILOT_WORKSPACE_PROVIDERS = [
export const COPILOT_RESOLVER_PROVIDERS = [
CopilotResolver,
UserCopilotResolver,
CopilotContextRootResolver,
WorkspaceByokResolver,
];
export const COPILOT_JOB_PROVIDERS = [CopilotEmbeddingJob, CopilotCronJobs];
export const COPILOT_JOB_PROVIDERS = [CopilotCronJobs];
export const COPILOT_MCP_PROVIDERS = [WorkspaceMcpProvider];
@@ -123,7 +123,6 @@ export const COPILOT_KERNEL_PROVIDERS = [
export const COPILOT_FEATURE_PROVIDERS = [
TurnOrchestrator,
...COPILOT_CONTEXT_PROVIDERS,
...COPILOT_TRANSCRIPT_PROVIDERS,
...COPILOT_WORKSPACE_PROVIDERS,
...COPILOT_JOB_PROVIDERS,
packages/backend/server/src/plugins/copilot/providers/types.ts
+10 -4
View File
@@ -11,6 +11,7 @@ import {
type StreamObject,
StreamObjectSchema,
} from '../runtime/contracts/runtime-event-contract';
import { RetrievalScopeSchema } from '../runtime/contracts/shared';
// Owner map:
// - provider/profile/config schemas in this file are backend host ingress.
@@ -74,17 +75,21 @@ export const VertexSchema: JSONSchema = {
export const PromptToolsSchema = z
.enum([
'blobRead',
'artifactRead',
'artifactSearch',
'codeArtifact',
'conversationSummary',
// work with indexer
'docRead',
'docCanvasRead',
'docSearch',
'docCreate',
'docUpdate',
'docUpdateMeta',
'docKeywordSearch',
// work with embeddings
'docSemanticSearch',
'frontendGetEditorState',
'frontendReadSelection',
'frontendReadNodes',
'frontendSnapshotDocument',
// work with exa/model internal tools
'webSearch',
// artifact tools
@@ -280,6 +285,7 @@ const CopilotProviderOptionsSchema = z.object({
'transcript',
])
.optional(),
retrievalScope: RetrievalScopeSchema.optional(),
});
export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
packages/backend/server/src/plugins/copilot/providers/utils.ts
+7 -35
View File
@@ -174,8 +174,8 @@ export class TextStreamParser {
result += `\nCrawling the web "${chunk.input.url}"\n`;
break;
}
case 'doc_keyword_search': {
result += `\nSearching the keyword "${chunk.input.query}"\n`;
case 'doc_search': {
result += `\nSearching workspace documents for "${chunk.input.query}"\n`;
break;
}
case 'doc_read': {
@@ -196,27 +196,11 @@ export class TextStreamParser {
);
result = this.addPrefix(result);
switch (chunk.toolName) {
case 'doc_semantic_search': {
const output = chunk.output;
if (Array.isArray(output)) {
result += `\nFound ${output.length} document${output.length !== 1 ? 's' : ''} related to “${chunk.input.query}”.\n`;
} else if (typeof output === 'string') {
result += `\n${output}\n`;
} else {
const message = asRecord(output)?.message;
this.logger.warn(
`Unexpected result type for doc_semantic_search: ${
typeof message === 'string' ? message : 'Unknown error'
}`
);
}
break;
}
case 'doc_keyword_search': {
const output = chunk.output;
if (Array.isArray(output)) {
result += `\nFound ${output.length} document${output.length !== 1 ? 's' : ''} related to “${chunk.input.query}”.\n`;
result += `\n${this.getKeywordSearchLinks(output)}\n`;
case 'doc_search': {
const output = asRecord(chunk.output);
const hits = output?.hits;
if (Array.isArray(hits)) {
result += `\nFound ${hits.length} document${hits.length !== 1 ? 's' : ''} related to “${chunk.input.query}”.\n`;
}
break;
}
@@ -287,18 +271,6 @@ export class TextStreamParser {
}, '');
return links;
}
private getKeywordSearchLinks(
list: {
docId: string;
title: string;
}[]
): string {
const links = list.reduce((acc, result) => {
return acc + `\n\n[${result.title}](${result.docId})\n\n`;
}, '');
return links;
}
}
export class StreamObjectParser {
packages/backend/server/src/plugins/copilot/resolver.ts
+3
View File
@@ -236,6 +236,9 @@ class ChatMessageType implements Partial<ChatMessage> {
@Field(() => GraphQLJSON, { nullable: true })
params!: Record<string, string> | undefined;
@Field(() => GraphQLJSON, { nullable: true })
scopeSnapshot!: ChatMessage['scopeSnapshot'];
@Field(() => Date)
createdAt!: Date;
}
packages/backend/server/src/plugins/copilot/retrieval/artifact.ts
+187
View File
@@ -0,0 +1,187 @@
import { Injectable } from '@nestjs/common';
import { PrismaClient } from '@prisma/client';
import { AccessDenied } from '../../../base';
import { PermissionAccess } from '../../../core/permission';
import type { RuntimeRetrievalScope } from '../../../native';
import { NativeEmbeddingService } from '../embedding/native';
@Injectable()
export class ArtifactRetrievalService {
constructor(
private readonly access: PermissionAccess,
private readonly embedding: NativeEmbeddingService,
private readonly db: PrismaClient
) {}
private async authorize(userId: string, workspaceId: string) {
return await this.access
.user(userId)
.workspace(workspaceId)
.allowLocal()
.can('Workspace.Read');
}
async search(options: {
userId: string;
workspaceId: string;
query: string;
retrieval: RuntimeRetrievalScope;
limit: number;
messageId?: string;
signal?: AbortSignal;
}) {
if (!(await this.authorize(options.userId, options.workspaceId))) {
throw new AccessDenied();
}
let degraded = false;
let matched: Awaited<ReturnType<NativeEmbeddingService['match']>> = [];
try {
matched = await this.embedding.match(
options.workspaceId,
options.query,
'artifact',
options.retrieval,
options.limit,
options.signal
);
} catch (error) {
if (options.signal?.aborted) throw error;
degraded = true;
}
const matchedIds = new Set(matched.map(hit => hit.artifactId));
const missingRequired =
options.retrieval.mode === 'required'
? options.retrieval.requiredArtifactIds
.filter(id => !matchedIds.has(id))
.slice(0, Math.max(0, options.limit - matched.length))
: [];
const directAttempts = await Promise.allSettled(
missingRequired.map(async artifactId => {
const source = await this.embedding.readSourceContent(
options.workspaceId,
'artifact',
artifactId,
options.retrieval,
20_000
);
return {
sourceKind: 'artifact',
sourceKey: artifactId,
artifactId,
content: source.content,
distance: 0,
chunk: 0,
};
})
);
const direct = directAttempts.flatMap(result =>
result.status === 'fulfilled' ? [result.value] : []
);
degraded ||= direct.length !== directAttempts.length;
const hits = [...matched, ...direct].map(hit => ({
...hit,
artifactId: hit.artifactId ?? hit.sourceKey,
}));
const metadata = await this.loadMetadata(
options.workspaceId,
hits.map(hit => hit.artifactId),
options.retrieval,
options.messageId
);
return {
hits: hits.map(hit => ({ ...hit, ...metadata.get(hit.artifactId) })),
degraded,
} as const;
}
async read(options: {
userId: string;
workspaceId: string;
artifactId: string;
retrieval: RuntimeRetrievalScope;
messageId?: string;
maxChars?: number;
cursor?: string;
}) {
if (!(await this.authorize(options.userId, options.workspaceId))) {
throw new AccessDenied();
}
const [result, metadata] = await Promise.all([
this.embedding.readSourceContent(
options.workspaceId,
'artifact',
options.artifactId,
options.retrieval,
options.maxChars,
options.cursor
),
this.loadMetadata(
options.workspaceId,
[options.artifactId],
options.retrieval,
options.messageId
),
]);
return {
...result,
name: metadata.get(options.artifactId)?.name ?? result.name,
mimeType: metadata.get(options.artifactId)?.mimeType ?? result.mimeType,
};
}
private async loadMetadata(
workspaceId: string,
artifactIds: string[],
retrieval: RuntimeRetrievalScope,
messageId?: string
) {
const ids = [...new Set(artifactIds)];
if (!ids.length) {
return new Map<string, { name?: string; mimeType: string }>();
}
const [artifacts, occurrences] = await Promise.all([
this.db.workspaceArtifact.findMany({
where: {
workspaceId,
id: { in: ids },
...(retrieval.mode === 'workspace' ? { libraryOwned: true } : {}),
},
select: {
id: true,
displayName: true,
canonicalMediaType: true,
},
}),
retrieval.mode === 'required' && messageId
? this.db.aiMessageArtifact.findMany({
where: {
workspaceId,
messageId,
artifactId: { in: ids },
role: 'attachment',
},
select: { artifactId: true, displayName: true },
})
: [],
]);
const occurrenceNames = new Map(
occurrences.map(occurrence => [
occurrence.artifactId,
occurrence.displayName ?? undefined,
])
);
return new Map(
artifacts.map(artifact => [
artifact.id,
{
name:
occurrenceNames.get(artifact.id) ??
artifact.displayName ??
undefined,
mimeType: artifact.canonicalMediaType,
},
])
);
}
}
packages/backend/server/src/plugins/copilot/retrieval/document.ts
+264
View File
@@ -0,0 +1,264 @@
import { Inject, Injectable } from '@nestjs/common';
import { Config, SearchProviderNotFound } from '../../../base';
import { PermissionAccess } from '../../../core/permission';
import type { DocVisibility } from '../../../core/utils/blocksuite';
import { type DocChunkSimilarity, Models } from '../../../models';
import { IndexerService } from '../../indexer/service';
import type { SearchDoc } from '../../indexer/types';
import type { EmbeddingRouteContext } from '../embedding/route-context';
type DocumentSearchContext =
| {
user?: string;
workspace?: string;
byokLeaseId?: string;
}
| undefined;
type DocumentVectorSearch = {
readonly canEmbedding: boolean;
matchWorkspaceDocCandidates(
workspaceId: string,
content: string,
topK?: number,
docIds?: string[]
): Promise<DocChunkSimilarity[]>;
rerankWorkspaceDocs(
workspaceId: string,
content: string,
candidates: DocChunkSimilarity[],
topK?: number,
routeContext?: EmbeddingRouteContext
): Promise<DocChunkSimilarity[]>;
};
export const DOCUMENT_VECTOR_SEARCH = Symbol('DOCUMENT_VECTOR_SEARCH');
export type DocumentSearchHit = {
docId: string;
title: string;
excerpt: string;
visibility: DocVisibility;
blockId?: string;
elementId?: string;
frameId?: string;
updatedAt?: Date;
score: number;
unitId: string;
};
type Candidate = DocumentSearchHit & { channels: Set<'lexical' | 'vector'> };
type ProjectedSearchDoc = SearchDoc &
Required<
Pick<
SearchDoc,
'unitId' | 'projectionVersion' | 'sourceHash' | 'visibility'
>
>;
function hasProjectionMetadata(hit: SearchDoc): hit is ProjectedSearchDoc {
return Boolean(
hit.unitId && hit.projectionVersion && hit.sourceHash && hit.visibility
);
}
function hasVectorProjectionMetadata(hit: DocChunkSimilarity) {
return Boolean(hit.unitId && hit.visibility);
}
@Injectable()
export class DocumentRetrievalService {
constructor(
private readonly config: Config,
private readonly ac: PermissionAccess,
private readonly indexer: IndexerService,
@Inject(DOCUMENT_VECTOR_SEARCH)
private readonly context: DocumentVectorSearch,
private readonly models: Models
) {}
async search(
options: DocumentSearchContext,
query: string,
docIds: string[] | undefined,
requestedLimit: number,
signal?: AbortSignal
) {
if (!options?.user || !options.workspace) {
throw new Error('INVALID_SEARCH_CONTEXT');
}
const userId = options.user;
const workspaceId = options.workspace;
const limit = Math.min(requestedLimit, 20);
const routeContext = {
userId,
byokLeaseId: options.byokLeaseId,
};
const [lexicalAttempt, vectorAttempt] = await Promise.allSettled([
this.config.indexer.enabled
? this.indexer
.searchDocsByKeyword(workspaceId, query, {
limit: Math.max(limit * 3, 20),
docIds,
})
.catch(error => {
if (error instanceof SearchProviderNotFound) return null;
throw error;
})
: null,
this.context.canEmbedding
? this.context.matchWorkspaceDocCandidates(
workspaceId,
query,
Math.max(limit * 3, 20),
docIds
)
: null,
]);
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
const lexicalResult =
lexicalAttempt.status === 'fulfilled' ? lexicalAttempt.value : null;
const vectorResult =
vectorAttempt.status === 'fulfilled' ? vectorAttempt.value : null;
const lexical = lexicalResult
? await this.readable(
userId,
workspaceId,
lexicalResult.filter(hasProjectionMetadata)
)
: [];
const vectorScoped = (vectorResult ?? []).filter(
candidate =>
hasVectorProjectionMetadata(candidate) &&
(!docIds || docIds.includes(candidate.docId))
);
const readableVector = vectorScoped.length
? await this.readable(userId, workspaceId, vectorScoped)
: [];
let vector = null;
if (vectorResult !== null) {
try {
vector = await this.context.rerankWorkspaceDocs(
workspaceId,
query,
readableVector,
Math.max(limit * 3, 20),
routeContext
);
} catch {
vector = null;
}
}
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
const metas = await this.models.doc.findMetas(
(vector ?? []).map(candidate => ({
workspaceId,
docId: candidate.docId,
})),
{ select: { title: true } }
);
const metaByDoc = new Map(
metas
.filter((meta): meta is NonNullable<typeof meta> => meta !== null)
.map(meta => [meta.docId, meta])
);
const candidates = new Map<string, Candidate>();
const merge = (
hit: DocumentSearchHit,
channel: 'lexical' | 'vector',
rank: number
) => {
const key = `${hit.docId}:${hit.unitId}`;
const score = 1 / (60 + rank);
const existing = candidates.get(key);
if (existing) {
existing.score += score;
existing.channels.add(channel);
} else {
candidates.set(key, { ...hit, score, channels: new Set([channel]) });
}
};
lexical.forEach((hit, index) =>
merge(this.fromLexical(hit), 'lexical', index + 1)
);
vector?.forEach((hit, index) => {
const meta = metaByDoc.get(hit.docId);
merge(
{
docId: hit.docId,
title: meta?.title ?? '',
excerpt: hit.content,
visibility: hit.visibility as DocVisibility,
blockId: hit.blockId,
elementId: hit.elementId,
frameId: hit.frameId,
score: 0,
unitId: hit.unitId,
},
'vector',
index + 1
);
});
if (lexicalResult === null && vector === null) {
throw new Error('SEARCH_UNAVAILABLE');
}
const perDoc = new Map<string, number>();
const hits = [...candidates.values()]
.sort(
(left, right) =>
right.score - left.score || left.unitId.localeCompare(right.unitId)
)
.filter(hit => {
const count = perDoc.get(hit.docId) ?? 0;
if (count >= 3) return false;
perDoc.set(hit.docId, count + 1);
return true;
})
.slice(0, limit)
.map(({ channels: _, ...hit }) => hit);
const hasLexical = lexicalResult !== null;
const retrievalMode =
hasLexical && vector ? 'hybrid' : hasLexical ? 'lexical' : 'vector';
return {
retrievalMode,
degradedReason:
retrievalMode === 'hybrid'
? undefined
: lexicalResult
? 'VECTOR_UNAVAILABLE'
: 'LEXICAL_UNAVAILABLE',
hits,
} as const;
}
private async readable<T extends { docId: string }>(
userId: string,
workspaceId: string,
candidates: T[]
) {
return (
(await this.ac
.user(userId)
.workspace(workspaceId)
.docs(candidates, 'Doc.Read')) ?? []
);
}
private fromLexical(hit: ProjectedSearchDoc): DocumentSearchHit {
return {
docId: hit.docId,
title: hit.title,
excerpt: hit.highlight || '',
visibility: hit.visibility as DocVisibility,
blockId: hit.blockId,
elementId: hit.elementId,
frameId: hit.frameId,
updatedAt: hit.updatedAt,
score: 0,
unitId: hit.unitId,
};
}
}
packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts
+5
View File
@@ -1,4 +1,6 @@
/* oxlint-disable import/no-cycle -- Tool callbacks can invoke nested Copilot prompts. */
import { randomUUID } from 'node:crypto';
import { Injectable } from '@nestjs/common';
import { Config } from '../../../base/config';
@@ -198,6 +200,7 @@ export class CapabilityRuntime {
options: RuntimeOptions
) {
const { request, toolSet } = await this.prepareChat(messages, options);
const runId = randomUUID();
const rawStream = this.backend.streamCopilot<
LlmToolLoopStreamEvent | CopilotRuntimeEvent
>(
@@ -218,6 +221,8 @@ export class CapabilityRuntime {
await executeToolCall(toolSet, toolRequest, {
signal: options.signal,
messages,
runId,
toolCallId: toolRequest.callId,
})
);
},
packages/backend/server/src/plugins/copilot/runtime/contracts/shared.ts
+49
View File
@@ -33,6 +33,55 @@ export const JsonObjectSchema = z.record(JsonValueSchema);
export const NonEmptyStringSchema = z.string().trim().min(1);
export const ScopeSelectorSchema = z
.object({
kind: z.enum(['document', 'tag', 'collection', 'favorite', 'artifact']),
id: NonEmptyStringSchema,
name: z.string().optional(),
source: z.enum(['draft', 'focus', 'message']),
})
.strict();
export const ScopeSelectorsSchema = ScopeSelectorSchema.array().max(100);
export const ClientScopeSelectorSchema = ScopeSelectorSchema.omit({
kind: true,
source: true,
}).extend({
kind: z.enum(['document', 'tag', 'collection', 'favorite']),
});
export const RetrievalScopeSchema = z
.object({
mode: z.enum(['workspace', 'required']),
requiredDocIds: z.array(z.string()),
requiredArtifactIds: z.array(z.string()),
preferredSourceIds: z.array(z.string()),
})
.strict();
export const TurnScopeSnapshotSchema = z
.object({
version: z.number().int().positive(),
resolvedAt: z.string(),
selectors: ScopeSelectorsSchema,
requiredDocIds: z.array(z.string()),
requiredArtifactIds: z.array(z.string()),
preferredSourceIds: z.array(z.string()),
retrieval: RetrievalScopeSchema,
})
.strict();
export const SessionFocusSchema = z
.object({
selectors: ScopeSelectorsSchema,
})
.strict();
export type ScopeSelector = z.infer<typeof ScopeSelectorSchema>;
export type TurnScopeSnapshot = z.infer<typeof TurnScopeSnapshotSchema>;
export type SessionFocus = z.infer<typeof SessionFocusSchema>;
export const ToolDefinitionBaseSchema = z
.object({
name: NonEmptyStringSchema,
packages/backend/server/src/plugins/copilot/runtime/copilot-runtime-event-consumer.ts
+15 -7
View File
@@ -53,7 +53,9 @@ export class CopilotRuntimeEventConsumer {
) {
for (const event of events) {
try {
if (event.type === 'usage') {
if (event.type === 'route_selected') {
await this.recordSelection(event, context);
} else if (event.type === 'usage') {
await this.recordUsage(event, context);
} else if (event.type === 'route_failed') {
await this.recordFailure(event, context);
@@ -68,6 +70,18 @@ export class CopilotRuntimeEventConsumer {
}
}
private async recordSelection(
event: Extract<CopilotRuntimeEvent, { type: 'route_selected' }>,
context: CopilotRuntimeEventContext
) {
if (context.workspaceId && event.route.source === 'server') {
await this.models.copilotWorkspaceByokConfig.touchUsed(
context.workspaceId,
event.route.profileId
);
}
}
private async recordUsage(
event: Extract<CopilotRuntimeEvent, { type: 'usage' }>,
context: CopilotRuntimeEventContext
@@ -100,12 +114,6 @@ export class CopilotRuntimeEventConsumer {
totalTokens: usage.total_tokens ?? 0,
cachedTokens: usage.cached_tokens ?? 0,
});
if (event.route.source === 'server') {
await this.models.copilotWorkspaceByokConfig.touchUsed(
context.workspaceId,
event.route.profileId
);
}
}
private async recordFailure(
packages/backend/server/src/plugins/copilot/runtime/hosts/conversation-host.ts
+128 -3
View File
@@ -2,19 +2,30 @@ import { Injectable } from '@nestjs/common';
import {
CopilotMessageNotFound,
CopilotSelectedSourcesLimitExceeded,
CopilotSessionNotFound,
Mutex,
} from '../../../../base';
import { BackendRuntimeProvider } from '../../../../core/backend-runtime';
import { CompatSubmissionStore } from '../../compat/submission-store';
import { ConversationPolicy } from '../../conversation/policy';
import {
canonicalizeTurnTrace,
promptMessageFromTurn,
type Turn,
turnFromChatMessage,
} from '../../core';
import type { PromptParams } from '../../providers/types';
import { ChatSession, ChatSessionService } from '../../session';
import { ChatQuerySchema } from '../../types';
import {
ClientScopeSelectorSchema,
type ScopeSelector,
ScopeSelectorSchema,
type SessionFocus,
TurnScopeSnapshotSchema,
} from '../contracts/shared';
import { AttachmentAdmissionHost } from './attachment-admission';
export type PreparedConversationTurn = {
messageId?: string;
@@ -35,9 +46,114 @@ export class ConversationHost {
private readonly sessions: ChatSessionService,
private readonly submissions: CompatSubmissionStore,
private readonly mutex: Mutex,
private readonly policy: ConversationPolicy
private readonly policy: ConversationPolicy,
private readonly runtime: BackendRuntimeProvider,
private readonly attachmentAdmission: AttachmentAdmissionHost
) {}
private selectors(
value: unknown,
source: ScopeSelector['source']
): ScopeSelector[] {
if (value === undefined) return [];
return ClientScopeSelectorSchema.array()
.max(100)
.parse(value)
.map(selector => ({ ...selector, source }));
}
private mergeSelectors(...groups: ScopeSelector[][]): ScopeSelector[] {
const merged = new Map<string, ScopeSelector>();
for (const selector of groups.flat()) {
merged.set(`${selector.kind}:${selector.id}`, selector);
}
return [...merged.values()];
}
private async prepareMessageState(
session: ChatSession,
params: Record<string, any>,
attachments: NonNullable<
Parameters<AttachmentAdmissionHost['admitPromptAttachments']>[0]
>
) {
const {
scopeSelectors: rawSelectors,
focusSelectors: rawFocus,
preferredSourceIds: rawPreferred,
...metadata
} = params;
const focus: SessionFocus =
rawFocus === undefined
? session.config.focus
: { selectors: this.selectors(rawFocus, 'focus') };
const admitted = await this.attachmentAdmission.admitPromptAttachments(
attachments,
{
userId: session.config.userId,
workspaceId: session.config.workspaceId,
sessionId: session.config.sessionId,
}
);
const artifacts = await Promise.all(
admitted.map(async source => {
const artifact = await this.runtime.putWorkspaceArtifact(
{
workspaceId: session.config.workspaceId,
mimeType: source.mimeType,
fileName: source.fileName,
libraryOwned: false,
},
Buffer.from(source.data, 'base64')
);
return {
artifactId: artifact.id,
role: 'attachment',
displayName: source.fileName,
metadata: { mimeType: artifact.canonicalMediaType },
};
})
);
const artifactSelectors = artifacts.map(
({ artifactId, displayName }): ScopeSelector => ({
kind: 'artifact',
id: artifactId,
name: displayName,
source: 'message',
})
);
const selectors = this.mergeSelectors(
focus.selectors,
this.selectors(rawSelectors, 'draft'),
artifactSelectors
);
const preferredSourceIds =
rawPreferred === undefined
? []
: ScopeSelectorSchema.shape.id.array().max(100).parse(rawPreferred);
let compiledScope: Awaited<
ReturnType<BackendRuntimeProvider['compileTurnScope']>
>;
try {
compiledScope = await this.runtime.compileTurnScope({
workspaceId: session.config.workspaceId,
userId: session.config.userId,
selectors,
preferredSourceIds,
});
} catch (error) {
if (
error instanceof Error &&
error.message.includes('scope_required_document_limit_exceeded')
) {
throw new CopilotSelectedSourcesLimitExceeded();
}
throw error;
}
const scopeSnapshot = TurnScopeSnapshotSchema.parse(compiledScope);
return { artifacts, focus, metadata, scopeSnapshot };
}
private async loadAcceptedTurn(
session: ChatSession,
sessionId: string,
@@ -180,16 +296,25 @@ export class ConversationHost {
session.revertLatestMessage(true);
}
const prepared = await this.prepareMessageState(
session,
submission.params ?? {},
submission.attachments ?? []
);
const turn = await this.sessions.appendTurn({
sessionId,
userId: session.config.userId,
compatSubmissionId: messageId,
focus: prepared.focus,
artifacts: prepared.artifacts,
turn: {
conversationId: sessionId,
role: 'user',
content: submission.content ?? '',
attachments: submission.attachments ?? [],
metadata: submission.params ?? {},
metadata: prepared.metadata,
scopeSnapshot: prepared.scopeSnapshot,
renderTrace: [],
toolEvents: [],
createdAt: submission.createdAt,
@@ -245,7 +370,7 @@ export class ConversationHost {
return {
...latestTurn.metadata,
content: latestTurn.content,
attachments: latestTurn.attachments,
attachments: promptMessageFromTurn(latestTurn).attachments ?? [],
};
}
packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts
+94 -45
View File
@@ -1,38 +1,43 @@
/* oxlint-disable import/no-cycle -- Tools can invoke nested prompts and semantic search. */
import { Injectable } from '@nestjs/common';
import { forwardRef, Inject, Injectable } from '@nestjs/common';
import { Config } from '../../../base';
import { DocReader, DocWriter } from '../../../core/doc';
import { PermissionAccess } from '../../../core/permission';
import { Models } from '../../../models';
import { IndexerService } from '../../indexer';
import { CopilotContextService } from '../context/service';
import { DelegatedEditorService } from '../delegated/service';
import {
type CopilotChatOptions,
type CopilotChatTools,
} from '../providers/types';
import { ArtifactRetrievalService } from '../retrieval/artifact';
import { DocumentRetrievalService } from '../retrieval/document';
import {
buildBlobContentGetter,
buildDocCanvasGetter,
buildDocContentGetter,
buildDocCreateHandler,
buildDocKeywordSearchGetter,
buildDocSearchGetter,
buildDocumentSearch,
buildDocUpdateHandler,
buildDocUpdateMetaHandler,
type CopilotTool,
type CopilotToolSet,
createBlobReadTool,
createArtifactReadTool,
createArtifactSearchTool,
createCodeArtifactTool,
createConversationSummaryTool,
createDocCanvasReadTool,
createDocComposeTool,
createDocCreateTool,
createDocKeywordSearchTool,
createDocReadTool,
createDocSemanticSearchTool,
createDocSearchTool,
createDocUpdateMetaTool,
createDocUpdateTool,
createExaCrawlTool,
createExaSearchTool,
createFrontendEditorStateTool,
createFrontendNodesTool,
createFrontendSelectionTool,
createFrontendSnapshotTool,
createSectionEditTool,
} from '../tools';
import { PromptRuntime } from './prompt-runtime';
@@ -47,12 +52,14 @@ export class ToolRuntime {
constructor(
private readonly config: Config,
private readonly ac: PermissionAccess,
private readonly context: CopilotContextService,
private readonly docReader: DocReader,
private readonly docWriter: DocWriter,
private readonly models: Models,
private readonly promptRuntime: PromptRuntime,
private readonly indexerService: IndexerService
@Inject(forwardRef(() => PromptRuntime))
private readonly promptRuntime: Pick<PromptRuntime, 'runText'>,
private readonly retrieval: DocumentRetrievalService,
private readonly artifactRetrieval: ArtifactRetrievalService,
private readonly delegated: DelegatedEditorService
) {}
async getTools(
@@ -80,6 +87,14 @@ export class ToolRuntime {
},
});
const documentScope =
options.retrievalScope?.mode === 'required'
? {
mode: 'selected' as const,
allowedDocIds: options.retrievalScope.requiredDocIds,
}
: undefined;
for (const tool of options.tools) {
const toolDef = resolveProviderSpecificTool?.(tool, model);
if (toolDef) {
@@ -97,13 +112,17 @@ export class ToolRuntime {
}
switch (tool) {
case 'blobRead': {
const docContext = options.session
? await this.context.getBySessionId(options.session)
: null;
const getBlobContent = buildBlobContentGetter(this.ac, docContext);
tools.blob_read = createBlobReadTool(
getBlobContent.bind(null, options)
case 'artifactRead': {
tools.artifact_read = createArtifactReadTool(
this.artifactRetrieval,
options
);
break;
}
case 'artifactSearch': {
tools.artifact_search = createArtifactSearchTool(
this.artifactRetrieval,
options
);
break;
}
@@ -118,40 +137,70 @@ export class ToolRuntime {
);
break;
}
case 'docSemanticSearch': {
const searchDocs = buildDocSearchGetter(
this.ac,
this.context,
options.session,
this.models
);
tools.doc_semantic_search = createDocSemanticSearchTool(
searchDocs.bind(null, options)
);
break;
}
case 'docKeywordSearch': {
if (this.config.indexer.enabled) {
const searchDocs = buildDocKeywordSearchGetter(
this.ac,
this.indexerService,
this.models
);
tools.doc_keyword_search = createDocKeywordSearchTool(
searchDocs.bind(null, options)
);
}
break;
}
case 'docRead': {
const getDoc = buildDocContentGetter(
this.ac,
this.docReader,
this.models
this.models,
documentScope
);
tools.doc_read = createDocReadTool(getDoc.bind(null, options));
break;
}
case 'docCanvasRead': {
const readCanvas = buildDocCanvasGetter(
this.ac,
this.docReader,
this.models,
documentScope
);
tools.doc_canvas_read = createDocCanvasReadTool(
readCanvas.bind(null, options)
);
break;
}
case 'docSearch': {
tools.doc_search = createDocSearchTool(
buildDocumentSearch(this.retrieval, options, documentScope)
);
break;
}
case 'frontendGetEditorState': {
if (this.delegated.getLease(options, 'frontend_get_editor_state')) {
tools.frontend_get_editor_state = createFrontendEditorStateTool(
this.delegated,
options
);
}
break;
}
case 'frontendReadSelection': {
if (this.delegated.getLease(options, 'frontend_read_selection')) {
tools.frontend_read_selection = createFrontendSelectionTool(
this.delegated,
options
);
}
break;
}
case 'frontendReadNodes': {
if (this.delegated.getLease(options, 'frontend_read_nodes')) {
tools.frontend_read_nodes = createFrontendNodesTool(
this.delegated,
options
);
}
break;
}
case 'frontendSnapshotDocument': {
if (this.delegated.getLease(options, 'frontend_snapshot_document')) {
tools.frontend_snapshot_document = createFrontendSnapshotTool(
this.delegated,
options
);
}
break;
}
case 'docCreate': {
const createDoc = buildDocCreateHandler(this.ac, this.docWriter);
tools.doc_create = createDocCreateTool(createDoc.bind(null, options));
packages/backend/server/src/plugins/copilot/runtime/tool/footnotes.ts
+146
View File
@@ -0,0 +1,146 @@
import type { DocSource } from '../../tools/types';
import type { EnrichedToolResultEvent } from './native-runtime-adapter';
export type AttachmentFootnote = {
artifactId: string;
fileName: string;
fileType: string;
};
function pickAttachmentFootnote(value: unknown): AttachmentFootnote | null {
if (!value || typeof value !== 'object') return null;
const record = value as Record<string, unknown>;
if (record.source && typeof record.source === 'object') {
const source = pickAttachmentFootnote(record.source);
if (source) return source;
}
const artifactId =
typeof record.artifactId === 'string'
? record.artifactId
: typeof record.artifact_id === 'string'
? record.artifact_id
: undefined;
const fileName =
typeof record.fileName === 'string'
? record.fileName
: typeof record.name === 'string'
? record.name
: 'Attachment';
const fileType =
typeof record.fileType === 'string'
? record.fileType
: typeof record.mimeType === 'string'
? record.mimeType
: typeof record.mime_type === 'string'
? record.mime_type
: 'application/octet-stream';
return artifactId ? { artifactId, fileName, fileType } : null;
}
export function collectAttachmentFootnotes(
event: EnrichedToolResultEvent
): AttachmentFootnote[] {
if (!['artifact_read', 'artifact_search'].includes(event.name)) return [];
if (!event.output || typeof event.output !== 'object') return [];
const output = event.output as Record<string, unknown>;
if (event.name === 'artifact_search' && Array.isArray(output.hits)) {
return output.hits
.map(pickAttachmentFootnote)
.filter((item): item is AttachmentFootnote => item !== null);
}
const item = pickAttachmentFootnote(output);
return item ? [item] : [];
}
export function formatAttachmentFootnotes(
attachments: AttachmentFootnote[],
options: { includeReferences?: boolean } = {}
) {
const references =
options.includeReferences === false
? ''
: attachments.map((_, index) => `[^attachment-${index + 1}]`).join('');
const definitions = attachments
.map(
(attachment, index) =>
`[^attachment-${index + 1}]: ${JSON.stringify({
type: 'attachment',
artifactId: attachment.artifactId,
fileName: attachment.fileName,
fileType: attachment.fileType,
})}`
)
.join('\n');
return references
? `\n\n${references}\n\n${definitions}`
: `\n\n${definitions}`;
}
function pickDocumentFootnote(value: unknown): DocSource | null {
if (!value || typeof value !== 'object') return null;
const source = value as Record<string, unknown>;
if (source.type !== 'document') return null;
const workspaceId = source.workspace_id ?? source.workspaceId;
const docId = source.doc_id ?? source.docId;
if (typeof workspaceId !== 'string' || typeof docId !== 'string') return null;
const optional = (snake: string, camel: string) => {
const candidate = source[snake] ?? source[camel];
return typeof candidate === 'string' ? candidate : undefined;
};
return {
type: 'document',
workspace_id: workspaceId,
doc_id: docId,
title: typeof source.title === 'string' ? source.title : '',
revision: optional('revision', 'revision'),
visibility: optional('visibility', 'visibility') as
| DocSource['visibility']
| undefined,
block_id: optional('block_id', 'blockId'),
element_id: optional('element_id', 'elementId'),
frame_id: optional('frame_id', 'frameId'),
};
}
export function collectDocumentFootnotes(event: EnrichedToolResultEvent) {
if (
![
'doc_read',
'doc_canvas_read',
'doc_search',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
].includes(event.name)
)
return [];
if (!event.output || typeof event.output !== 'object') return [];
const output = event.output as Record<string, unknown>;
const direct = pickDocumentFootnote(output.source);
if (direct) return [direct];
return Array.isArray(output.hits)
? output.hits
.map(hit =>
pickDocumentFootnote((hit as Record<string, unknown>)?.source)
)
.filter((source): source is DocSource => source !== null)
: [];
}
export function formatDocumentFootnotes(documents: DocSource[]) {
const unique = [
...new Map(documents.map(document => [document.doc_id, document])).values(),
];
const references = unique.map((_, index) => `[^doc-${index + 1}]`).join('');
const definitions = unique
.map(
(document, index) =>
`[^doc-${index + 1}]: ${JSON.stringify({
type: 'doc',
docId: document.doc_id,
...(document.title ? { title: document.title } : {}),
})}`
)
.join('\n');
return `\n\n${references}\n\n${definitions}`;
}
packages/backend/server/src/plugins/copilot/runtime/tool/native-adapter.ts
+50 -88
View File
@@ -7,19 +7,21 @@ import {
CitationFootnoteFormatter,
TextStreamParser,
} from '../../providers/utils';
import type { DocSource } from '../../tools/types';
import { projectRuntimeEventToStreamObject } from '../contracts/runtime-event-contract';
import {
type AttachmentFootnote,
collectAttachmentFootnotes,
collectDocumentFootnotes,
formatAttachmentFootnotes,
formatDocumentFootnotes,
} from './footnotes';
import {
type EnrichedToolCallEvent,
type EnrichedToolResultEvent,
NativeRuntimeAdapter,
} from './native-runtime-adapter';
type AttachmentFootnote = {
blobId: string;
fileName: string;
fileType: string;
};
export type NativeProviderAdapterOptions = {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
@@ -34,79 +36,6 @@ type NativeStreamDispatch = ConstructorParameters<
typeof NativeRuntimeAdapter
>[0];
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: EnrichedToolResultEvent
): 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[],
options: { includeReferences?: boolean } = {}
) {
const references =
options.includeReferences === false
? ''
: 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 references
? `\n\n${references}\n\n${definitions}`
: `\n\n${definitions}`;
}
export class NativeProviderAdapter {
readonly logger = new Logger(NativeProviderAdapter.name);
readonly #runtime: NativeRuntimeAdapter;
@@ -180,6 +109,9 @@ export class NativeProviderAdapter {
const citationFormatter = this.#enableCitationFootnote
? new CitationFootnoteFormatter()
: null;
const attachmentFootnotes = new Map<string, AttachmentFootnote>();
const documentFootnotes = new Map<string, DocSource>();
let hasAttachmentFootnoteReference = false;
let streamPartId = 0;
const usageState: {
model?: string;
@@ -210,6 +142,9 @@ export class NativeProviderAdapter {
}
case 'text_delta': {
const textEvent = event as unknown as { text: string };
if (textEvent.text.includes('[^attachment-')) {
hasAttachmentFootnoteReference = true;
}
if (textParser) {
yield textParser.parse({
type: 'text-delta',
@@ -247,8 +182,14 @@ export class NativeProviderAdapter {
break;
}
case 'tool_result': {
if (!textParser) break;
const normalized = event as EnrichedToolResultEvent;
collectAttachmentFootnotes(normalized).forEach(attachment => {
attachmentFootnotes.set(attachment.artifactId, attachment);
});
collectDocumentFootnotes(normalized).forEach(document => {
documentFootnotes.set(JSON.stringify(document), document);
});
if (!textParser) break;
yield textParser.parse({
type: 'tool-result',
toolCallId: normalized.call_id,
@@ -280,7 +221,17 @@ export class NativeProviderAdapter {
usageState.usage = doneEvent.usage ?? usageState.usage;
const footnotes = textParser?.end() ?? '';
const citations = citationFormatter?.end() ?? '';
const tails = [citations, footnotes].filter(Boolean).join('\n');
const attachments = attachmentFootnotes.size
? formatAttachmentFootnotes([...attachmentFootnotes.values()], {
includeReferences: !hasAttachmentFootnoteReference,
})
: '';
const documents = documentFootnotes.size
? formatDocumentFootnotes([...documentFootnotes.values()])
: '';
const tails = [citations, attachments, documents, footnotes]
.filter(Boolean)
.join('\n');
if (tails) {
yield `\n${tails}`;
}
@@ -310,7 +261,8 @@ export class NativeProviderAdapter {
? new CitationFootnoteFormatter()
: null;
const fallbackAttachmentFootnotes = new Map<string, AttachmentFootnote>();
let hasFootnoteReference = false;
const fallbackDocumentFootnotes = new Map<string, DocSource>();
let hasAttachmentFootnoteReference = false;
const usageState: {
model?: string;
usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
@@ -340,8 +292,8 @@ export class NativeProviderAdapter {
}
case 'text_delta': {
const textEvent = event as unknown as { text: string };
if (textEvent.text.includes('[^')) {
hasFootnoteReference = true;
if (textEvent.text.includes('[^attachment-')) {
hasAttachmentFootnoteReference = true;
}
yield { type: 'text-delta', textDelta: textEvent.text };
break;
@@ -363,7 +315,10 @@ export class NativeProviderAdapter {
const normalized = event as EnrichedToolResultEvent;
const attachments = collectAttachmentFootnotes(normalized);
attachments.forEach(attachment => {
fallbackAttachmentFootnotes.set(attachment.blobId, attachment);
fallbackAttachmentFootnotes.set(attachment.artifactId, attachment);
});
collectDocumentFootnotes(normalized).forEach(document => {
fallbackDocumentFootnotes.set(JSON.stringify(document), document);
});
const streamObject = projectRuntimeEventToStreamObject(
event as LlmToolLoopStreamEvent
@@ -394,18 +349,25 @@ export class NativeProviderAdapter {
usageState.usage = doneEvent.usage ?? usageState.usage;
const citations = citationFormatter?.end() ?? '';
if (citations) {
hasFootnoteReference = true;
yield { type: 'text-delta', textDelta: `\n${citations}` };
}
if (!citations && fallbackAttachmentFootnotes.size > 0) {
if (fallbackAttachmentFootnotes.size > 0) {
yield {
type: 'text-delta',
textDelta: formatAttachmentFootnotes(
Array.from(fallbackAttachmentFootnotes.values()),
{ includeReferences: !hasFootnoteReference }
{ includeReferences: !hasAttachmentFootnoteReference }
),
};
}
if (fallbackDocumentFootnotes.size > 0) {
yield {
type: 'text-delta',
textDelta: formatDocumentFootnotes([
...fallbackDocumentFootnotes.values(),
]),
};
}
break;
}
case 'provider_selected':
packages/backend/server/src/plugins/copilot/runtime/turn-orchestrator.ts
+14 -31
View File
@@ -1,6 +1,6 @@
import { Injectable } from '@nestjs/common';
import { CopilotContextService } from '../context/service';
import { BackendRuntimeEmbeddingJob } from '../../../core/backend-runtime';
import { type Turn } from '../core';
import {
type ModelConditions,
@@ -20,32 +20,14 @@ import { TurnPersistence } from './hosts/turn-persistence';
export class TurnOrchestrator {
constructor(
private readonly conversations: ConversationHost,
private readonly context: CopilotContextService,
private readonly runtime: CapabilityRuntime,
private readonly imageResults: ImageResultHost,
private readonly turnPersistence: TurnPersistence
private readonly turnPersistence: TurnPersistence,
private readonly embeddings: BackendRuntimeEmbeddingJob
) {}
private async buildPromptParams(
sessionId: string,
options: {
latestTurn?: Turn;
includeContextFiles?: boolean;
} = {}
): Promise<Record<string, unknown>> {
const current = await this.context.getBySessionId(sessionId);
const contextFiles =
options.includeContextFiles &&
current &&
(current.files.length > 0 || current.blobs.length > 0)
? [...current.files, ...(await current.getBlobMetadata())]
: [];
const latestTurn = options.latestTurn;
return {
...this.conversations.buildLatestTurnPromptParams(latestTurn),
...(contextFiles.length ? { contextFiles } : {}),
};
private buildPromptParams(latestTurn?: Turn): Record<string, unknown> {
return this.conversations.buildLatestTurnPromptParams(latestTurn);
}
private async prepareChatSelection(
@@ -54,7 +36,6 @@ export class TurnOrchestrator {
query: Record<string, string | string[]>,
selection: {
responseMode: 'text' | 'object' | 'image';
includeContextFiles?: boolean;
}
) {
const prepared = await this.conversations.prepareTurn(
@@ -71,15 +52,18 @@ export class TurnOrchestrator {
toolsConfig,
byokLeaseId,
} = ChatQuerySchema.parse(query);
const promptParams = await this.buildPromptParams(sessionId, {
latestTurn: prepared.latestTurn,
includeContextFiles: selection.includeContextFiles,
});
const promptParams = this.buildPromptParams(prepared.latestTurn);
const scope = prepared.latestTurn?.scopeSnapshot?.retrieval;
if (scope?.mode === 'required' && scope.requiredDocIds.length) {
await this.embeddings.prepareSelectedDocuments(
prepared.session.config.workspaceId,
scope.requiredDocIds
);
}
const finalMessage = prepared.session.finish({
...prepared.params,
...promptParams,
});
return {
prepared,
finalMessage,
@@ -97,6 +81,7 @@ export class TurnOrchestrator {
builtInRouteId: prepared.session.config.promptName,
managedTargetId: routeTargetId,
quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed,
retrievalScope: prepared.latestTurn?.scopeSnapshot?.retrieval,
featureKind:
selection.responseMode === 'image'
? 'image'
@@ -124,7 +109,6 @@ export class TurnOrchestrator {
const { prepared, finalMessage, selection } =
await this.prepareChatSelection(userId, sessionId, query, {
responseMode: 'text',
includeContextFiles: true,
});
const stream = this.streamTextResult(
@@ -175,7 +159,6 @@ export class TurnOrchestrator {
const { prepared, finalMessage, selection } =
await this.prepareChatSelection(userId, sessionId, query, {
responseMode: 'object',
includeContextFiles: true,
});
return {
packages/backend/server/src/plugins/copilot/session.ts
+22 -1
View File
@@ -70,10 +70,19 @@ export class ChatSession implements AsyncDisposable {
userId,
workspaceId,
docId,
focus,
prompt: { name: promptName, config: promptConfig },
} = this.state;
return { sessionId, userId, workspaceId, docId, promptName, promptConfig };
return {
sessionId,
userId,
workspaceId,
docId,
focus,
promptName,
promptConfig,
};
}
get stashTurns() {
@@ -144,11 +153,13 @@ export class ChatSession implements AsyncDisposable {
export type ConversationState = {
conversation: Conversation;
turns: Turn[];
focus: ChatSessionState['focus'];
prompt: ResolvedPrompt;
};
export type ConversationMetaState = {
conversation: Conversation;
focus: ChatSessionState['focus'];
prompt: ResolvedPrompt;
};
@@ -196,6 +207,7 @@ export class ChatSessionService {
return {
conversation,
turns: session.turns,
focus: session.focus,
prompt,
};
}
@@ -208,6 +220,7 @@ export class ChatSessionService {
return {
conversation: session.conversation,
focus: session.focus,
prompt,
};
}
@@ -417,6 +430,13 @@ export class ChatSessionService {
userId: string;
turn: Turn;
compatSubmissionId?: string;
focus?: ChatSessionState['focus'];
artifacts?: Array<{
artifactId: string;
role: string;
displayName?: string;
metadata?: Record<string, unknown>;
}>;
}) {
return await this.store.appendTurn(input);
}
@@ -463,6 +483,7 @@ export class ChatSessionService {
workspaceId: state.conversation.workspaceId,
docId: state.conversation.docId,
turns: state.turns,
focus: state.focus,
prompt: state.prompt,
},
(prompt, turns, params, sessionId) =>
packages/backend/server/src/plugins/copilot/tools/artifact.ts
+111
View File
@@ -0,0 +1,111 @@
import { Logger } from '@nestjs/common';
import { z } from 'zod';
import type { ArtifactRetrievalService } from '../retrieval/artifact';
import { toolError } from './error';
import { defineTool } from './tool';
import type { ArtifactSource, CopilotChatOptions } from './types';
const logger = new Logger('ArtifactTool');
export const createArtifactSearchTool = (
retrieval: ArtifactRetrievalService,
options: CopilotChatOptions
) =>
defineTool({
description:
'Search workspace artifacts and message attachments within the current retrieval scope. This tool never searches documents or the web.',
inputSchema: z
.object({
query: z.string().trim().min(1).max(2000),
limit: z.number().int().min(1).max(10).optional(),
})
.strict(),
execute: async ({ query, limit }, execution) => {
if (!options?.user || !options.workspace || !options.retrievalScope) {
return toolError('Artifact Search Failed', 'Missing retrieval scope.', {
code: 'INVALID_CONTEXT',
retryable: false,
});
}
const result = await retrieval.search({
userId: options.user,
workspaceId: options.workspace,
query,
retrieval: options.retrievalScope,
limit: limit ?? 5,
messageId: options.billingUnitId,
signal: execution.signal,
});
return {
degraded: result.degraded,
hits: result.hits.map(hit => ({
artifact_id: hit.artifactId,
excerpt: hit.content,
distance: hit.distance,
chunk: hit.chunk,
source: {
type: 'artifact',
workspace_id: options.workspace as string,
artifact_id: hit.artifactId as string,
name: hit.name,
mime_type: hit.mimeType,
} satisfies ArtifactSource,
})),
};
},
});
export const createArtifactReadTool = (
retrieval: ArtifactRetrievalService,
options: CopilotChatOptions
) =>
defineTool({
description:
'Read extracted content from an artifact within the current retrieval scope. Use cursor to continue a truncated result.',
inputSchema: z
.object({
artifact_id: z.string().uuid(),
max_chars: z.number().int().min(1).max(100_000).optional(),
cursor: z.string().max(128).optional(),
})
.strict(),
execute: async ({ artifact_id, max_chars, cursor }) => {
if (!options?.user || !options.workspace || !options.retrievalScope) {
return toolError('Artifact Read Failed', 'Missing retrieval scope.', {
code: 'INVALID_CONTEXT',
retryable: false,
});
}
try {
const result = await retrieval.read({
userId: options.user,
workspaceId: options.workspace,
artifactId: artifact_id,
retrieval: options.retrievalScope,
messageId: options.billingUnitId,
maxChars: max_chars,
cursor,
});
return {
artifact_id,
...result,
source: {
type: 'artifact',
workspace_id: options.workspace,
artifact_id,
name: result.name,
mime_type: result.mimeType,
revision: result.revision,
} satisfies ArtifactSource,
};
} catch (error) {
logger.warn('Artifact read denied or unavailable', error);
return toolError(
'Artifact Read Failed',
'The artifact is unavailable in the current scope.',
{ code: 'ARTIFACT_UNAVAILABLE', retryable: false }
);
}
},
});
packages/backend/server/src/plugins/copilot/tools/blob-read.ts
-91
View File
@@ -1,91 +0,0 @@
import { Logger } from '@nestjs/common';
import { z } from 'zod';
import { PermissionAccess } from '../../../core/permission';
import { toolError } from './error';
import { defineTool } from './tool';
import type { ContextSession, CopilotChatOptions } from './types';
const logger = new Logger('ContextBlobReadTool');
export const buildBlobContentGetter = (
ac: PermissionAccess,
context: ContextSession | null
) => {
const getBlobContent = async (
options: CopilotChatOptions,
blobId?: string,
chunk?: number
) => {
if (!options?.user || !options?.workspace || !blobId || !context) {
return toolError(
'Blob Read Failed',
'Missing workspace, user, blob id, or copilot context for blob_read.'
);
}
const canAccess = await ac
.user(options.user)
.workspace(options.workspace)
.allowLocal()
.can('Workspace.Read');
if (!canAccess || context.workspaceId !== options.workspace) {
logger.warn(
`User ${options.user} does not have access workspace ${options.workspace}`
);
return toolError(
'Blob Read Failed',
'You do not have permission to access this workspace attachment.'
);
}
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([
targetFileId ? context.getFileContent(targetFileId, chunk) : undefined,
context.getBlobContent(canonicalBlobId, chunk),
]);
const content = file?.trim() || blob?.trim();
if (!content) {
return toolError(
'Blob Read Failed',
`Attachment ${canonicalBlobId} is not available for reading in the current copilot context.`
);
}
const info = contextFile
? { fileName: contextFile.name, fileType: contextFile.mimeType }
: {};
return { blobId: canonicalBlobId, chunk, content, ...info };
};
return getBlobContent;
};
export const createBlobReadTool = (
getBlobContent: (targetId?: string, chunk?: number) => Promise<object>
) => {
return defineTool({
description:
'Return the content and basic metadata of a single attachment identified by blobId; more inclined to use search tools rather than this tool.',
inputSchema: z.object({
blob_id: z.string().describe('The target blob in context to read'),
chunk: z
.number()
.optional()
.describe(
'The chunk number to read, if not provided, read the whole content, start from 0'
),
}),
execute: async ({ blob_id, chunk }) => {
try {
const blob = await getBlobContent(blob_id, chunk);
return { ...blob };
} catch (err: any) {
logger.error(`Failed to read the blob ${blob_id} in context`, err);
return toolError('Blob Read Failed', err.message ?? String(err));
}
},
});
};
packages/backend/server/src/plugins/copilot/tools/doc-canvas-read.ts
+354
View File
@@ -0,0 +1,354 @@
import { createHash } from 'node:crypto';
import { Logger } from '@nestjs/common';
import { z } from 'zod';
import type { DocReader } from '../../../core/doc';
import type { PermissionAccess } from '../../../core/permission';
import type {
CanvasProjectionBlock,
CanvasProjectionElement,
CanvasProjectionV1,
DocBounds,
} from '../../../core/utils/blocksuite';
import type { Models } from '../../../models';
import {
documentSyncPendingError,
workspaceSyncRequiredError,
} from './doc-sync';
import { toolError } from './error';
import { defineTool } from './tool';
import {
type CopilotChatOptions,
type DocSource,
type DocumentScope,
isDocumentInScope,
} from './types';
const logger = new Logger('DocCanvasReadTool');
const MAX_LIMIT = 100;
const MAX_PREVIEW_CHARS = 4_000;
const MAX_RELATION_IDS = 200;
const boundsSchema = z
.object({
x: z.number().finite().min(-10_000_000).max(10_000_000),
y: z.number().finite().min(-10_000_000).max(10_000_000),
width: z.number().finite().positive().max(10_000_000),
height: z.number().finite().positive().max(10_000_000),
})
.strict();
const targetSchema = z.discriminatedUnion('kind', [
z.object({ kind: z.literal('overview') }).strict(),
z
.object({ kind: z.literal('frame'), frame_id: z.string().min(1).max(128) })
.strict(),
z
.object({
kind: z.literal('elements'),
element_ids: z.array(z.string().min(1).max(128)).min(1).max(MAX_LIMIT),
})
.strict(),
z.object({ kind: z.literal('region'), bounds: boundsSchema }).strict(),
]);
type CanvasTarget = z.infer<typeof targetSchema>;
const cursorSchema = z
.object({
version: z.literal(1),
projectionVersion: z.number().int().positive(),
revision: z.string(),
targetHash: z.string(),
offset: z.number().int().nonnegative(),
})
.strict();
type Cursor = z.infer<typeof cursorSchema>;
function targetHash(target: CanvasTarget) {
return createHash('sha256').update(JSON.stringify(target)).digest('hex');
}
function encodeCursor(cursor: Cursor) {
return Buffer.from(JSON.stringify(cursor)).toString('base64url');
}
function decodeCursor(value: string): Cursor | null {
try {
const parsed = cursorSchema.safeParse(
JSON.parse(Buffer.from(value, 'base64url').toString())
);
return parsed.success ? parsed.data : null;
} catch {
return null;
}
}
function boundedContent(value: { text?: string; title?: string }) {
const text = value.text?.slice(0, MAX_PREVIEW_CHARS);
const title = value.title?.slice(0, MAX_PREVIEW_CHARS);
const contentTruncated =
(value.text?.length ?? 0) > (text?.length ?? 0) ||
(value.title?.length ?? 0) > (title?.length ?? 0);
return {
text,
title,
...(contentTruncated ? { content_truncated: true } : {}),
};
}
function boundedBlock(value: CanvasProjectionBlock) {
return {
id: value.id,
type: value.type,
visibility: value.visibility,
bounds: value.bounds,
...boundedContent(value),
child_ids: value.childIds.slice(0, MAX_RELATION_IDS),
child_ids_truncated: value.childIds.length > MAX_RELATION_IDS,
};
}
function boundedElement(value: CanvasProjectionElement) {
return {
id: value.id,
type: value.type,
bounds: value.bounds,
...boundedContent(value),
frame_id: value.frameId,
child_ids: value.childIds.slice(0, MAX_RELATION_IDS),
child_ids_truncated: value.childIds.length > MAX_RELATION_IDS,
source_id: value.sourceId,
target_id: value.targetId,
parent_id: value.parentId,
index: value.index,
point_count: value.pointCount,
color: value.color,
line_width: value.lineWidth,
};
}
function intersects(left: DocBounds | undefined, right: DocBounds) {
return Boolean(
left &&
left.x < right.x + right.width &&
left.x + left.width > right.x &&
left.y < right.y + right.height &&
left.y + left.height > right.y
);
}
function selectProjection(
projection: CanvasProjectionV1,
target: CanvasTarget
) {
const canvasBlocks = projection.blocks.filter(
block => block.visibility !== 'page'
);
switch (target.kind) {
case 'overview': {
const ownedIds = new Set(
canvasBlocks
.filter(block => block.type === 'frame')
.flatMap(block => block.childIds)
);
return {
blocks: canvasBlocks.filter(
block => block.type === 'frame' || !ownedIds.has(block.id)
),
elements: projection.elements.filter(element => !element.frameId),
};
}
case 'frame': {
const frame = canvasBlocks.find(block => block.id === target.frame_id);
const childIds = new Set(frame?.childIds ?? []);
return {
blocks: canvasBlocks.filter(
block => block.id === target.frame_id || childIds.has(block.id)
),
elements: projection.elements.filter(
element =>
element.frameId === target.frame_id || childIds.has(element.id)
),
};
}
case 'elements': {
const ids = new Set(target.element_ids);
return {
blocks: canvasBlocks.filter(block => ids.has(block.id)),
elements: projection.elements.filter(element => ids.has(element.id)),
};
}
case 'region':
return {
blocks: canvasBlocks.filter(block =>
intersects(block.bounds, target.bounds)
),
elements: projection.elements.filter(element =>
intersects(element.bounds, target.bounds)
),
};
}
}
export const buildDocCanvasGetter = (
ac: PermissionAccess,
docReader: DocReader,
models: Models,
documentScope?: DocumentScope
) => {
return async (
options: CopilotChatOptions,
docId: string,
target: CanvasTarget,
cursorValue: string | undefined,
requestedLimit: number | undefined
) => {
if (!options?.user || !options.workspace) {
return toolError('Doc Canvas Read Failed', 'Missing workspace or user.', {
code: 'INVALID_CONTEXT',
retryable: false,
});
}
if (!isDocumentInScope(documentScope, docId)) {
return toolError(
'Doc Canvas Read Failed',
'The document is outside the user-selected document scope.',
{
code: 'DOC_SCOPE_DENIED',
retryable: false,
locator: { doc_id: docId },
}
);
}
if (!(await models.workspace.get(options.workspace))) {
return workspaceSyncRequiredError();
}
const canAccess = await ac
.user(options.user)
.workspace(options.workspace)
.doc(docId)
.can('Doc.Read');
if (!canAccess) {
return toolError('Doc Canvas Read Failed', 'Document access denied.', {
code: 'DOC_ACCESS_DENIED',
retryable: false,
locator: { doc_id: docId },
});
}
const projection = await docReader.getDocCanvas(options.workspace, docId);
if (!projection) {
return documentSyncPendingError(docId);
}
const cursor = cursorValue ? decodeCursor(cursorValue) : null;
if (cursorValue && !cursor) {
return toolError('Doc Canvas Read Failed', 'Invalid canvas cursor.', {
code: 'INVALID_CURSOR',
retryable: false,
locator: { doc_id: docId },
});
}
const fingerprint = targetHash(target);
if (
cursor &&
(cursor.projectionVersion !== projection.version ||
cursor.revision !== projection.revision ||
cursor.targetHash !== fingerprint)
) {
return toolError(
'Doc Canvas Read Failed',
'The document changed after this cursor was issued.',
{
code: 'REVISION_CHANGED',
retryable: true,
locator: { doc_id: docId, revision: projection.revision },
}
);
}
const selected = selectProjection(projection, target);
const items = [
...selected.blocks.map(value => ({ kind: 'block' as const, value })),
...selected.elements.map(value => ({ kind: 'element' as const, value })),
].sort((left, right) => left.value.id.localeCompare(right.value.id));
const offset = cursor?.offset ?? 0;
const limit = Math.min(requestedLimit ?? 50, MAX_LIMIT);
const page = items.slice(offset, offset + limit);
const nextOffset = offset + page.length;
const truncated = nextOffset < items.length;
return {
doc_id: projection.docId,
revision: projection.revision,
target,
bounds: projection.bounds,
counts: projection.counts,
blocks: page
.filter(item => item.kind === 'block')
.map(item => boundedBlock(item.value)),
elements: page
.filter(item => item.kind === 'element')
.map(item => boundedElement(item.value)),
truncated,
next_cursor: truncated
? encodeCursor({
version: 1,
projectionVersion: projection.version,
revision: projection.revision,
targetHash: fingerprint,
offset: nextOffset,
})
: undefined,
warnings: projection.warnings.slice(0, MAX_LIMIT),
warnings_truncated: projection.warnings.length > MAX_LIMIT,
source: {
type: 'document' as const,
workspace_id: options.workspace,
doc_id: projection.docId,
title: projection.title,
revision: projection.revision,
visibility: 'edgeless' as const,
} satisfies DocSource,
};
};
};
type CanvasReadResult = Awaited<
ReturnType<ReturnType<typeof buildDocCanvasGetter>>
>;
export const createDocCanvasReadTool = (
readCanvas: (
docId: string,
target: CanvasTarget,
cursor?: string,
limit?: number
) => Promise<CanvasReadResult>
) =>
defineTool({
description:
'Read bounded structure from a persisted document canvas. Use overview first, then a frame, element ids, or region for detail. Use doc_read for Page text and frontend tools for unsynced editor state. Cursors are revision-bound.',
inputSchema: z
.object({
doc_id: z.string().min(1).max(128),
target: targetSchema,
cursor: z.string().max(2048).optional(),
limit: z.number().int().min(1).max(MAX_LIMIT).optional(),
})
.strict(),
execute: async ({ doc_id, target, cursor, limit }) => {
try {
return await readCanvas(doc_id, target, cursor, limit);
} catch {
logger.error(`Failed to read canvas ${doc_id}: DOC_CANVAS_READ_FAILED`);
return toolError(
'Doc Canvas Read Failed',
'The persisted canvas could not be read.',
{
code: 'DOC_CANVAS_READ_FAILED',
retryable: false,
locator: { doc_id },
}
);
}
},
});
packages/backend/server/src/plugins/copilot/tools/doc-keyword-search.ts
-87
View File
@@ -1,87 +0,0 @@
import { z } from 'zod';
import type { PermissionAccess } from '../../../core/permission';
import type { Models } from '../../../models';
import type { IndexerService, SearchDoc } from '../../indexer';
import { workspaceSyncRequiredError } from './doc-sync';
import { toolError } from './error';
import { defineTool } from './tool';
import type { CopilotChatOptions } from './types';
export const buildDocKeywordSearchGetter = (
ac: PermissionAccess,
indexerService: IndexerService,
models: Models
) => {
const searchDocs = async (options: CopilotChatOptions, query?: string) => {
const queryTrimmed = query?.trim();
if (!options || !queryTrimmed || !options.user || !options.workspace) {
return toolError(
'Doc Keyword Search Failed',
'Missing workspace, user, or query for doc_keyword_search.'
);
}
const workspace = await models.workspace.get(options.workspace);
if (!workspace) {
return workspaceSyncRequiredError();
}
const canAccess = await ac
.user(options.user)
.workspace(options.workspace)
.can('Workspace.Read');
if (!canAccess) {
return toolError(
'Doc Keyword Search Failed',
'You do not have permission to access this workspace.'
);
}
const docs = await indexerService.searchDocsByKeyword(
options.workspace,
queryTrimmed
);
// filter current user readable docs
const readableDocs = await ac
.user(options.user)
.workspace(options.workspace)
.docs(docs, 'Doc.Read');
return readableDocs ?? [];
};
return searchDocs;
};
export const createDocKeywordSearchTool = (
searchDocs: (
query: string
) => Promise<SearchDoc[] | ReturnType<typeof toolError>>
) => {
return defineTool({
description:
'Fuzzy search all workspace documents for the exact keyword or phrase supplied and return passages ranked by textual match. Use this tool by default whenever a straightforward term-based or keyword-base lookup is sufficient.',
inputSchema: z.object({
query: z
.string()
.describe(
'The query to search for, e.g. "meeting notes" or "project plan".'
),
}),
execute: async ({ query }) => {
try {
const docs = await searchDocs(query);
if (!Array.isArray(docs)) {
return docs;
}
return docs.map(doc => ({
docId: doc.docId,
title: doc.title,
createdAt: doc.createdAt,
updatedAt: doc.updatedAt,
createdByUser: doc.createdByUser,
updatedByUser: doc.updatedByUser,
}));
} catch (e: any) {
return toolError('Doc Keyword Search Failed', e.message);
}
},
});
};
packages/backend/server/src/plugins/copilot/tools/doc-read.ts
+69 -24
View File
@@ -10,7 +10,12 @@ import {
} from './doc-sync';
import { type ToolError, toolError } from './error';
import { defineTool } from './tool';
import type { CopilotChatOptions } from './types';
import {
type CopilotChatOptions,
type DocSource,
type DocumentScope,
isDocumentInScope,
} from './types';
const logger = new Logger('DocReadTool');
@@ -20,13 +25,30 @@ const isToolError = (result: ToolError | object): result is ToolError =>
export const buildDocContentGetter = (
ac: PermissionAccess,
docReader: DocReader,
models: Models
models: Models,
documentScope?: DocumentScope
) => {
const getDoc = async (options: CopilotChatOptions, docId?: string) => {
const getDoc = async (
options: CopilotChatOptions,
docId?: string,
maxChars = 40_000
) => {
if (!options?.user || !options?.workspace || !docId) {
return toolError(
'Doc Read Failed',
'Missing workspace, user, or document id for doc_read.'
'Missing workspace, user, or document id for doc_read.',
{ code: 'INVALID_CONTEXT', retryable: false }
);
}
if (!isDocumentInScope(documentScope, docId)) {
return toolError(
'Doc Read Failed',
'The document is outside the user-selected document scope.',
{
code: 'DOC_SCOPE_DENIED',
retryable: false,
locator: { doc_id: docId },
}
);
}
@@ -44,10 +66,11 @@ export const buildDocContentGetter = (
logger.warn(
`User ${options.user} does not have access to doc ${docId} in workspace ${options.workspace}`
);
return toolError(
'Doc Read Failed',
`You do not have permission to read document ${docId} in this workspace.`
);
return toolError('Doc Read Failed', 'Document access denied.', {
code: 'DOC_ACCESS_DENIED',
retryable: false,
locator: { doc_id: docId },
});
}
const docMeta = await models.doc.getAuthors(options.workspace, docId);
@@ -64,14 +87,22 @@ export const buildDocContentGetter = (
return documentSyncPendingError(docId);
}
const markdown = content.markdown.slice(0, maxChars);
return {
docId,
doc_id: docId,
title: content.title,
markdown: content.markdown,
createdAt: docMeta.createdAt,
updatedAt: docMeta.updatedAt,
createdByUser: docMeta.createdByUser,
updatedByUser: docMeta.updatedByUser,
markdown,
revision: content.revision,
max_chars: maxChars,
truncated: markdown.length < content.markdown.length,
source: {
type: 'document' as const,
workspace_id: options.workspace,
doc_id: docId,
title: content.title,
revision: content.revision,
visibility: 'page' as const,
} satisfies DocSource,
};
};
return getDoc;
@@ -82,21 +113,35 @@ type DocReadToolResult = Awaited<
>;
export const createDocReadTool = (
getDoc: (targetId?: string) => Promise<DocReadToolResult>
getDoc: (targetId?: string, maxChars?: number) => Promise<DocReadToolResult>
) => {
return defineTool({
description:
'Return the complete text and basic metadata of a single document identified by docId; use this when the user needs the full content of a specific file rather than a search result.',
inputSchema: z.object({
doc_id: z.string().describe('The target doc to read'),
}),
execute: async ({ doc_id }) => {
'Read Page-mode text from a persisted document. Use doc_canvas_read for canvas-only content and frontend read tools for unsynced editor state. The result includes its persisted revision and may be truncated.',
inputSchema: z
.object({
doc_id: z
.string()
.min(1)
.max(128)
.describe('The persisted document to read'),
max_chars: z.number().int().min(1).max(100_000).optional(),
})
.strict(),
execute: async ({ doc_id, max_chars }) => {
try {
const doc = await getDoc(doc_id);
const doc = await getDoc(
doc_id,
Math.min(max_chars ?? 40_000, 100_000)
);
return isToolError(doc) ? doc : { ...doc };
} catch (err: any) {
logger.error(`Failed to read the doc ${doc_id}`, err);
return toolError('Doc Read Failed', err.message ?? String(err));
} catch {
logger.error(`Failed to read doc ${doc_id}: DOC_READ_FAILED`);
return toolError(
'Doc Read Failed',
'The persisted Page content could not be read.',
{ code: 'DOC_READ_FAILED', retryable: false, locator: { doc_id } }
);
}
},
});
packages/backend/server/src/plugins/copilot/tools/doc-search.ts
+132
View File
@@ -0,0 +1,132 @@
import { Logger } from '@nestjs/common';
import { z } from 'zod';
import type { DocumentRetrievalService } from '../retrieval/document';
import { toolError } from './error';
import { defineTool } from './tool';
import type { CopilotChatOptions, DocSource, DocumentScope } from './types';
const logger = new Logger('DocSearchTool');
export const buildDocumentSearch = (
retrieval: DocumentRetrievalService,
options: CopilotChatOptions,
documentScope?: DocumentScope
) => {
return async (
query: string,
docIds: string[] | undefined,
limit: number,
signal?: AbortSignal
) => {
if (!options?.workspace || !options.user) {
return toolError('Document Search Failed', 'Missing workspace or user.', {
code: 'INVALID_CONTEXT',
retryable: false,
});
}
const workspaceId = options.workspace;
const allowed = documentScope
? new Set(documentScope.allowedDocIds)
: undefined;
const effectiveDocIds = allowed
? [...allowed]
: docIds?.length
? docIds
: undefined;
if (allowed?.size === 0) {
return {
scope_mode: 'selected' as const,
scope_doc_count: 0,
retrieval_mode: 'scoped',
degraded_reason: undefined,
hits: [],
};
}
try {
const result = await retrieval.search(
options,
query,
effectiveDocIds,
limit,
signal
);
return {
scope_mode: documentScope
? ('selected' as const)
: ('workspace' as const),
scope_doc_count: documentScope?.allowedDocIds.length,
retrieval_mode: result.retrievalMode,
degraded_reason: result.degradedReason,
hits: result.hits.map(hit => ({
doc_id: hit.docId,
title: hit.title,
excerpt: hit.excerpt,
visibility: hit.visibility,
block_id: hit.blockId,
element_id: hit.elementId,
frame_id: hit.frameId,
updated_at: hit.updatedAt,
score: hit.score,
source: {
type: 'document' as const,
workspace_id: workspaceId,
doc_id: hit.docId,
title: hit.title,
visibility: hit.visibility,
block_id: hit.blockId,
element_id: hit.elementId,
frame_id: hit.frameId,
} satisfies DocSource,
})),
};
} catch (error) {
const unavailable =
error instanceof Error && error.message === 'SEARCH_UNAVAILABLE';
const code = unavailable
? 'SEARCH_UNAVAILABLE'
: 'DOCUMENT_SEARCH_FAILED';
logger.error(`Document search failed: ${code}`);
return toolError(
'Document Search Failed',
'Document search is unavailable.',
{
code,
retryable: unavailable,
}
);
}
};
};
type DocumentSearchResult = Awaited<
ReturnType<ReturnType<typeof buildDocumentSearch>>
>;
export const createDocSearchTool = (
search: (
query: string,
docIds: string[] | undefined,
limit: number,
signal?: AbortSignal
) => Promise<DocumentSearchResult>
) =>
defineTool({
description:
'Search persisted workspace documents and return bounded passages with Page or canvas locators. The runtime chooses hybrid, lexical, or vector retrieval. This tool never searches files, blobs, session attachments, or the web.',
inputSchema: z
.object({
query: z.string().trim().min(1).max(2000),
doc_ids: z
.array(z.string().min(1).max(128))
.max(50)
.optional()
.describe(
'Restrict workspace search to these document ids. When the user selected documents above the chat input, the complete selected scope is always searched instead.'
),
limit: z.number().int().min(1).max(20).optional(),
})
.strict(),
execute: ({ query, doc_ids, limit }, options) =>
search(query, doc_ids, Math.min(limit ?? 10, 20), options.signal),
});
packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts
-162
View File
@@ -1,162 +0,0 @@
/* oxlint-disable import/no-cycle -- Semantic search uses the shared embedding runtime. */
import { omit } from 'lodash-es';
import { z } from 'zod';
import type { PermissionAccess } from '../../../core/permission';
import {
type ChunkSimilarity,
clearEmbeddingChunk,
type Models,
} from '../../../models';
import { CopilotContextService } from '../context/service';
import { workspaceSyncRequiredError } from './doc-sync';
import { toolError } from './error';
import { defineTool } from './tool';
import type { CopilotChatOptions } from './types';
const getEmbeddingRouteContext = (options: CopilotChatOptions) => ({
userId: options?.user,
byokLeaseId: options?.byokLeaseId,
});
export const buildDocSearchGetter = (
ac: PermissionAccess,
context: CopilotContextService,
sessionId: string | undefined,
models: Models
) => {
const searchDocs = async (
options: CopilotChatOptions,
query?: string,
signal?: AbortSignal
) => {
if (!options || !query?.trim() || !options.user || !options.workspace) {
return toolError(
'Doc Semantic Search Failed',
'Missing workspace, user, or query for doc_semantic_search.'
);
}
const workspace = await models.workspace.get(options.workspace);
if (!workspace) {
return workspaceSyncRequiredError();
}
const canAccess = await ac
.user(options.user)
.workspace(options.workspace)
.can('Workspace.Read');
if (!canAccess)
return toolError(
'Doc Semantic Search Failed',
'You do not have permission to access this workspace.'
);
const routeContext = getEmbeddingRouteContext(options);
const [chunks, contextChunks] = await Promise.all([
context.matchWorkspaceAll(
options.workspace,
query,
10,
signal,
0.8,
undefined,
0.85,
routeContext
),
sessionId
? context
.getBySessionId(sessionId)
.then(
current =>
current?.matchFiles(
query,
10,
signal,
0.85,
0.5,
routeContext
) ?? []
)
: [],
]);
const docChunks = await ac
.user(options.user)
.workspace(options.workspace)
.docs(
chunks.filter(c => 'docId' in c),
'Doc.Read'
);
const blobChunks = chunks.filter(c => 'blobId' in c);
const fileChunks = chunks.filter(c => 'fileId' in c);
if (contextChunks.length) {
fileChunks.push(...contextChunks);
}
if (!blobChunks.length && !docChunks.length && !fileChunks.length) {
return [];
}
const docIds = docChunks.map(c => ({
// oxlint-disable-next-line no-non-null-assertion
workspaceId: options.workspace!,
docId: c.docId,
}));
const docAuthors = await models.doc
.findAuthors(docIds)
.then(
docs =>
new Map(
docs
.filter(d => !!d)
.map(doc => [doc.id, omit(doc, ['id', 'workspaceId'])])
)
);
const docMetas = await models.doc
.findMetas(docIds, { select: { title: true } })
.then(
docs =>
new Map(
docs
.filter(d => !!d)
.map(doc => [
doc.docId,
Object.assign({}, doc, docAuthors.get(doc.docId)),
])
)
);
return [
...fileChunks.map(clearEmbeddingChunk),
...blobChunks.map(clearEmbeddingChunk),
...docChunks.map(c => ({
...c,
...docMetas.get(c.docId),
})),
] as ChunkSimilarity[];
};
return searchDocs;
};
export const createDocSemanticSearchTool = (
searchDocs: (
query: string,
signal?: AbortSignal
) => Promise<ChunkSimilarity[] | ReturnType<typeof toolError>>
) => {
return defineTool({
description:
'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts, recent documents).',
inputSchema: z.object({
query: z
.string()
.describe(
'The query statement to search for, e.g. "What is the capital of France?"\nWhen querying specific terms or IDs, you should provide the complete string instead of separating it with delimiters.\nFor example, if a user wants to look up the ID "sicDoe1is", use "What is sicDoe1is" instead of "si code 1is".'
),
}),
execute: async ({ query }, options) => {
try {
return await searchDocs(query, options.signal);
} catch (e: any) {
return toolError('Doc Semantic Search Failed', e.message);
}
},
});
};
packages/backend/server/src/plugins/copilot/tools/error.ts
+8 -5
View File
@@ -2,10 +2,13 @@ export interface ToolError {
type: 'error';
name: string;
message: string;
code?: string;
retryable?: boolean;
locator?: Record<string, string>;
}
export const toolError = (name: string, message: string): ToolError => ({
type: 'error',
name,
message,
});
export const toolError = (
name: string,
message: string,
details: Pick<ToolError, 'code' | 'retryable' | 'locator'> = {}
): ToolError => ({ type: 'error', name, message, ...details });
packages/backend/server/src/plugins/copilot/tools/frontend-read.ts
+87
View File
@@ -0,0 +1,87 @@
import type { DelegatedToolName } from '@affine/realtime';
import { z } from 'zod';
import type { DelegatedEditorService } from '../delegated/service';
import type { CopilotChatOptions } from '../providers/types';
import { type CopilotToolExecuteOptions, defineTool } from './tool';
const execute =
(
delegated: DelegatedEditorService,
options: CopilotChatOptions,
tool: DelegatedToolName
) =>
(args: Record<string, unknown>, execution: CopilotToolExecuteOptions) =>
delegated.execute(options, tool, args, execution.signal, execution);
export function createFrontendEditorStateTool(
delegated: DelegatedEditorService,
options: CopilotChatOptions
) {
const run = execute(delegated, options, 'frontend_get_editor_state');
return defineTool({
description:
'Get lightweight state for the focused live editor: mode, readonly state, selection locator, capabilities, and editor_state_id. Use before live reads when freshness matters. It does not return document content.',
inputSchema: z.object({}).strict(),
execute: run,
});
}
export function createFrontendSelectionTool(
delegated: DelegatedEditorService,
options: CopilotChatOptions
) {
const run = execute(delegated, options, 'frontend_read_selection');
return defineTool({
description:
'Read the current Page or Edgeless selection from the focused live editor. Use for unsynced selected content; results are bounded and include editor_state_id and truncation. Do not use for persisted documents outside the active editor.',
inputSchema: z
.object({
format: z.enum(['text', 'markdown', 'structure']).optional(),
limit: z.number().int().min(1).max(50_000).optional(),
neighborhood: z.number().int().min(0).max(20).optional(),
})
.strict(),
execute: run,
});
}
export function createFrontendNodesTool(
delegated: DelegatedEditorService,
options: CopilotChatOptions
) {
const run = execute(delegated, options, 'frontend_read_nodes');
return defineTool({
description:
'Read bounded live blocks or canvas elements by ids in the focused editor. Use locators returned by editor state, selection, or snapshot tools. Each item can return its own error; ids must belong to the active document.',
inputSchema: z
.object({
block_ids: z.array(z.string().min(1).max(128)).max(50).optional(),
element_ids: z.array(z.string().min(1).max(128)).max(50).optional(),
limit: z.number().int().min(1).max(50_000).optional(),
})
.strict()
.refine(value => value.block_ids?.length || value.element_ids?.length, {
message: 'block_ids or element_ids is required',
}),
execute: run,
});
}
export function createFrontendSnapshotTool(
delegated: DelegatedEditorService,
options: CopilotChatOptions
) {
const run = execute(delegated, options, 'frontend_snapshot_document');
return defineTool({
description:
'Get a lightweight view from the focused editor. Page mode returns an outline or selection neighborhood; Edgeless mode returns the visible viewport. The requested view is adapted to the active editor mode. Use it to locate content before targeted reads; it is bounded and is not a full document snapshot.',
inputSchema: z
.object({
view: z.enum(['outline', 'selection_neighborhood', 'viewport']),
limit: z.number().int().min(1).max(200).optional(),
})
.strict(),
execute: run,
});
}
packages/backend/server/src/plugins/copilot/tools/index.ts
+4 -4
View File
@@ -1,14 +1,14 @@
/* oxlint-disable import/no-cycle -- Tool exports include semantic search runtime dependencies. */
export * from './blob-read';
export * from './artifact';
export * from './code-artifact';
export * from './conversation-summary';
export * from './doc-canvas-read';
export * from './doc-compose';
export * from './doc-keyword-search';
export * from './doc-read';
export * from './doc-semantic-search';
export * from './doc-search';
export * from './doc-write';
export * from './error';
export * from './exa-crawl';
export * from './exa-search';
export * from './frontend-read';
export * from './section-edit';
export * from './tool';
packages/backend/server/src/plugins/copilot/tools/tool.ts
+2
View File
@@ -7,6 +7,8 @@ import { toToolJsonSchema } from './json-schema';
export type CopilotToolExecuteOptions = {
signal?: AbortSignal;
messages?: PromptMessage[];
runId?: string;
toolCallId?: string;
};
export type CopilotTool = {
packages/backend/server/src/plugins/copilot/tools/types.ts
+33 -2
View File
@@ -1,3 +1,34 @@
export type { CopilotContextService } from '../context/service';
export type { ContextSession } from '../context/session';
export type { CopilotChatOptions } from '../providers/types';
export type DocumentScope = {
mode: 'selected';
allowedDocIds: readonly string[];
};
export function isDocumentInScope(
scope: DocumentScope | undefined,
docId: string
) {
return !scope || scope.allowedDocIds.includes(docId);
}
export type DocSource = {
type: 'document';
workspace_id: string;
doc_id: string;
title: string;
revision?: string;
visibility?: 'page' | 'edgeless' | 'both';
block_id?: string;
element_id?: string;
frame_id?: string;
};
export type ArtifactSource = {
type: 'artifact';
workspace_id: string;
artifact_id: string;
name?: string;
mime_type?: string;
revision?: string;
};
packages/backend/server/src/plugins/copilot/types.ts
+6 -7
View File
@@ -3,6 +3,10 @@ import { z } from 'zod';
import type { Turn } from './core/types';
import type { ResolvedPrompt } from './prompt';
import { PromptMessageSchema, PureMessageSchema } from './providers/types';
import {
type SessionFocus,
TurnScopeSnapshotSchema,
} from './runtime/contracts/shared';
const takeFirst = (v: unknown) => (Array.isArray(v) ? v[0] : v);
@@ -93,6 +97,7 @@ export const ChatQuerySchema = z
export const ChatMessageSchema = PromptMessageSchema.extend({
id: z.string().optional(),
scopeSnapshot: TurnScopeSnapshotSchema.nullable().optional(),
createdAt: z.date(),
}).strict();
export type ChatMessage = z.infer<typeof ChatMessageSchema>;
@@ -148,12 +153,6 @@ export type ChatSessionState = {
workspaceId: string;
docId: string | null;
turns: Turn[];
focus: SessionFocus;
prompt: ResolvedPrompt;
};
export type CopilotContextFile = {
id: string; // fileId
created_at: number;
// embedding status
status: 'in_progress' | 'completed' | 'failed';
};
packages/backend/server/src/plugins/copilot/utils.ts
+9 -2
View File
@@ -67,14 +67,21 @@ export function getTools(
case 'searchWorkspace':
if (value === false) {
result = result.filter(tool => {
return tool !== 'docKeywordSearch' && tool !== 'docSemanticSearch';
return tool !== 'docSearch';
});
}
break;
case 'readingDocs':
if (value === false) {
result = result.filter(tool => {
return tool !== 'docRead';
return ![
'docRead',
'docCanvasRead',
'frontendGetEditorState',
'frontendReadSelection',
'frontendReadNodes',
'frontendSnapshotDocument',
].includes(tool);
});
}
break;
packages/backend/server/src/plugins/copilot/workspace/resolver.ts
+23 -36
View File
@@ -16,7 +16,7 @@ import GraphQLUpload, {
import {
BlobQuotaExceeded,
CopilotEmbeddingUnavailable,
CopilotFailedToAddWorkspaceFileEmbedding,
CopilotFailedToAddWorkspaceArtifact,
Mutex,
paginate,
PaginationInput,
@@ -31,9 +31,9 @@ import { COPILOT_LOCKER } from '../resolver';
import { MAX_EMBEDDABLE_SIZE } from '../utils';
import { CopilotWorkspaceService } from './service';
import {
CopilotWorkspaceFileType,
CopilotWorkspaceArtifactType,
CopilotWorkspaceIgnoredDocType,
PaginatedCopilotWorkspaceFileType,
PaginatedCopilotWorkspaceArtifactType,
PaginatedIgnoredDocsType,
} from './types';
@@ -132,34 +132,34 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
);
}
@ResolveField(() => PaginatedCopilotWorkspaceFileType, {
@ResolveField(() => PaginatedCopilotWorkspaceArtifactType, {
complexity: 2,
})
async files(
async artifacts(
@Parent() config: CopilotWorkspaceConfigType,
@Args('pagination', PaginationInput.decode) pagination: PaginationInput
): Promise<PaginatedCopilotWorkspaceFileType> {
const [files, totalCount] = await this.copilotWorkspace.listFiles(
): Promise<PaginatedCopilotWorkspaceArtifactType> {
const [artifacts, totalCount] = await this.copilotWorkspace.listArtifacts(
config.workspaceId,
pagination
);
return paginate(files, 'createdAt', pagination, totalCount);
return paginate(artifacts, 'createdAt', pagination, totalCount);
}
@Mutation(() => CopilotWorkspaceFileType, {
name: 'addWorkspaceEmbeddingFiles',
@Mutation(() => CopilotWorkspaceArtifactType, {
name: 'addWorkspaceArtifact',
complexity: 2,
description: 'Update workspace embedding files',
description: 'Add a workspace artifact',
})
async addFiles(
async addArtifact(
@Context() ctx: { req: Request },
@CurrentUser() user: CurrentUser,
@Args('workspaceId', { type: () => String })
workspaceId: string,
@Args({ name: 'blob', type: () => GraphQLUpload })
content: FileUpload
): Promise<CopilotWorkspaceFileType> {
): Promise<CopilotWorkspaceArtifactType> {
await this.ac
.user(user.id)
.workspace(workspaceId)
@@ -181,48 +181,35 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
}
try {
const { blobId, file } = await this.copilotWorkspace.addFile(
user.id,
workspaceId,
content
);
await this.copilotWorkspace.queueFileEmbedding({
userId: user.id,
workspaceId,
blobId,
fileId: file.fileId,
fileName: file.fileName,
});
return file;
} catch (e: any) {
return await this.copilotWorkspace.addArtifact(workspaceId, content);
} catch (e) {
// passthrough user friendly error
if (e instanceof UserFriendlyError) {
throw e;
}
throw new CopilotFailedToAddWorkspaceFileEmbedding({
message: e.message,
throw new CopilotFailedToAddWorkspaceArtifact({
message: e instanceof Error ? e.message : String(e),
});
}
}
@Mutation(() => Boolean, {
name: 'removeWorkspaceEmbeddingFiles',
name: 'removeWorkspaceArtifact',
complexity: 2,
description: 'Remove workspace embedding files',
description: 'Remove a workspace artifact',
})
async removeFiles(
async removeArtifact(
@CurrentUser() user: CurrentUser,
@Args('workspaceId', { type: () => String })
workspaceId: string,
@Args('fileId', { type: () => String })
fileId: string
@Args('artifactId', { type: () => String })
artifactId: string
): Promise<boolean> {
await this.ac
.user(user.id)
.workspace(workspaceId)
.assert('Workspace.Settings.Update');
return await this.copilotWorkspace.removeFile(workspaceId, fileId);
return await this.copilotWorkspace.removeArtifact(workspaceId, artifactId);
}
}
packages/backend/server/src/plugins/copilot/workspace/service.ts
+118 -44
View File
@@ -1,16 +1,11 @@
import { createHash } from 'node:crypto';
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import { Prisma, PrismaClient } from '@prisma/client';
import {
FileUpload,
JobQueue,
PaginationInput,
sniffMime,
} from '../../../base';
import { FileUpload, PaginationInput, sniffMime } from '../../../base';
import { ServerFeature, ServerService } from '../../../core';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { NativeEmbeddingService } from '../embedding/native';
import { readStream } from '../utils';
@Injectable()
@@ -20,14 +15,14 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
constructor(
private readonly server: ServerService,
private readonly models: Models,
private readonly queue: JobQueue,
private readonly storage: CopilotStorage
private readonly embedding: NativeEmbeddingService,
private readonly runtime: BackendRuntimeProvider,
private readonly db: PrismaClient
) {}
async onApplicationBootstrap() {
const supportEmbedding =
await this.models.copilotWorkspace.checkEmbeddingAvailable();
if (supportEmbedding) {
const health = await this.embedding.health();
if (health.enabled) {
this.server.enableFeature(ServerFeature.CopilotEmbedding);
this.supportEmbedding = true;
}
@@ -61,48 +56,127 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
]);
}
async addFile(userId: string, workspaceId: string, content: FileUpload) {
const fileName = content.filename;
async addArtifact(workspaceId: string, content: FileUpload) {
const buffer = await readStream(content.createReadStream());
const blobId = createHash('sha256').update(buffer).digest('base64url');
await this.storage.put(userId, workspaceId, blobId, buffer);
const file = await this.models.copilotWorkspace.addFile(workspaceId, {
fileName,
blobId,
mimeType: sniffMime(buffer, content.mimetype) || content.mimetype,
size: buffer.length,
const artifact = await this.runtime.putWorkspaceArtifact(
{
workspaceId,
mimeType: sniffMime(buffer, content.mimetype) || content.mimetype,
displayName: content.filename,
fileName: content.filename,
libraryOwned: true,
},
buffer
);
return await this.getArtifact(workspaceId, artifact.id);
}
async getArtifact(workspaceId: string, artifactId: string) {
const artifact = await this.db.workspaceArtifact.findUniqueOrThrow({
where: {
id: artifactId,
workspaceId,
libraryOwned: true,
status: 'ready',
},
});
return { blobId, file };
const statuses = await this.embeddingStatuses(workspaceId, [artifact.id]);
return this.projectArtifact(
artifact,
statuses.get(artifact.id) ?? 'processing'
);
}
async getFile(workspaceId: string, fileId: string) {
return await this.models.copilotWorkspace.getFile(workspaceId, fileId);
}
async listFiles(
async listArtifacts(
workspaceId: string,
pagination?: {
includeRead?: boolean;
} & PaginationInput
) {
return await Promise.all([
this.models.copilotWorkspace.listFiles(workspaceId, pagination),
this.models.copilotWorkspace.countFiles(workspaceId),
const where = { workspaceId, libraryOwned: true, status: 'ready' };
const [artifacts, count] = await Promise.all([
this.db.workspaceArtifact.findMany({
where,
orderBy: { createdAt: 'desc' },
skip: pagination?.offset,
take: pagination?.first,
}),
this.db.workspaceArtifact.count({ where }),
]);
}
async queueFileEmbedding(file: Jobs['copilot.embedding.files']) {
const { userId, workspaceId, blobId, fileId, fileName } = file;
await this.queue.add('copilot.embedding.files', {
userId,
const statuses = await this.embeddingStatuses(
workspaceId,
blobId,
fileId,
fileName,
});
artifacts.map(artifact => artifact.id)
);
return [
artifacts.map(artifact =>
this.projectArtifact(
artifact,
statuses.get(artifact.id) ?? 'processing'
)
),
count,
] as const;
}
async removeFile(workspaceId: string, fileId: string) {
return await this.models.copilotWorkspace.removeFile(workspaceId, fileId);
async removeArtifact(workspaceId: string, artifactId: string) {
await this.runtime.setArtifactLibraryOwned(workspaceId, artifactId, false);
return true;
}
private async embeddingStatuses(workspaceId: string, artifactIds: string[]) {
if (artifactIds.length === 0)
return new Map<string, 'processing' | 'ready' | 'failed'>();
const rows = await this.db.$queryRaw<
{ artifactId: string; status: 'processing' | 'ready' | 'failed' }[]
>`
SELECT artifact.id::text AS "artifactId",
CASE
WHEN projection.status='ready'
AND projection.applied_content_revision=source.content_revision THEN 'ready'
WHEN projection.status='failed' THEN 'failed'
ELSE 'processing'
END AS status
FROM workspace_artifacts artifact
LEFT JOIN embedding_sources source
ON source.workspace_id=artifact.workspace_id
AND source.source_kind='artifact'
AND source.source_key=artifact.id::text
AND source.deleted_at IS NULL
LEFT JOIN embedding_workspace_states state
ON state.workspace_id=artifact.workspace_id
LEFT JOIN embedding_projections projection
ON projection.source_id=source.id
AND projection.index_id=state.active_index_id
WHERE artifact.workspace_id=${workspaceId}
AND artifact.id::text IN (${Prisma.join(artifactIds)})
`;
return new Map(rows.map(row => [row.artifactId, row.status]));
}
private projectArtifact(
artifact: {
id: string;
workspaceId: string;
contentHash: string;
displayName: string | null;
canonicalMediaType: string;
sizeBytes: bigint;
createdAt: Date;
},
embeddingStatus: 'processing' | 'ready' | 'failed'
) {
if (!artifact.displayName) {
throw new Error('Library artifact display name is missing');
}
return {
workspaceId: artifact.workspaceId,
artifactId: artifact.id,
contentHash: artifact.contentHash,
fileName: artifact.displayName,
embeddingStatus,
mediaType: artifact.canonicalMediaType,
size: Number(artifact.sizeBytes),
createdAt: artifact.createdAt,
};
}
}
packages/backend/server/src/plugins/copilot/workspace/types.ts
+11 -19
View File
@@ -2,18 +2,7 @@ import { Field, ObjectType } from '@nestjs/graphql';
import { SafeIntResolver } from 'graphql-scalars';
import { Paginated } from '../../../base';
import { CopilotWorkspaceFile, IgnoredDoc } from '../../../models';
declare global {
interface Events {
'workspace.file.embedding.finished': {
jobId: string;
};
'workspace.file.embedding.failed': {
jobId: string;
};
}
}
import { CopilotWorkspaceArtifact, IgnoredDoc } from '../../../models';
@ObjectType('CopilotWorkspaceIgnoredDoc')
export class CopilotWorkspaceIgnoredDocType implements IgnoredDoc {
@@ -47,22 +36,25 @@ export class PaginatedIgnoredDocsType extends Paginated(
CopilotWorkspaceIgnoredDocType
) {}
@ObjectType('CopilotWorkspaceFile')
export class CopilotWorkspaceFileType implements CopilotWorkspaceFile {
@ObjectType('CopilotWorkspaceArtifact')
export class CopilotWorkspaceArtifactType implements CopilotWorkspaceArtifact {
@Field(() => String)
workspaceId!: string;
@Field(() => String)
fileId!: string;
artifactId!: string;
@Field(() => String)
blobId!: string;
contentHash!: string;
@Field(() => String)
fileName!: string;
@Field(() => String)
mimeType!: string;
embeddingStatus!: 'processing' | 'ready' | 'failed';
@Field(() => String)
mediaType!: string;
@Field(() => SafeIntResolver)
size!: number;
@@ -72,6 +64,6 @@ export class CopilotWorkspaceFileType implements CopilotWorkspaceFile {
}
@ObjectType()
export class PaginatedCopilotWorkspaceFileType extends Paginated(
CopilotWorkspaceFileType
export class PaginatedCopilotWorkspaceArtifactType extends Paginated(
CopilotWorkspaceArtifactType
) {}
packages/backend/server/src/plugins/indexer/__tests__/__snapshots__/service.spec.ts.md
+10 -1
View File
@@ -461,7 +461,16 @@ Generated by [AVA](https://avajs.dev).
{
summary: [
'AFFiNE is an open source all in one workspace, an operating system for all the building blocks of your team wiki, knowledge management and digital assets and a better alternative to Notion and Miro. You own your data, with no compromisesLocal-first & Real-time collaborativeWe love the idea proposed by Ink & Switch in the famous article about you owning your data, despite the cloud. Furthermore, AFFiNE is the first all-in-one workspace that keeps your data ownership with no compromises on real-time collaboration and editing experience.AFFiNE is a local-first application upon CRDTs with real-time collaboration support. Your data is always stored locally while multiple nodes remain synced in real-time.Blocks that assemble your next docs, tasks kanban or whiteboardThere is a large overlap of their atomic "building blocks" between these apps. They are neither open source nor have a plugin system like VS Code for contributors to customize. We want to have something that contains all the features we love and goes one step further. ',
`We are building AFFiNE to be a fundamental open source platform that contains all the building blocks for docs, task management and visual collaboration, hoping you can shape your next workflow with us that can make your life better and also connect others, too.␊
Airtable & Miro with their no-code programable datasheets␊
␊
For developer or installation guides, please go to AFFiNE Development␊
Blocks that assemble your next docs, tasks kanban or whiteboard␊
␊
Trello with their Kanban␊
Remnote & Capacities with their object-based tag system␊
AFFiNE is an open source all in one workspace, an operating system for all the building blocks of your team wiki, knowledge management and digital assets and a better alternative to Notion and Miro. ␊
There is a large overlap of their atomic "building blocks" between these apps. They are neither open source nor have a plugin system like VS Code for contributors to customize. We want to have something that contains all the features we love and goes one step fu`,
],
title: [
'Write, Draw, Plan all at Once.',
packages/backend/server/src/plugins/indexer/__tests__/__snapshots__/service.spec.ts.snap
BIN
View File
Binary file not shown.
packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts
+29 -5
View File
@@ -2044,7 +2044,6 @@ test('should list doc ids work', async t => {
// #region indexDoc()
test('should index doc work', async t => {
const count = module.queue.count('copilot.embedding.updateDoc');
const docSnapshot = await module.create(Mockers.DocSnapshot, {
workspaceId: workspace.id,
user,
@@ -2092,7 +2091,16 @@ test('should index doc work', async t => {
],
},
options: {
fields: ['workspaceId', 'docId', 'blockId', 'content', 'flavour'],
fields: [
'workspaceId',
'docId',
'blockId',
'unitId',
'projectionVersion',
'sourceHash',
'content',
'flavour',
],
highlights: [
{
field: 'content',
@@ -2107,10 +2115,26 @@ test('should index doc work', async t => {
});
t.is(result2.nodes.length, 2);
t.snapshot(
result2.nodes.map(node => omit(node.fields, ['workspaceId', 'docId']))
t.true(
result2.nodes.every(
node =>
node.fields.unitId.length === 1 &&
node.fields.projectionVersion[0] === 1 &&
node.fields.sourceHash.length === 1
)
);
t.is(new Set(result2.nodes.map(node => node.fields.sourceHash[0])).size, 1);
t.snapshot(
result2.nodes.map(node =>
omit(node.fields, [
'workspaceId',
'docId',
'unitId',
'projectionVersion',
'sourceHash',
])
)
);
t.is(module.queue.count('copilot.embedding.updateDoc'), count + 1);
});
// #endregion
packages/backend/server/src/plugins/indexer/service.ts
+55 -23
View File
@@ -6,7 +6,7 @@ import {
JobQueue,
SearchProviderNotFound,
} from '../../base';
import { readAllBlocksFromDocSnapshot } from '../../core/utils/blocksuite';
import { projectDocSearch } from '../../core/utils/blocksuite';
import { Models } from '../../models';
import { SearchProviderType } from './config';
import { SearchProviderFactory } from './factory';
@@ -284,9 +284,10 @@ export class IndexerService {
};
try {
const result = await readAllBlocksFromDocSnapshot(
const projection = projectDocSearch(
docSnapshot.blob,
docId,
docSnapshot.blob
docSnapshot.updatedAt.getTime().toString()
);
await this.write(
SearchTable.doc,
@@ -294,8 +295,11 @@ export class IndexerService {
{
workspaceId,
docId,
title: result.title,
summary: result.summary,
title: projection.title,
summary: projection.units
.map(unit => unit.text)
.join('\n')
.slice(0, 1000),
// NOTE(@fengmk): journal is not supported yet
// journal: result.journal,
createdByUserId: docSnapshot.createdBy ?? '',
@@ -309,20 +313,25 @@ export class IndexerService {
await this.deleteBlocksByDocId(workspaceId, docId, options);
await this.write(
SearchTable.block,
result.blocks.map(block => ({
projection.units.map(unit => ({
workspaceId,
docId,
blockId: block.blockId,
content: block.content ?? '',
flavour: block.flavour,
blob: block.blob,
refDocId: block.refDocId,
ref: block.ref,
parentFlavour: block.parentFlavour,
parentBlockId: block.parentBlockId,
additional: block.additional
? JSON.stringify(block.additional)
: undefined,
blockId: unit.blockId ?? unit.unitId,
unitId: unit.unitId,
projectionVersion: projection.version,
sourceHash: projection.sourceHash,
visibility: unit.visibility,
elementId: unit.elementId,
frameId: unit.frameId,
sourceBlockId: unit.blockId,
blob: unit.blobId,
refDocId: unit.refDocIds.length ? unit.refDocIds : undefined,
ref: unit.refs.length ? unit.refs : undefined,
content: unit.text,
flavour: `affine:${unit.type}`,
parentFlavour: unit.parentFlavour,
parentBlockId: unit.parentBlockId,
additional: unit.additional,
markdownPreview: undefined,
createdByUserId: docSnapshot.createdBy ?? '',
updatedByUserId: docSnapshot.updatedBy ?? '',
@@ -332,12 +341,8 @@ export class IndexerService {
options
);
await this.queue.add('copilot.embedding.updateDoc', {
workspaceId,
docId,
});
this.logger.verbose(
`synced doc ${workspaceId}/${docId} with ${result.blocks.length} blocks`
`synced doc ${workspaceId}/${docId} with ${projection.units.length} search units`
);
} catch (err) {
this.logger.warn(
@@ -559,6 +564,13 @@ export class IndexerService {
hits: {
fields: [
'blockId',
'unitId',
'projectionVersion',
'sourceHash',
'visibility',
'elementId',
'frameId',
'sourceBlockId',
'flavour',
'content',
'createdAt',
@@ -590,6 +602,20 @@ export class IndexerService {
for (const bucket of result.buckets) {
const docId = bucket.key;
const blockId = bucket.hits.nodes[0].fields.blockId[0] as string;
const unitId = bucket.hits.nodes[0].fields.unitId[0] as string;
const projectionVersion = bucket.hits.nodes[0].fields
.projectionVersion[0] as number;
const sourceHash = bucket.hits.nodes[0].fields.sourceHash[0] as string;
const visibility = bucket.hits.nodes[0].fields.visibility[0] as string;
const elementId = bucket.hits.nodes[0].fields.elementId?.[0] as
| string
| undefined;
const frameId = bucket.hits.nodes[0].fields.frameId?.[0] as
| string
| undefined;
const sourceBlockId = bucket.hits.nodes[0].fields.sourceBlockId?.[0] as
| string
| undefined;
const flavour = bucket.hits.nodes[0].fields.flavour[0] as string;
const content = bucket.hits.nodes[0].fields.content[0] as string;
const createdAt = bucket.hits.nodes[0].fields.createdAt[0] as Date;
@@ -611,7 +637,13 @@ export class IndexerService {
docs.push({
docId,
blockId,
blockId: sourceBlockId || blockId,
...(unitId ? { unitId } : {}),
...(projectionVersion ? { projectionVersion } : {}),
...(sourceHash ? { sourceHash } : {}),
...(visibility ? { visibility } : {}),
...(elementId ? { elementId } : {}),
...(frameId ? { frameId } : {}),
title,
highlight,
createdAt,
packages/backend/server/src/plugins/indexer/tables/block.ts
+21
View File
@@ -4,6 +4,13 @@ export const BlockSchema = z.object({
workspace_id: z.string(),
doc_id: z.string(),
block_id: z.string(),
unit_id: z.string().optional(),
projection_version: z.number().int().optional(),
source_hash: z.string().optional(),
visibility: z.string().optional(),
element_id: z.string().optional(),
frame_id: z.string().optional(),
source_block_id: z.string().optional(),
content: z.union([z.string(), z.string().array()]),
flavour: z.string(),
blob: z.union([z.string(), z.string().array()]).optional(),
@@ -75,6 +82,13 @@ export const blockMapping = {
block_id: {
type: 'keyword',
},
unit_id: { type: 'keyword' },
projection_version: { type: 'integer' },
source_hash: { type: 'keyword' },
visibility: { type: 'keyword' },
element_id: { type: 'keyword' },
frame_id: { type: 'keyword' },
source_block_id: { type: 'keyword' },
content: {
type: 'text',
analyzer: 'standard_with_cjk',
@@ -128,6 +142,13 @@ CREATE TABLE IF NOT EXISTS block (
workspace_id string attribute,
doc_id string attribute,
block_id string attribute,
unit_id string attribute,
projection_version int,
source_hash string attribute,
visibility string attribute,
element_id string attribute,
frame_id string attribute,
source_block_id string attribute,
content text,
flavour string attribute,
-- use flavour_indexed to match with boost
packages/backend/server/src/plugins/indexer/types.ts
+6
View File
@@ -45,6 +45,12 @@ registerEnumType(SearchQueryOccur, {
export interface SearchDoc {
docId: string;
blockId: string;
unitId?: string;
projectionVersion?: number;
sourceHash?: string;
visibility?: string;
elementId?: string;
frameId?: string;
title: string;
highlight: string;
createdAt: Date;