Skip to content

Commit 3eb5b32

Browse files
darkskygitdarkskygit
committed
feat: match workspace docs
1 parent 2d44780 commit 3eb5b32

File tree

10 files changed

+230
-48
lines changed

10 files changed

+230
-48
lines changed

‎packages/backend/server/schema.prisma‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -409,10 +409,10 @@ model AiSession {
409409
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
410410
deletedAt DateTime? @map("deleted_at") @db.Timestamptz(3)
411411
412-
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
413-
prompt AiPrompt @relation(fields: [promptName], references: [name], onDelete: Cascade)
414-
messages AiSessionMessage[]
415-
AiContext AiContext[]
412+
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
413+
prompt AiPrompt @relation(fields: [promptName], references: [name], onDelete: Cascade)
414+
messages AiSessionMessage[]
415+
context AiContext[]
416416
417417
@@index([userId])
418418
@@index([userId, workspaceId])
@@ -453,11 +453,11 @@ model AiContextEmbedding {
453453
}
454454

455455
model AiWorkspaceEmbedding {
456-
workspaceId String @map("workspace_id") @db.VarChar
457-
docId String @map("doc_id") @db.VarChar
456+
workspaceId String @map("workspace_id") @db.VarChar
457+
docId String @map("doc_id") @db.VarChar
458458
// a doc can be divided into multiple chunks and embedded separately.
459-
chunk Int @db.Integer
460-
content String @db.VarChar
459+
chunk Int @db.Integer
460+
content String @db.VarChar
461461
embedding Unsupported("vector(512)")
462462
463463
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)

‎packages/backend/server/src/plugins/copilot/context/resolver.ts‎

Lines changed: 53 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ import { CopilotContextService } from './service';
3737
import {
3838
type ContextFile,
3939
ContextFileStatus,
40+
DocChunkSimilarity,
4041
FileChunkSimilarity,
4142
} from './types';
4243

@@ -70,18 +71,6 @@ class RemoveContextFileInput {
7071
fileId!: string;
7172
}
7273

73-
@InputType()
74-
class MatchContextInput {
75-
@Field(() => String)
76-
contextId!: string;
77-
78-
@Field(() => String)
79-
content!: string;
80-
81-
@Field(() => SafeIntResolver, { nullable: true })
82-
limit?: number;
83-
}
84-
8574
@ObjectType('CopilotContext')
8675
export class CopilotContextType {
8776
@Field(() => ID)
@@ -132,6 +121,21 @@ class ContextWorkspaceEmbeddingStatus {
132121
embedded!: number;
133122
}
134123

124+
@ObjectType()
125+
class ContextMatchedDocChunk implements DocChunkSimilarity {
126+
@Field(() => String)
127+
docId!: string;
128+
129+
@Field(() => SafeIntResolver)
130+
chunk!: number;
131+
132+
@Field(() => String)
133+
content!: string;
134+
135+
@Field(() => Float, { nullable: true })
136+
distance!: number | null;
137+
}
138+
135139
@Throttle()
136140
@Resolver(() => CopilotType)
137141
export class CopilotContextRootResolver {
@@ -238,6 +242,7 @@ export class CopilotContextRootResolver {
238242
export class CopilotContextResolver {
239243
constructor(
240244
private readonly mutex: RequestMutex,
245+
private readonly permissions: PermissionService,
241246
private readonly context: CopilotContextService
242247
) {}
243248

@@ -388,27 +393,52 @@ export class CopilotContextResolver {
388393
@CallMetric('ai', 'context_file_remove')
389394
async matchContext(
390395
@Context() ctx: { req: Request },
391-
@Args({ name: 'options', type: () => MatchContextInput })
392-
options: MatchContextInput
396+
@Args('contextId') contextId: string,
397+
@Args('content') content: string,
398+
@Args('limit', { type: () => SafeIntResolver, nullable: true })
399+
limit?: number
393400
) {
394-
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
401+
const lockFlag = `${COPILOT_LOCKER}:context:${contextId}`;
395402
await using lock = await this.mutex.acquire(lockFlag);
396403
if (!lock) {
397404
return new TooManyRequest('Server is busy');
398405
}
399-
const session = await this.context.get(options.contextId);
406+
const session = await this.context.get(contextId);
400407

401408
try {
402-
return await session.match(
403-
options.content,
404-
options.limit,
405-
this.getSignal(ctx.req)
406-
);
409+
return await session.match(content, limit, this.getSignal(ctx.req));
407410
} catch (e: any) {
408411
throw new CopilotFailedToMatchContext({
409-
contextId: options.contextId,
412+
contextId,
413+
// don't record the large content
414+
content: content.slice(0, 512),
415+
message: e.message,
416+
});
417+
}
418+
}
419+
420+
@Mutation(() => ContextMatchedDocChunk, {
421+
description: 'match workspace doc',
422+
})
423+
@CallMetric('ai', 'context_match_workspace_doc')
424+
async matchWorkspaceContext(
425+
@CurrentUser() user: CurrentUser,
426+
@Context() ctx: { req: Request },
427+
@Args('contextId') contextId: string,
428+
@Args('content') content: string,
429+
@Args('limit', { type: () => SafeIntResolver, nullable: true })
430+
limit?: number
431+
) {
432+
const session = await this.context.get(contextId);
433+
await this.permissions.checkCloudWorkspace(session.workspaceId, user.id);
434+
435+
try {
436+
return await session.match(content, limit, this.getSignal(ctx.req));
437+
} catch (e: any) {
438+
throw new CopilotFailedToMatchContext({
439+
contextId,
410440
// don't record the large content
411-
content: options.content.slice(0, 512),
441+
content: content.slice(0, 512),
412442
message: e.message,
413443
});
414444
}

‎packages/backend/server/src/plugins/copilot/context/service.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import OpenAI from 'openai';
55
import {
66
Config,
77
CopilotInvalidContext,
8+
CopilotSessionNotFound,
89
NoCopilotProviderAvailable,
910
} from '../../../base';
1011
import { OpenAIEmbeddingClient } from './embedding';
@@ -46,9 +47,21 @@ export class CopilotContextService {
4647
}
4748

4849
async create(sessionId: string): Promise<ContextSession> {
50+
const session = await this.db.aiSession.findFirst({
51+
where: { id: sessionId },
52+
select: { workspaceId: true },
53+
});
54+
if (!session) {
55+
throw new CopilotSessionNotFound();
56+
}
57+
4958
const context = await this.db.aiContext.create({
50-
data: { sessionId, config: { files: [] } },
59+
data: {
60+
sessionId,
61+
config: { workspaceId: session.workspaceId, files: [] },
62+
},
5163
});
64+
5265
const config = ContextConfigSchema.parse(context.config);
5366
return this.cacheSession(context.id, config);
5467
}

‎packages/backend/server/src/plugins/copilot/context/session.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
ContextConfig,
1212
ContextFile,
1313
ContextFileStatus,
14+
DocChunkSimilarity,
1415
Embedding,
1516
EmbeddingClient,
1617
FileChunkSimilarity,
@@ -28,6 +29,10 @@ export class ContextSession implements AsyncDisposable {
2829
return this.contextId;
2930
}
3031

32+
get workspaceId() {
33+
return this.config.workspaceId;
34+
}
35+
3136
async listDocs() {
3237
return [...this.config.docs];
3338
}
@@ -180,6 +185,24 @@ export class ContextSession implements AsyncDisposable {
180185
`;
181186
}
182187

188+
async matchWorkspace(
189+
content: string,
190+
topK: number = 5,
191+
signal?: AbortSignal
192+
) {
193+
const embedding = await this.client
194+
.getEmbeddings([content], signal)
195+
.then(r => r?.[0]?.embedding);
196+
if (!embedding) return [];
197+
return await this.db.$queryRaw<Array<DocChunkSimilarity>>`
198+
SELECT "doc_id" as "docId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance"
199+
FROM "ai_workspace_embeddings"
200+
WHERE "workspace_id" = ${this.workspaceId}
201+
ORDER BY "distance" ASC
202+
LIMIT ${topK};
203+
`;
204+
}
205+
183206
private async saveFileRecord(
184207
fileId: string,
185208
cb: (

‎packages/backend/server/src/plugins/copilot/context/types.ts‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ export enum ContextFileStatus {
1111
}
1212

1313
export const ContextConfigSchema = z.object({
14+
workspaceId: z.string(),
1415
files: z
1516
.object({
1617
id: z.string(),
@@ -30,13 +31,20 @@ export const ContextConfigSchema = z.object({
3031
export type ContextConfig = z.infer<typeof ContextConfigSchema>;
3132
export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];
3233

33-
export type FileChunkSimilarity = {
34-
fileId: string;
34+
export type ChunkSimilarity = {
3535
chunk: number;
3636
content: string;
3737
distance: number | null;
3838
};
3939

40+
export type FileChunkSimilarity = ChunkSimilarity & {
41+
fileId: string;
42+
};
43+
44+
export type DocChunkSimilarity = ChunkSimilarity & {
45+
docId: string;
46+
};
47+
4048
export type Embedding = {
4149
/**
4250
* The index of the embedding in the list of embeddings.

‎packages/backend/server/src/schema.gql‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,13 @@ enum ContextFileStatus {
4242
processing
4343
}
4444

45+
type ContextMatchedDocChunk {
46+
chunk: SafeInt!
47+
content: String!
48+
distance: Float
49+
docId: String!
50+
}
51+
4552
type ContextMatchedFileChunk {
4653
chunk: SafeInt!
4754
content: String!
@@ -539,12 +546,6 @@ input ManageUserInput {
539546
name: String
540547
}
541548

542-
input MatchContextInput {
543-
content: String!
544-
contextId: String!
545-
limit: SafeInt
546-
}
547-
548549
type MemberNotFoundInSpaceDataType {
549550
spaceId: String!
550551
}
@@ -612,7 +613,10 @@ type Mutation {
612613
leaveWorkspace(sendLeaveMail: Boolean, workspaceId: String!, workspaceName: String @deprecated(reason: "no longer used")): Boolean!
613614

614615
"""remove a file from context"""
615-
matchContext(options: MatchContextInput!): [ContextMatchedFileChunk!]!
616+
matchContext(content: String!, contextId: String!, limit: SafeInt): [ContextMatchedFileChunk!]!
617+
618+
"""match workspace doc"""
619+
matchWorkspaceContext(content: String!, contextId: String!, limit: SafeInt): ContextMatchedDocChunk!
616620
publishPage(mode: PublicPageMode = Page, pageId: String!, workspaceId: String!): WorkspacePage!
617621

618622
"""queue workspace doc embedding"""

‎packages/frontend/graphql/src/graphql/copilot-context-file-match.gql‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
mutation matchContext($contextId: String!, $content: String!, $limit: SafeInt) {
2+
matchContext(contextId: $contextId, content: $content, limit: $limit) {
3+
fileId
4+
chunk
5+
content
6+
distance
7+
}
8+
}

‎packages/frontend/graphql/src/graphql/copilot-context-workspace-match.gql‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
mutation matchWorkspaceContext($contextId: String!, $content: String!, $limit: SafeInt) {
2+
matchWorkspaceContext(contextId: $contextId, content: $content, limit: $limit) {
3+
docId
4+
chunk
5+
content
6+
distance
7+
}
8+
}

‎packages/frontend/graphql/src/graphql/index.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,22 @@ query listContextFiles($workspaceId: String!, $sessionId: String!, $contextId: S
212212
}`,
213213
};
214214

215+
export const matchContextMutation = {
216+
id: 'matchContextMutation' as const,
217+
operationName: 'matchContext',
218+
definitionName: 'matchContext',
219+
containsFile: false,
220+
query: `
221+
mutation matchContext($contextId: String!, $content: String!, $limit: SafeInt) {
222+
matchContext(contextId: $contextId, content: $content, limit: $limit) {
223+
fileId
224+
chunk
225+
content
226+
distance
227+
}
228+
}`,
229+
};
230+
215231
export const removeContextFileMutation = {
216232
id: 'removeContextFileMutation' as const,
217233
operationName: 'removeContextFile',
@@ -240,6 +256,22 @@ query listContext($workspaceId: String!, $sessionId: String!) {
240256
}`,
241257
};
242258

259+
export const matchWorkspaceContextMutation = {
260+
id: 'matchWorkspaceContextMutation' as const,
261+
operationName: 'matchWorkspaceContext',
262+
definitionName: 'matchWorkspaceContext',
263+
containsFile: false,
264+
query: `
265+
mutation matchWorkspaceContext($contextId: String!, $content: String!, $limit: SafeInt) {
266+
matchWorkspaceContext(contextId: $contextId, content: $content, limit: $limit) {
267+
docId
268+
chunk
269+
content
270+
distance
271+
}
272+
}`,
273+
};
274+
243275
export const getWorkspaceEmbeddingStatusQuery = {
244276
id: 'getWorkspaceEmbeddingStatusQuery' as const,
245277
operationName: 'getWorkspaceEmbeddingStatus',

0 commit comments

Comments
 (0)