Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: record createdAt for docs and files
Browse files Browse the repository at this point in the history
darkskygit committed Jan 27, 2025

Verified

This commit was signed with the committer’s verified signature.
makenowjust Hiroya Fujinami
1 parent e9a566c commit 460e3e6
Showing 12 changed files with 154 additions and 35 deletions.
6 changes: 5 additions & 1 deletion packages/backend/server/src/__tests__/copilot.e2e.ts
Original file line number Diff line number Diff line change
@@ -763,7 +763,11 @@ test('should be able to manage context', async t => {
await t.notThrowsAsync(context, 'should create context with chat session');

const list = await listContext(app, token, workspaceId, sessionId);
t.deepEqual(list, [{ id: await context }], 'should list context');
t.deepEqual(
list.map(f => ({ id: f.id })),
[{ id: await context }],
'should list context'
);
}

const fs = await import('node:fs');
8 changes: 4 additions & 4 deletions packages/backend/server/src/__tests__/copilot.spec.ts
Original file line number Diff line number Diff line change
@@ -1294,8 +1294,8 @@ test('should be able to manage context', async t => {
{
const session = await context.create(chatSession);

const fileId = await session.add(file, randomUUID());
const list = await session.listFiles();
const fileId = await session.addFile(file, randomUUID());
const list = session.listFiles();
t.deepEqual(
list.map(f => f.chunk_size),
[3],
@@ -1309,10 +1309,10 @@ test('should be able to manage context', async t => {

const docId = randomUUID();
await session.addDocRecord(randomUUID());
const docs = await session.listDocs();
const docs = session.listDocs();
t.deepEqual(docs, [docId], 'should list doc id');

const result = await session.match('test', 2);
const result = await session.matchFileChunks('test', 2);
t.is(result.length, 2, 'should match context');
t.is(result[0].fileId, fileId!, 'should match file id');
}
10 changes: 9 additions & 1 deletion packages/backend/server/src/__tests__/utils/copilot.ts
Original file line number Diff line number Diff line change
@@ -306,7 +306,12 @@ export async function listContext(
userToken: string,
workspaceId: string,
sessionId: string
): Promise<{ id: string }[]> {
): Promise<
{
id: string;
createdAt: number;
}[]
> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
@@ -318,6 +323,7 @@ export async function listContext(
copilot(workspaceId: "${workspaceId}") {
contexts(sessionId: "${sessionId}") {
id
createdAt
}
}
}
@@ -407,6 +413,7 @@ export async function listContextFiles(
blobId: string;
chunk_size: number;
status: string;
createdAt: number;
}[]
| undefined
> {
@@ -426,6 +433,7 @@ export async function listContextFiles(
blobId
chunk_size
status
createdAt
}
}
}
38 changes: 31 additions & 7 deletions packages/backend/server/src/plugins/copilot/context/resolver.ts
Original file line number Diff line number Diff line change
@@ -35,6 +35,7 @@ import { COPILOT_LOCKER, CopilotType } from '../resolver';
import { ChatSessionService } from '../session';
import { CopilotContextService } from './service';
import {
ContextDoc,
type ContextFile,
ContextFileStatus,
DocChunkSimilarity,
@@ -75,10 +76,22 @@ class RemoveContextFileInput {
export class CopilotContextType {
@Field(() => ID)
id!: string;

@Field(() => SafeIntResolver)
createdAt!: number;
}

registerEnumType(ContextFileStatus, { name: 'ContextFileStatus' });

@ObjectType()
class CopilotContextDoc implements ContextDoc {
@Field(() => ID)
id!: string;

@Field(() => SafeIntResolver)
createdAt!: number;
}

@ObjectType()
class CopilotContextFile implements ContextFile {
@Field(() => ID)
@@ -95,6 +108,9 @@ class CopilotContextFile implements ContextFile {

@Field(() => String)
blobId!: string;

@Field(() => SafeIntResolver)
createdAt!: number;
}

@ObjectType()
@@ -252,17 +268,17 @@ export class CopilotContextResolver {
return controller.signal;
}

@ResolveField(() => [String], {
@ResolveField(() => [CopilotContextDoc], {
description: 'list files in context',
})
@CallMetric('ai', 'context_file_list')
async docs(
@Parent() context: CopilotContextType,
@Args('contextId', { nullable: true }) contextId?: string
): Promise<string[]> {
): Promise<ContextDoc[]> {
const id = contextId || context.id;
const session = await this.context.get(id);
return await session.listDocs();
return session.listDocs();
}

@Mutation(() => SafeIntResolver, {
@@ -325,7 +341,7 @@ export class CopilotContextResolver {
): Promise<CopilotContextFile[]> {
const id = contextId || context.id;
const session = await this.context.get(id);
return await session.listFiles();
return session.listFiles();
}

@Mutation(() => String, {
@@ -378,7 +394,7 @@ export class CopilotContextResolver {
const session = await this.context.get(options.contextId);

try {
return await session.remove(options.fileId);
return await session.removeFile(options.fileId);
} catch (e: any) {
throw new CopilotFailedToModifyContext({
contextId: options.contextId,
@@ -406,7 +422,11 @@ export class CopilotContextResolver {
const session = await this.context.get(contextId);

try {
return await session.match(content, limit, this.getSignal(ctx.req));
return await session.matchFileChunks(
content,
limit,
this.getSignal(ctx.req)
);
} catch (e: any) {
throw new CopilotFailedToMatchContext({
contextId,
@@ -433,7 +453,11 @@ export class CopilotContextResolver {
await this.permissions.checkCloudWorkspace(session.workspaceId, user.id);

try {
return await session.match(content, limit, this.getSignal(ctx.req));
return await session.matchFileChunks(
content,
limit,
this.getSignal(ctx.req)
);
} catch (e: any) {
throw new CopilotFailedToMatchContext({
contextId,
Original file line number Diff line number Diff line change
@@ -84,11 +84,14 @@ export class CopilotContextService {
throw new CopilotInvalidContext({ contextId: id });
}

async list(sessionId: string): Promise<{ id: string }[]> {
async list(sessionId: string): Promise<{ id: string; createdAt: number }[]> {
const contexts = await this.db.aiContext.findMany({
where: { sessionId },
select: { id: true },
select: { id: true, createdAt: true },
});
return contexts;
return contexts.map(c => ({
id: c.id,
createdAt: c.createdAt.getTime(),
}));
}
}
27 changes: 15 additions & 12 deletions packages/backend/server/src/plugins/copilot/context/session.ts
Original file line number Diff line number Diff line change
@@ -8,7 +8,9 @@ import { nanoid } from 'nanoid';
import { BlobQuotaExceeded, PrismaTransaction } from '../../../base';
import { OneMB } from '../../../core/quota/constant';
import {
ChunkSimilarity,
ContextConfig,
ContextDoc,
ContextFile,
ContextFileStatus,
DocChunkSimilarity,
@@ -33,11 +35,11 @@ export class ContextSession implements AsyncDisposable {
return this.config.workspaceId;
}

async listDocs() {
listDocs(): ContextDoc[] {
return [...this.config.docs];
}

async listFiles() {
listFiles() {
return this.config.files.map(f => ({ ...f }));
}

@@ -65,6 +67,7 @@ export class ContextSession implements AsyncDisposable {
blobId,
chunk_size: embeddings.length,
name,
createdAt: Date.now(),
}));

const values = this.processEmbeddings(fileId, embeddings);
@@ -116,15 +119,15 @@ export class ContextSession implements AsyncDisposable {
}

async addDocRecord(docId: string) {
if (!this.config.docs.includes(docId)) {
this.config.docs.push(docId);
if (!this.config.docs.some(f => f.id === docId)) {
this.config.docs.push({ id: docId, createdAt: Date.now() });
await this.save();
}
return this.config.docs.length;
return this.config.docs;
}

async removeDocRecord(docId: string) {
const index = this.config.docs.indexOf(docId);
const index = this.config.docs.findIndex(f => f.id === docId);
if (index >= 0) {
this.config.docs.splice(index, 1);
await this.save();
@@ -142,10 +145,10 @@ export class ContextSession implements AsyncDisposable {
if (signal?.aborted) return;
const buffer = await this.readStream(readable, 50 * OneMB);
const file = new File([buffer], name);
return await this.add(file, blobId, signal);
return await this.addFile(file, blobId, signal);
}

async add(
async addFile(
file: File,
blobId: string,
signal?: AbortSignal
@@ -157,7 +160,7 @@ export class ContextSession implements AsyncDisposable {
return undefined;
}

async remove(fileId: string) {
async removeFile(fileId: string) {
return await this.db.$transaction(async tx => {
const ret = await tx.aiContextEmbedding.deleteMany({
where: { contextId: this.contextId, fileId },
@@ -168,7 +171,7 @@ export class ContextSession implements AsyncDisposable {
});
}

async match(
async matchFileChunks(
content: string,
topK: number = 5,
signal?: AbortSignal
@@ -185,11 +188,11 @@ export class ContextSession implements AsyncDisposable {
`;
}

async matchWorkspace(
async matchWorkspaceChunks(
content: string,
topK: number = 5,
signal?: AbortSignal
) {
): Promise<ChunkSimilarity[]> {
const embedding = await this.client
.getEmbeddings([content], signal)
.then(r => r?.[0]?.embedding);
9 changes: 8 additions & 1 deletion packages/backend/server/src/plugins/copilot/context/types.ts
Original file line number Diff line number Diff line change
@@ -32,12 +32,19 @@ export const ContextConfigSchema = z.object({
ContextFileStatus.failed,
]),
blobId: z.string(),
createdAt: z.number(),
})
.array(),
docs: z
.object({
id: z.string(),
createdAt: z.number(),
})
.array(),
docs: z.string().array(),
});

export type ContextConfig = z.infer<typeof ContextConfigSchema>;
export type ContextDoc = z.infer<typeof ContextConfigSchema>['docs'][number];
export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];

export type ChunkSimilarity = {
10 changes: 9 additions & 1 deletion packages/backend/server/src/schema.gql
Original file line number Diff line number Diff line change
@@ -78,17 +78,25 @@ type Copilot {
}

type CopilotContext {
createdAt: SafeInt!

"""list files in context"""
docs(contextId: String): [String!]!
docs(contextId: String): [CopilotContextDoc!]!

"""list files in context"""
files(contextId: String): [CopilotContextFile!]!
id: ID!
}

type CopilotContextDoc {
createdAt: SafeInt!
id: ID!
}

type CopilotContextFile {
blobId: String!
chunk_size: SafeInt!
createdAt: SafeInt!
id: ID!
name: String!
status: ContextFileStatus!
Original file line number Diff line number Diff line change
@@ -6,13 +6,17 @@ query listContextFiles(
currentUser {
copilot(workspaceId: $workspaceId) {
contexts(sessionId: $sessionId) {
docs(contextId: $contextId)
docs(contextId: $contextId) {
id
createdAt
}
files(contextId: $contextId) {
id
name
blobId
chunk_size
status
createdAt
}
}
}
Original file line number Diff line number Diff line change
@@ -3,6 +3,7 @@ query listContext($workspaceId: String!, $sessionId: String!) {
copilot(workspaceId: $workspaceId) {
contexts(sessionId: $sessionId) {
id
createdAt
}
}
}
7 changes: 6 additions & 1 deletion packages/frontend/graphql/src/graphql/index.ts
Original file line number Diff line number Diff line change
@@ -198,13 +198,17 @@ query listContextFiles($workspaceId: String!, $sessionId: String!, $contextId: S
currentUser {
copilot(workspaceId: $workspaceId) {
contexts(sessionId: $sessionId) {
docs(contextId: $contextId)
docs(contextId: $contextId) {
id
createdAt
}
files(contextId: $contextId) {
id
name
blobId
chunk_size
status
createdAt
}
}
}
@@ -250,6 +254,7 @@ query listContext($workspaceId: String!, $sessionId: String!) {
copilot(workspaceId: $workspaceId) {
contexts(sessionId: $sessionId) {
id
createdAt
}
}
}
58 changes: 55 additions & 3 deletions packages/frontend/graphql/src/schema.ts
Original file line number Diff line number Diff line change
@@ -127,8 +127,9 @@ export interface CopilotHistoriesArgs {

export interface CopilotContext {
__typename?: 'CopilotContext';
createdAt: Scalars['SafeInt']['output'];
/** list files in context */
docs: Array<Scalars['String']['output']>;
docs: Array<CopilotContextDoc>;
/** list files in context */
files: Array<CopilotContextFile>;
id: Scalars['ID']['output'];
@@ -142,10 +143,17 @@ export interface CopilotContextFilesArgs {
contextId?: InputMaybe<Scalars['String']['input']>;
}

export interface CopilotContextDoc {
__typename?: 'CopilotContextDoc';
createdAt: Scalars['SafeInt']['output'];
id: Scalars['ID']['output'];
}

export interface CopilotContextFile {
__typename?: 'CopilotContextFile';
blobId: Scalars['String']['output'];
chunk_size: Scalars['SafeInt']['output'];
createdAt: Scalars['SafeInt']['output'];
id: Scalars['ID']['output'];
name: Scalars['String']['output'];
status: ContextFileStatus;
@@ -368,6 +376,7 @@ export type ErrorDataUnion =
| MemberNotFoundInSpaceDataType
| MissingOauthQueryParameterDataType
| NotInSpaceDataType
| QueryTooLongDataType
| RuntimeConfigNotFoundDataType
| SameSubscriptionRecurringDataType
| SpaceAccessDeniedDataType
@@ -445,6 +454,7 @@ export enum ErrorNames {
OAUTH_STATE_EXPIRED = 'OAUTH_STATE_EXPIRED',
PAGE_IS_NOT_PUBLIC = 'PAGE_IS_NOT_PUBLIC',
PASSWORD_REQUIRED = 'PASSWORD_REQUIRED',
QUERY_TOO_LONG = 'QUERY_TOO_LONG',
RUNTIME_CONFIG_NOT_FOUND = 'RUNTIME_CONFIG_NOT_FOUND',
SAME_EMAIL_PROVIDED = 'SAME_EMAIL_PROVIDED',
SAME_SUBSCRIPTION_RECURRING = 'SAME_SUBSCRIPTION_RECURRING',
@@ -622,6 +632,15 @@ export interface InvoiceType {
updatedAt: Scalars['DateTime']['output'];
}

export interface License {
__typename?: 'License';
expiredAt: Maybe<Scalars['DateTime']['output']>;
installedAt: Scalars['DateTime']['output'];
quantity: Scalars['Int']['output'];
recurring: SubscriptionRecurring;
validatedAt: Scalars['DateTime']['output'];
}

export interface LimitedUserType {
__typename?: 'LimitedUserType';
/** User email */
@@ -663,6 +682,7 @@ export interface MissingOauthQueryParameterDataType {
export interface Mutation {
__typename?: 'Mutation';
acceptInviteById: Scalars['Boolean']['output'];
activateLicense: License;
/** add a doc to context */
addContextDoc: Scalars['SafeInt']['output'];
/** add a file to context */
@@ -689,10 +709,12 @@ export interface Mutation {
/** Create a stripe customer portal to manage payment methods */
createCustomerPortal: Scalars['String']['output'];
createInviteLink: InviteLink;
createSelfhostWorkspaceCustomerPortal: Scalars['String']['output'];
/** Create a new user */
createUser: UserType;
/** Create a new workspace */
createWorkspace: WorkspaceType;
deactivateLicense: Scalars['Boolean']['output'];
deleteAccount: DeleteAccount;
deleteBlob: Scalars['Boolean']['output'];
/** Delete a user account */
@@ -763,6 +785,11 @@ export interface MutationAcceptInviteByIdArgs {
workspaceId: Scalars['String']['input'];
}

export interface MutationActivateLicenseArgs {
license: Scalars['String']['input'];
workspaceId: Scalars['String']['input'];
}

export interface MutationAddContextDocArgs {
options: AddContextDocInput;
}
@@ -834,6 +861,10 @@ export interface MutationCreateInviteLinkArgs {
workspaceId: Scalars['String']['input'];
}

export interface MutationCreateSelfhostWorkspaceCustomerPortalArgs {
workspaceId: Scalars['String']['input'];
}

export interface MutationCreateUserArgs {
input: CreateUserInput;
}
@@ -842,6 +873,10 @@ export interface MutationCreateWorkspaceArgs {
init?: InputMaybe<Scalars['Upload']['input']>;
}

export interface MutationDeactivateLicenseArgs {
workspaceId: Scalars['String']['input'];
}

export interface MutationDeleteBlobArgs {
hash?: InputMaybe<Scalars['String']['input']>;
key?: InputMaybe<Scalars['String']['input']>;
@@ -1187,6 +1222,11 @@ export interface QueryChatHistoriesInput {
skip?: InputMaybe<Scalars['Int']['input']>;
}

export interface QueryTooLongDataType {
__typename?: 'QueryTooLongDataType';
max: Scalars['Int']['output'];
}

export interface QuotaQueryType {
__typename?: 'QuotaQueryType';
blobLimit: Scalars['SafeInt']['output'];
@@ -1559,6 +1599,8 @@ export interface WorkspaceType {
/** Get user invoice count */
invoiceCount: Scalars['Int']['output'];
invoices: Array<InvoiceType>;
/** The selfhost license of the workspace */
license: Maybe<License>;
/** member count of workspace */
memberCount: Scalars['Int']['output'];
/** Members of workspace */
@@ -1600,6 +1642,7 @@ export interface WorkspaceTypeInvoicesArgs {
}

export interface WorkspaceTypeMembersArgs {
query?: InputMaybe<Scalars['String']['input']>;
skip?: InputMaybe<Scalars['Int']['input']>;
take?: InputMaybe<Scalars['Int']['input']>;
}
@@ -1792,14 +1835,19 @@ export type ListContextFilesQuery = {
__typename?: 'Copilot';
contexts: Array<{
__typename?: 'CopilotContext';
docs: Array<string>;
docs: Array<{
__typename?: 'CopilotContextDoc';
id: string;
createdAt: number;
}>;
files: Array<{
__typename?: 'CopilotContextFile';
id: string;
name: string;
blobId: string;
chunk_size: number;
status: ContextFileStatus;
createdAt: number;
}>;
}>;
};
@@ -1843,7 +1891,11 @@ export type ListContextQuery = {
__typename?: 'UserType';
copilot: {
__typename?: 'Copilot';
contexts: Array<{ __typename?: 'CopilotContext'; id: string }>;
contexts: Array<{
__typename?: 'CopilotContext';
id: string;
createdAt: number;
}>;
};
} | null;
};

0 comments on commit 460e3e6

Please sign in to comment.