Skip to content

Commit e821638

Browse files
darkskygitdarkskygit
committed
feat: add doc record for gql
1 parent ab52bf1 commit e821638

File tree

10 files changed

+198
-8
lines changed

10 files changed

+198
-8
lines changed

‎packages/backend/server/src/__tests__/copilot.spec.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1222,7 +1222,7 @@ test('should be able to manage context', async t => {
12221222
const session = await context.create(chatSession);
12231223

12241224
const fileId = await session.add(file, randomUUID());
1225-
const list = await session.list();
1225+
const list = await session.listFiles();
12261226
t.deepEqual(
12271227
list.map(f => f.chunk_size),
12281228
[3],
@@ -1234,6 +1234,11 @@ test('should be able to manage context', async t => {
12341234
'should list file id'
12351235
);
12361236

1237+
const docId = randomUUID();
1238+
await session.addDocRecord(randomUUID());
1239+
const docs = await session.listDocs();
1240+
t.deepEqual(docs, [docId], 'should list doc id');
1241+
12371242
const result = await session.match('test', 2);
12381243
t.is(result.length, 2, 'should match context');
12391244
t.is(result[0].fileId, fileId!, 'should match file id');

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

Lines changed: 72 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,15 @@ import {
4040
FileChunkSimilarity,
4141
} from './types';
4242

43+
@InputType()
44+
class AddContextDocInput {
45+
@Field(() => String)
46+
contextId!: string;
47+
48+
@Field(() => String)
49+
docId!: string;
50+
}
51+
4352
@InputType()
4453
class AddContextFileInput {
4554
@Field(() => String)
@@ -238,23 +247,80 @@ export class CopilotContextResolver {
238247
return controller.signal;
239248
}
240249

241-
@ResolveField(() => [CopilotContextFile], {
250+
@ResolveField(() => [String], {
242251
description: 'list files in context',
243252
})
244253
@CallMetric('ai', 'context_file_list')
245-
async files(
254+
async docs(
246255
@Parent() context: CopilotContextType,
247256
@Args('contextId', { nullable: true }) contextId?: string
248-
): Promise<CopilotContextFile[] | TooManyRequest> {
257+
): Promise<string[]> {
249258
const id = contextId || context.id;
250-
const lockFlag = `${COPILOT_LOCKER}:context:${id}`;
259+
const session = await this.context.get(id);
260+
return await session.listDocs();
261+
}
262+
263+
@Mutation(() => SafeIntResolver, {
264+
description: 'add a doc to context',
265+
})
266+
@CallMetric('ai', 'context_doc_add')
267+
async addContextDoc(
268+
@Args({ name: 'options', type: () => AddContextDocInput })
269+
options: AddContextDocInput
270+
) {
271+
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
251272
await using lock = await this.mutex.acquire(lockFlag);
252273
if (!lock) {
253274
return new TooManyRequest('Server is busy');
254275
}
255-
const session = await this.context.get(id);
276+
const session = await this.context.get(options.contextId);
277+
278+
try {
279+
return await session.addDocRecord(options.docId);
280+
} catch (e: any) {
281+
throw new CopilotFailedToModifyContext({
282+
contextId: options.contextId,
283+
message: e.message,
284+
});
285+
}
286+
}
287+
288+
@Mutation(() => Boolean, {
289+
description: 'remove a doc from context',
290+
})
291+
@CallMetric('ai', 'context_doc_remove')
292+
async removeContextDoc(
293+
@Args({ name: 'options', type: () => RemoveContextFileInput })
294+
options: RemoveContextFileInput
295+
) {
296+
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
297+
await using lock = await this.mutex.acquire(lockFlag);
298+
if (!lock) {
299+
return new TooManyRequest('Server is busy');
300+
}
301+
const session = await this.context.get(options.contextId);
302+
303+
try {
304+
return await session.removeDocRecord(options.fileId);
305+
} catch (e: any) {
306+
throw new CopilotFailedToModifyContext({
307+
contextId: options.contextId,
308+
message: e.message,
309+
});
310+
}
311+
}
256312

257-
return await session.list();
313+
@ResolveField(() => [CopilotContextFile], {
314+
description: 'list files in context',
315+
})
316+
@CallMetric('ai', 'context_file_list')
317+
async files(
318+
@Parent() context: CopilotContextType,
319+
@Args('contextId', { nullable: true }) contextId?: string
320+
): Promise<CopilotContextFile[]> {
321+
const id = contextId || context.id;
322+
const session = await this.context.get(id);
323+
return await session.listFiles();
258324
}
259325

260326
@Mutation(() => String, {

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,11 @@ export class ContextSession implements AsyncDisposable {
2828
return this.contextId;
2929
}
3030

31-
async list() {
31+
async listDocs() {
32+
return [...this.config.docs];
33+
}
34+
35+
async listFiles() {
3236
return this.config.files.map(f => ({ ...f }));
3337
}
3438

@@ -106,6 +110,24 @@ export class ContextSession implements AsyncDisposable {
106110
});
107111
}
108112

113+
async addDocRecord(docId: string) {
114+
if (!this.config.docs.includes(docId)) {
115+
this.config.docs.push(docId);
116+
await this.save();
117+
}
118+
return this.config.docs.length;
119+
}
120+
121+
async removeDocRecord(docId: string) {
122+
const index = this.config.docs.indexOf(docId);
123+
if (index >= 0) {
124+
this.config.docs.splice(index, 1);
125+
await this.save();
126+
return true;
127+
}
128+
return false;
129+
}
130+
109131
async addStream(
110132
readable: Readable,
111133
name: string,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ export const ContextConfigSchema = z.object({
2424
blobId: z.string(),
2525
})
2626
.array(),
27+
docs: z.string().array(),
2728
});
2829

2930
export type ContextConfig = z.infer<typeof ContextConfigSchema>;

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,11 @@
22
# THIS FILE WAS AUTOMATICALLY GENERATED (DO NOT MODIFY)
33
# ------------------------------------------------------
44

5+
input AddContextDocInput {
6+
contextId: String!
7+
docId: String!
8+
}
9+
510
input AddContextFileInput {
611
blobId: String!
712
contextId: String!
@@ -66,6 +71,9 @@ type Copilot {
6671
}
6772

6873
type CopilotContext {
74+
"""list files in context"""
75+
docs(contextId: String): [String!]!
76+
6977
"""list files in context"""
7078
files(contextId: String): [CopilotContextFile!]!
7179
id: ID!
@@ -548,6 +556,9 @@ type MissingOauthQueryParameterDataType {
548556
type Mutation {
549557
acceptInviteById(inviteId: String!, sendAcceptMail: Boolean, workspaceId: String!): Boolean!
550558

559+
"""add a doc to context"""
560+
addContextDoc(options: AddContextDocInput!): SafeInt!
561+
551562
"""add a file to context"""
552563
addContextFile(content: Upload!, options: AddContextFileInput!): String!
553564
addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int!
@@ -612,6 +623,9 @@ type Mutation {
612623
"""Remove user avatar"""
613624
removeAvatar: RemoveAvatar!
614625

626+
"""remove a doc from context"""
627+
removeContextDoc(options: RemoveContextFileInput!): Boolean!
628+
615629
"""remove a file from context"""
616630
removeContextFile(options: RemoveContextFileInput!): Boolean!
617631
removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int!

‎packages/frontend/graphql/src/graphql/copilot-context-doc-add.gql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
mutation addContextDoc($options: AddContextDocInput!) {
2+
addContextDoc(options: $options)
3+
}

‎packages/frontend/graphql/src/graphql/copilot-context-doc-remove.gql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
mutation removeContextDoc($options: RemoveContextFileInput!) {
2+
removeContextDoc(options: $options)
3+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ query listContextFiles(
66
currentUser {
77
copilot(workspaceId: $workspaceId) {
88
contexts(sessionId: $sessionId) {
9+
docs(contextId: $contextId)
910
files(contextId: $contextId) {
1011
id
1112
name

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,28 @@ mutation createCopilotContext($workspaceId: String!, $sessionId: String!) {
155155
}`,
156156
};
157157

158+
export const addContextDocMutation = {
159+
id: 'addContextDocMutation' as const,
160+
operationName: 'addContextDoc',
161+
definitionName: 'addContextDoc',
162+
containsFile: false,
163+
query: `
164+
mutation addContextDoc($options: AddContextDocInput!) {
165+
addContextDoc(options: $options)
166+
}`,
167+
};
168+
169+
export const removeContextDocMutation = {
170+
id: 'removeContextDocMutation' as const,
171+
operationName: 'removeContextDoc',
172+
definitionName: 'removeContextDoc',
173+
containsFile: false,
174+
query: `
175+
mutation removeContextDoc($options: RemoveContextFileInput!) {
176+
removeContextDoc(options: $options)
177+
}`,
178+
};
179+
158180
export const addContextFileMutation = {
159181
id: 'addContextFileMutation' as const,
160182
operationName: 'addContextFile',
@@ -176,6 +198,7 @@ query listContextFiles($workspaceId: String!, $sessionId: String!, $contextId: S
176198
currentUser {
177199
copilot(workspaceId: $workspaceId) {
178200
contexts(sessionId: $sessionId) {
201+
docs(contextId: $contextId)
179202
files(contextId: $contextId) {
180203
id
181204
name

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ export interface Scalars {
3737
Upload: { input: File; output: File };
3838
}
3939

40+
export interface AddContextDocInput {
41+
contextId: Scalars['String']['input'];
42+
docId: Scalars['String']['input'];
43+
}
44+
4045
export interface AddContextFileInput {
4146
blobId: Scalars['String']['input'];
4247
contextId: Scalars['String']['input'];
@@ -115,10 +120,16 @@ export interface CopilotHistoriesArgs {
115120
export interface CopilotContext {
116121
__typename?: 'CopilotContext';
117122
/** list files in context */
123+
docs: Array<Scalars['String']['output']>;
124+
/** list files in context */
118125
files: Array<CopilotContextFile>;
119126
id: Scalars['ID']['output'];
120127
}
121128

129+
export interface CopilotContextDocsArgs {
130+
contextId?: InputMaybe<Scalars['String']['input']>;
131+
}
132+
122133
export interface CopilotContextFilesArgs {
123134
contextId?: InputMaybe<Scalars['String']['input']>;
124135
}
@@ -636,6 +647,8 @@ export interface MissingOauthQueryParameterDataType {
636647
export interface Mutation {
637648
__typename?: 'Mutation';
638649
acceptInviteById: Scalars['Boolean']['output'];
650+
/** add a doc to context */
651+
addContextDoc: Scalars['SafeInt']['output'];
639652
/** add a file to context */
640653
addContextFile: Scalars['String']['output'];
641654
addWorkspaceFeature: Scalars['Int']['output'];
@@ -684,6 +697,8 @@ export interface Mutation {
684697
releaseDeletedBlobs: Scalars['Boolean']['output'];
685698
/** Remove user avatar */
686699
removeAvatar: RemoveAvatar;
700+
/** remove a doc from context */
701+
removeContextDoc: Scalars['Boolean']['output'];
687702
/** remove a file from context */
688703
removeContextFile: Scalars['Boolean']['output'];
689704
removeWorkspaceFeature: Scalars['Int']['output'];
@@ -729,6 +744,10 @@ export interface MutationAcceptInviteByIdArgs {
729744
workspaceId: Scalars['String']['input'];
730745
}
731746

747+
export interface MutationAddContextDocArgs {
748+
options: AddContextDocInput;
749+
}
750+
732751
export interface MutationAddContextFileArgs {
733752
content: Scalars['Upload']['input'];
734753
options: AddContextFileInput;
@@ -873,6 +892,10 @@ export interface MutationReleaseDeletedBlobsArgs {
873892
workspaceId: Scalars['String']['input'];
874893
}
875894

895+
export interface MutationRemoveContextDocArgs {
896+
options: RemoveContextFileInput;
897+
}
898+
876899
export interface MutationRemoveContextFileArgs {
877900
options: RemoveContextFileInput;
878901
}
@@ -1690,6 +1713,24 @@ export type CreateCopilotContextMutation = {
16901713
createCopilotContext: string;
16911714
};
16921715

1716+
export type AddContextDocMutationVariables = Exact<{
1717+
options: AddContextDocInput;
1718+
}>;
1719+
1720+
export type AddContextDocMutation = {
1721+
__typename?: 'Mutation';
1722+
addContextDoc: number;
1723+
};
1724+
1725+
export type RemoveContextDocMutationVariables = Exact<{
1726+
options: RemoveContextFileInput;
1727+
}>;
1728+
1729+
export type RemoveContextDocMutation = {
1730+
__typename?: 'Mutation';
1731+
removeContextDoc: boolean;
1732+
};
1733+
16931734
export type AddContextFileMutationVariables = Exact<{
16941735
content: Scalars['Upload']['input'];
16951736
options: AddContextFileInput;
@@ -1714,6 +1755,7 @@ export type ListContextFilesQuery = {
17141755
__typename?: 'Copilot';
17151756
contexts: Array<{
17161757
__typename?: 'CopilotContext';
1758+
docs: Array<string>;
17171759
files: Array<{
17181760
__typename?: 'CopilotContextFile';
17191761
id: string;
@@ -3311,6 +3353,16 @@ export type Mutations =
33113353
variables: CreateCopilotContextMutationVariables;
33123354
response: CreateCopilotContextMutation;
33133355
}
3356+
| {
3357+
name: 'addContextDocMutation';
3358+
variables: AddContextDocMutationVariables;
3359+
response: AddContextDocMutation;
3360+
}
3361+
| {
3362+
name: 'removeContextDocMutation';
3363+
variables: RemoveContextDocMutationVariables;
3364+
response: RemoveContextDocMutation;
3365+
}
33143366
| {
33153367
name: 'addContextFileMutation';
33163368
variables: AddContextFileMutationVariables;

0 commit comments

Comments
 (0)