feat: add pagination support for workspace config (#11859)

fix AI-78
This commit is contained in:
darkskygit
2025-04-23 11:25:41 +00:00
parent 5397fba897
commit ddb739fa13
13 changed files with 460 additions and 113 deletions

View File

@@ -0,0 +1,56 @@
# Snapshot report for `src/__tests__/models/copilot-workspace.spec.ts`
The actual snapshot is saved in `copilot-workspace.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## should manage copilot workspace ignored docs
> should add ignored doc
1
> should return added doc
[
{
docId: 'doc1',
},
]
> should return ignored docs in workspace
[
'doc1',
]
> should not add ignored doc again
[
{
docId: 'doc1',
},
]
> should add new ignored doc
2
> should add ignored doc
[
{
docId: 'new_doc',
},
{
docId: 'doc1',
},
]
> should remove ignored doc
[
{
docId: 'new_doc',
},
]

View File

@@ -6,6 +6,7 @@ import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace';
import { UserModel } from '../../models/user';
import { WorkspaceModel } from '../../models/workspace';
import { createTestingModule, type TestingModule } from '../utils';
import { cleanObject } from '../utils/copilot';
interface Context {
config: Config;
@@ -56,16 +57,16 @@ test('should manage copilot workspace ignored docs', async t => {
workspace.id,
[docId]
);
t.is(count, 1, 'should add ignored doc');
t.snapshot(count, 'should add ignored doc');
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.deepEqual(ret, [docId], 'should return added doc');
t.snapshot(cleanObject(ret), 'should return added doc');
const check = await t.context.copilotWorkspace.checkIgnoredDocs(
workspace.id,
[docId]
);
t.deepEqual(check, [docId], 'should return ignored docs in workspace');
t.snapshot(check, 'should return ignored docs in workspace');
}
{
@@ -76,7 +77,7 @@ test('should manage copilot workspace ignored docs', async t => {
t.is(count, 1, 'should not add ignored doc again');
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.deepEqual(ret, [docId], 'should not add ignored doc again');
t.snapshot(cleanObject(ret), 'should not add ignored doc again');
}
{
@@ -84,10 +85,10 @@ test('should manage copilot workspace ignored docs', async t => {
workspace.id,
['new_doc']
);
t.is(count, 2, 'should add new ignored doc');
t.snapshot(count, 'should add new ignored doc');
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.deepEqual(ret, [docId, 'new_doc'], 'should add ignored doc');
t.snapshot(cleanObject(ret), 'should add ignored doc');
}
{
@@ -98,7 +99,7 @@ test('should manage copilot workspace ignored docs', async t => {
);
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.deepEqual(ret, ['new_doc'], 'should remove ignored doc');
t.snapshot(cleanObject(ret), 'should remove ignored doc');
}
});

View File

@@ -4,6 +4,7 @@ import { Injectable } from '@nestjs/common';
import { Transactional } from '@nestjs-cls/transactional';
import { Prisma } from '@prisma/client';
import { PaginationInput } from '../base';
import { BaseModel } from './base';
import type {
CopilotWorkspaceFile,
@@ -22,7 +23,7 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
) {
const removed = new Set(remove);
const ignored = await this.listIgnoredDocs(workspaceId).then(
r => new Set(r.filter(id => !removed.has(id)))
r => new Set(r.map(r => r.docId).filter(id => !removed.has(id)))
);
const added = add.filter(id => !ignored.has(id));
@@ -49,22 +50,40 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
return added.length + ignored.size;
}
async listIgnoredDocs(workspaceId: string): Promise<string[]> {
async listIgnoredDocs(
workspaceId: string,
options?: {
includeRead?: boolean;
} & PaginationInput
): Promise<{ docId: string; createdAt: Date }[]> {
const row = await this.db.aiWorkspaceIgnoredDocs.findMany({
where: {
workspaceId,
},
select: {
docId: true,
createdAt: true,
},
orderBy: { createdAt: 'desc' },
skip: options?.offset,
take: options?.first,
});
return row;
}
async countIgnoredDocs(workspaceId: string): Promise<number> {
const count = await this.db.aiWorkspaceIgnoredDocs.count({
where: {
workspaceId,
},
});
return row.map(r => r.docId);
return count;
}
@Transactional()
async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
const ignored = await this.listIgnoredDocs(workspaceId).then(
r => new Set(r)
r => new Set(r.map(r => r.docId))
);
return docIds.filter(id => ignored.has(id));
@@ -133,16 +152,31 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
}
async listWorkspaceFiles(
workspaceId: string
workspaceId: string,
options?: {
includeRead?: boolean;
} & PaginationInput
): Promise<CopilotWorkspaceFile[]> {
const files = await this.db.aiWorkspaceFiles.findMany({
where: {
workspaceId,
},
orderBy: { createdAt: 'desc' },
skip: options?.offset,
take: options?.first,
});
return files;
}
async countWorkspaceFiles(workspaceId: string): Promise<number> {
const count = await this.db.aiWorkspaceFiles.count({
where: {
workspaceId,
},
});
return count;
}
async matchWorkspaceFileEmbedding(
workspaceId: string,
embedding: number[],

View File

@@ -9,7 +9,6 @@ import {
Resolver,
} from '@nestjs/graphql';
import type { Request } from 'express';
import { SafeIntResolver } from 'graphql-scalars';
import GraphQLUpload, {
type FileUpload,
} from 'graphql-upload/GraphQLUpload.mjs';
@@ -19,16 +18,22 @@ import {
CopilotEmbeddingUnavailable,
CopilotFailedToAddWorkspaceFileEmbedding,
Mutex,
paginate,
PaginationInput,
TooManyRequest,
UserFriendlyError,
} from '../../../base';
import { CurrentUser } from '../../../core/auth';
import { AccessController } from '../../../core/permission';
import { WorkspaceType } from '../../../core/workspaces';
import { CopilotWorkspaceFile, Models } from '../../../models';
import { COPILOT_LOCKER } from '../resolver';
import { MAX_EMBEDDABLE_SIZE } from '../types';
import { CopilotWorkspaceService } from './service';
import {
CopilotWorkspaceFileType,
PaginatedCopilotWorkspaceFileType,
PaginatedIgnoredDocsType,
} from './types';
@ObjectType('CopilotWorkspaceConfig')
export class CopilotWorkspaceConfigType {
@@ -36,27 +41,6 @@ export class CopilotWorkspaceConfigType {
workspaceId!: string;
}
@ObjectType('CopilotWorkspaceFile')
export class CopilotWorkspaceFileType implements CopilotWorkspaceFile {
@Field(() => String)
workspaceId!: string;
@Field(() => String)
fileId!: string;
@Field(() => String)
fileName!: string;
@Field(() => String)
mimeType!: string;
@Field(() => SafeIntResolver)
size!: number;
@Field(() => Date)
createdAt!: Date;
}
/**
* Workspace embedding config resolver
* Public apis rate limit: 10 req/m
@@ -86,18 +70,24 @@ export class CopilotWorkspaceEmbeddingResolver {
export class CopilotWorkspaceEmbeddingConfigResolver {
constructor(
private readonly ac: AccessController,
private readonly models: Models,
private readonly mutex: Mutex,
private readonly copilotWorkspace: CopilotWorkspaceService
) {}
@ResolveField(() => [String], {
@ResolveField(() => PaginatedIgnoredDocsType, {
complexity: 2,
})
async ignoredDocs(
@Parent() config: CopilotWorkspaceConfigType
): Promise<string[]> {
return this.models.copilotWorkspace.listIgnoredDocs(config.workspaceId);
@Parent() config: CopilotWorkspaceConfigType,
@Args('pagination', PaginationInput.decode) pagination: PaginationInput
): Promise<PaginatedIgnoredDocsType> {
const [ignoredDocs, totalCount] =
await this.copilotWorkspace.listIgnoredDocs(
config.workspaceId,
pagination
);
return paginate(ignoredDocs, 'createdAt', pagination, totalCount);
}
@Mutation(() => Number, {
@@ -118,20 +108,26 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
.user(user.id)
.workspace(workspaceId)
.assert('Workspace.Settings.Update');
return await this.models.copilotWorkspace.updateIgnoredDocs(
return await this.copilotWorkspace.updateIgnoredDocs(
workspaceId,
add,
remove
);
}
@ResolveField(() => [CopilotWorkspaceFileType], {
@ResolveField(() => PaginatedCopilotWorkspaceFileType, {
complexity: 2,
})
async files(
@Parent() config: CopilotWorkspaceConfigType
): Promise<CopilotWorkspaceFileType[]> {
return this.models.copilotWorkspace.listWorkspaceFiles(config.workspaceId);
@Parent() config: CopilotWorkspaceConfigType,
@Args('pagination', PaginationInput.decode) pagination: PaginationInput
): Promise<PaginatedCopilotWorkspaceFileType> {
const [files, totalCount] = await this.copilotWorkspace.listWorkspaceFiles(
config.workspaceId,
pagination
);
return paginate(files, 'createdAt', pagination, totalCount);
}
@Mutation(() => CopilotWorkspaceFileType, {
@@ -210,9 +206,6 @@ export class CopilotWorkspaceEmbeddingConfigResolver {
.workspace(workspaceId)
.assert('Workspace.Settings.Update');
return await this.models.copilotWorkspace.removeWorkspaceFile(
workspaceId,
fileId
);
return await this.copilotWorkspace.removeWorkspaceFile(workspaceId, fileId);
}
}

View File

@@ -2,31 +2,11 @@ import { createHash } from 'node:crypto';
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import { FileUpload, JobQueue } from '../../../base';
import { FileUpload, JobQueue, PaginationInput } from '../../../base';
import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
declare global {
interface Events {
'workspace.file.embedding.finished': {
jobId: string;
};
'workspace.file.embedding.failed': {
jobId: string;
};
}
interface Jobs {
'copilot.workspace.embedding.files': {
userId: string;
workspaceId: string;
blobId: string;
fileId: string;
fileName: string;
};
}
}
@Injectable()
export class CopilotWorkspaceService implements OnApplicationBootstrap {
private supportEmbedding = false;
@@ -49,6 +29,30 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
return this.supportEmbedding;
}
async updateIgnoredDocs(
workspaceId: string,
add?: string[],
remove?: string[]
) {
return await this.models.copilotWorkspace.updateIgnoredDocs(
workspaceId,
add,
remove
);
}
async listIgnoredDocs(
workspaceId: string,
pagination?: {
includeRead?: boolean;
} & PaginationInput
) {
return await Promise.all([
this.models.copilotWorkspace.listIgnoredDocs(workspaceId, pagination),
this.models.copilotWorkspace.countIgnoredDocs(workspaceId),
]);
}
async addWorkspaceFile(
userId: string,
workspaceId: string,
@@ -70,6 +74,18 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
return await this.models.copilotWorkspace.getFile(workspaceId, fileId);
}
async listWorkspaceFiles(
workspaceId: string,
pagination?: {
includeRead?: boolean;
} & PaginationInput
) {
return await Promise.all([
this.models.copilotWorkspace.listWorkspaceFiles(workspaceId, pagination),
this.models.copilotWorkspace.countIgnoredDocs(workspaceId),
]);
}
async addWorkspaceFileEmbeddingQueue(
file: Jobs['copilot.workspace.embedding.files']
) {
@@ -84,4 +100,11 @@ export class CopilotWorkspaceService implements OnApplicationBootstrap {
fileName,
});
}
async removeWorkspaceFile(workspaceId: string, fileId: string) {
return await this.models.copilotWorkspace.removeWorkspaceFile(
workspaceId,
fileId
);
}
}

View File

@@ -0,0 +1,65 @@
import { Field, ObjectType } from '@nestjs/graphql';
import { SafeIntResolver } from 'graphql-scalars';
import { Paginated } from '../../../base';
import { CopilotWorkspaceFile } from '../../../models';
declare global {
interface Events {
'workspace.file.embedding.finished': {
jobId: string;
};
'workspace.file.embedding.failed': {
jobId: string;
};
}
interface Jobs {
'copilot.workspace.embedding.files': {
userId: string;
workspaceId: string;
blobId: string;
fileId: string;
fileName: string;
};
}
}
@ObjectType('CopilotWorkspaceIgnoredDoc')
export class CopilotWorkspaceIgnoredDocType {
@Field(() => String)
docId!: string;
@Field(() => Date)
createdAt!: Date;
}
@ObjectType()
export class PaginatedIgnoredDocsType extends Paginated(
CopilotWorkspaceIgnoredDocType
) {}
@ObjectType('CopilotWorkspaceFile')
export class CopilotWorkspaceFileType implements CopilotWorkspaceFile {
@Field(() => String)
workspaceId!: string;
@Field(() => String)
fileId!: string;
@Field(() => String)
fileName!: string;
@Field(() => String)
mimeType!: string;
@Field(() => SafeIntResolver)
size!: number;
@Field(() => Date)
createdAt!: Date;
}
@ObjectType()
export class PaginatedCopilotWorkspaceFileType extends Paginated(
CopilotWorkspaceFileType
) {}

View File

@@ -273,8 +273,8 @@ type CopilotSessionType {
}
type CopilotWorkspaceConfig {
files: [CopilotWorkspaceFile!]!
ignoredDocs: [String!]!
files(pagination: PaginationInput!): PaginatedCopilotWorkspaceFileType!
ignoredDocs(pagination: PaginationInput!): PaginatedIgnoredDocsType!
workspaceId: String!
}
@@ -287,6 +287,21 @@ type CopilotWorkspaceFile {
workspaceId: String!
}
type CopilotWorkspaceFileTypeEdge {
cursor: String!
node: CopilotWorkspaceFile!
}
type CopilotWorkspaceIgnoredDoc {
createdAt: DateTime!
docId: String!
}
type CopilotWorkspaceIgnoredDocTypeEdge {
cursor: String!
node: CopilotWorkspaceIgnoredDoc!
}
input CreateChatMessageInput {
attachments: [String!]
blobs: [Upload!]
@@ -1149,12 +1164,24 @@ type PageInfo {
startCursor: String
}
type PaginatedCopilotWorkspaceFileType {
edges: [CopilotWorkspaceFileTypeEdge!]!
pageInfo: PageInfo!
totalCount: Int!
}
type PaginatedGrantedDocUserType {
edges: [GrantedDocUserTypeEdge!]!
pageInfo: PageInfo!
totalCount: Int!
}
type PaginatedIgnoredDocsType {
edges: [CopilotWorkspaceIgnoredDocTypeEdge!]!
pageInfo: PageInfo!
totalCount: Int!
}
type PaginatedNotificationObjectType {
edges: [NotificationObjectTypeEdge!]!
pageInfo: PageInfo!

View File

@@ -0,0 +1,22 @@
query getWorkspaceEmbeddingFiles($workspaceId: String!, $pagination: PaginationInput!) {
workspace(id: $workspaceId) {
embedding {
files(pagination: $pagination) {
totalCount
pageInfo {
endCursor
hasNextPage
}
edges {
node {
fileId
fileName
mimeType
size
createdAt
}
}
}
}
}
}

View File

@@ -0,0 +1,19 @@
query getWorkspaceEmbeddingIgnoredDocs($workspaceId: String!, $pagination: PaginationInput!) {
workspace(id: $workspaceId) {
embedding {
ignoredDocs(pagination: $pagination) {
totalCount
pageInfo {
endCursor
hasNextPage
}
edges {
node {
docId
createdAt
}
}
}
}
}
}

View File

@@ -1,14 +0,0 @@
query getWorkspaceEmbeddingConfig($workspaceId: String!) {
workspace(id: $workspaceId) {
embedding {
files {
fileId
fileName
mimeType
size
createdAt
}
ignoredDocs
}
}
}

View File

@@ -764,20 +764,52 @@ export const removeWorkspaceEmbeddingFilesMutation = {
}`,
};
export const getWorkspaceEmbeddingConfigQuery = {
id: 'getWorkspaceEmbeddingConfigQuery' as const,
op: 'getWorkspaceEmbeddingConfig',
query: `query getWorkspaceEmbeddingConfig($workspaceId: String!) {
export const getWorkspaceEmbeddingFilesQuery = {
id: 'getWorkspaceEmbeddingFilesQuery' as const,
op: 'getWorkspaceEmbeddingFiles',
query: `query getWorkspaceEmbeddingFiles($workspaceId: String!, $pagination: PaginationInput!) {
workspace(id: $workspaceId) {
embedding {
files {
fileId
fileName
mimeType
size
createdAt
files(pagination: $pagination) {
totalCount
pageInfo {
endCursor
hasNextPage
}
edges {
node {
fileId
fileName
mimeType
size
createdAt
}
}
}
}
}
}`,
};
export const getWorkspaceEmbeddingIgnoredDocsQuery = {
id: 'getWorkspaceEmbeddingIgnoredDocsQuery' as const,
op: 'getWorkspaceEmbeddingIgnoredDocs',
query: `query getWorkspaceEmbeddingIgnoredDocs($workspaceId: String!, $pagination: PaginationInput!) {
workspace(id: $workspaceId) {
embedding {
ignoredDocs(pagination: $pagination) {
totalCount
pageInfo {
endCursor
hasNextPage
}
edges {
node {
docId
createdAt
}
}
}
ignoredDocs
}
}
}`,

View File

@@ -366,11 +366,19 @@ export interface CopilotSessionType {
export interface CopilotWorkspaceConfig {
__typename?: 'CopilotWorkspaceConfig';
files: Array<CopilotWorkspaceFile>;
ignoredDocs: Array<Scalars['String']['output']>;
files: PaginatedCopilotWorkspaceFileType;
ignoredDocs: PaginatedIgnoredDocsType;
workspaceId: Scalars['String']['output'];
}
export interface CopilotWorkspaceConfigFilesArgs {
pagination: PaginationInput;
}
export interface CopilotWorkspaceConfigIgnoredDocsArgs {
pagination: PaginationInput;
}
export interface CopilotWorkspaceFile {
__typename?: 'CopilotWorkspaceFile';
createdAt: Scalars['DateTime']['output'];
@@ -381,6 +389,24 @@ export interface CopilotWorkspaceFile {
workspaceId: Scalars['String']['output'];
}
export interface CopilotWorkspaceFileTypeEdge {
__typename?: 'CopilotWorkspaceFileTypeEdge';
cursor: Scalars['String']['output'];
node: CopilotWorkspaceFile;
}
export interface CopilotWorkspaceIgnoredDoc {
__typename?: 'CopilotWorkspaceIgnoredDoc';
createdAt: Scalars['DateTime']['output'];
docId: Scalars['String']['output'];
}
export interface CopilotWorkspaceIgnoredDocTypeEdge {
__typename?: 'CopilotWorkspaceIgnoredDocTypeEdge';
cursor: Scalars['String']['output'];
node: CopilotWorkspaceIgnoredDoc;
}
export interface CreateChatMessageInput {
attachments?: InputMaybe<Array<Scalars['String']['input']>>;
blobs?: InputMaybe<Array<Scalars['Upload']['input']>>;
@@ -1622,6 +1648,13 @@ export interface PageInfo {
startCursor: Maybe<Scalars['String']['output']>;
}
export interface PaginatedCopilotWorkspaceFileType {
__typename?: 'PaginatedCopilotWorkspaceFileType';
edges: Array<CopilotWorkspaceFileTypeEdge>;
pageInfo: PageInfo;
totalCount: Scalars['Int']['output'];
}
export interface PaginatedGrantedDocUserType {
__typename?: 'PaginatedGrantedDocUserType';
edges: Array<GrantedDocUserTypeEdge>;
@@ -1629,6 +1662,13 @@ export interface PaginatedGrantedDocUserType {
totalCount: Scalars['Int']['output'];
}
export interface PaginatedIgnoredDocsType {
__typename?: 'PaginatedIgnoredDocsType';
edges: Array<CopilotWorkspaceIgnoredDocTypeEdge>;
pageInfo: PageInfo;
totalCount: Scalars['Int']['output'];
}
export interface PaginatedNotificationObjectType {
__typename?: 'PaginatedNotificationObjectType';
edges: Array<NotificationObjectTypeEdge>;
@@ -3243,25 +3283,69 @@ export type RemoveWorkspaceEmbeddingFilesMutation = {
removeWorkspaceEmbeddingFiles: boolean;
};
export type GetWorkspaceEmbeddingConfigQueryVariables = Exact<{
export type GetWorkspaceEmbeddingFilesQueryVariables = Exact<{
workspaceId: Scalars['String']['input'];
pagination: PaginationInput;
}>;
export type GetWorkspaceEmbeddingConfigQuery = {
export type GetWorkspaceEmbeddingFilesQuery = {
__typename?: 'Query';
workspace: {
__typename?: 'WorkspaceType';
embedding: {
__typename?: 'CopilotWorkspaceConfig';
ignoredDocs: Array<string>;
files: Array<{
__typename?: 'CopilotWorkspaceFile';
fileId: string;
fileName: string;
mimeType: string;
size: number;
createdAt: string;
}>;
files: {
__typename?: 'PaginatedCopilotWorkspaceFileType';
totalCount: number;
pageInfo: {
__typename?: 'PageInfo';
endCursor: string | null;
hasNextPage: boolean;
};
edges: Array<{
__typename?: 'CopilotWorkspaceFileTypeEdge';
node: {
__typename?: 'CopilotWorkspaceFile';
fileId: string;
fileName: string;
mimeType: string;
size: number;
createdAt: string;
};
}>;
};
};
};
};
export type GetWorkspaceEmbeddingIgnoredDocsQueryVariables = Exact<{
workspaceId: Scalars['String']['input'];
pagination: PaginationInput;
}>;
export type GetWorkspaceEmbeddingIgnoredDocsQuery = {
__typename?: 'Query';
workspace: {
__typename?: 'WorkspaceType';
embedding: {
__typename?: 'CopilotWorkspaceConfig';
ignoredDocs: {
__typename?: 'PaginatedIgnoredDocsType';
totalCount: number;
pageInfo: {
__typename?: 'PageInfo';
endCursor: string | null;
hasNextPage: boolean;
};
edges: Array<{
__typename?: 'CopilotWorkspaceIgnoredDocTypeEdge';
node: {
__typename?: 'CopilotWorkspaceIgnoredDoc';
docId: string;
createdAt: string;
};
}>;
};
};
};
};
@@ -4540,9 +4624,14 @@ export type Queries =
response: GetCopilotSessionsQuery;
}
| {
name: 'getWorkspaceEmbeddingConfigQuery';
variables: GetWorkspaceEmbeddingConfigQueryVariables;
response: GetWorkspaceEmbeddingConfigQuery;
name: 'getWorkspaceEmbeddingFilesQuery';
variables: GetWorkspaceEmbeddingFilesQueryVariables;
response: GetWorkspaceEmbeddingFilesQuery;
}
| {
name: 'getWorkspaceEmbeddingIgnoredDocsQuery';
variables: GetWorkspaceEmbeddingIgnoredDocsQueryVariables;
response: GetWorkspaceEmbeddingIgnoredDocsQuery;
}
| {
name: 'getDocRolePermissionsQuery';