feat(server): basic context api (#10056)

fix CLOUD-97
fix CLOUD-98

packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md

@@ -0,0 +1,26 @@
+# Snapshot report for `src/__tests__/copilot.e2e.ts`
+
+The actual snapshot is saved in `copilot.e2e.ts.snap`.
+
+Generated by [AVA](https://avajs.dev).
+
+## should be able to manage context
+
+> should list context files
+
+    [
+      {
+        id: 'docId1',
+      },
+    ]
+
+> should list context docs
+
+    [
+      {
+        blobId: 'fileId1',
+        chunkSize: 3,
+        name: 'sample.pdf',
+        status: 'finished',
+      },
+    ]

packages/backend/server/src/__tests__/copilot.e2e.ts

@@ -8,6 +8,7 @@ import { ConfigModule } from '../base/config';
 import { AuthService } from '../core/auth';
 import { WorkspaceModule } from '../core/workspaces';
 import { CopilotModule } from '../plugins/copilot';
+import { CopilotContextService } from '../plugins/copilot/context';
 import { prompts, PromptService } from '../plugins/copilot/prompt';
 import {
   CopilotProviderService,
@@ -27,15 +28,19 @@ import {
   TestUser,
 } from './utils';
 import {
+  addContextDoc,
   array2sse,
   chatWithImages,
   chatWithText,
   chatWithTextStream,
   chatWithWorkflow,
+  createCopilotContext,
   createCopilotMessage,
   createCopilotSession,
   forkCopilotSession,
   getHistories,
+  listContext,
+  listContextFiles,
   MockCopilotTestProvider,
   sse2array,
   textToEventStream,
@@ -46,6 +51,7 @@ import {
 const test = ava as TestFn<{
   auth: AuthService;
   app: TestingApp;
+  context: CopilotContextService;
   prompt: PromptService;
   provider: CopilotProviderService;
   storage: CopilotStorage;
@@ -77,11 +83,13 @@ test.before(async t => {
   });
 
   const auth = app.get(AuthService);
+  const context = app.get(CopilotContextService);
   const prompt = app.get(PromptService);
   const storage = app.get(CopilotStorage);
 
   t.context.app = app;
   t.context.auth = auth;
+  t.context.context = context;
   t.context.prompt = prompt;
   t.context.storage = storage;
 });
@@ -678,3 +686,46 @@ test('should be able to search image from unsplash', async t => {
   const resp = await unsplashSearch(app);
   t.not(resp.status, 404, 'route should be exists');
 });
+
+test('should be able to manage context', async t => {
+  const { app } = t.context;
+
+  const { id: workspaceId } = await createWorkspace(app);
+  const sessionId = await createCopilotSession(
+    app,
+    workspaceId,
+    randomUUID(),
+    promptName
+  );
+
+  {
+    await t.throwsAsync(
+      createCopilotContext(app, workspaceId, randomUUID()),
+      { instanceOf: Error },
+      'should throw error if create context with invalid session id'
+    );
+
+    const context = createCopilotContext(app, workspaceId, sessionId);
+    await t.notThrowsAsync(context, 'should create context with chat session');
+
+    const list = await listContext(app, workspaceId, sessionId);
+    t.deepEqual(
+      list.map(f => ({ id: f.id })),
+      [{ id: await context }],
+      'should list context'
+    );
+  }
+
+  {
+    const contextId = await createCopilotContext(app, workspaceId, sessionId);
+
+    await addContextDoc(app, contextId, 'docId1');
+
+    const { docs } =
+      (await listContextFiles(app, workspaceId, sessionId, contextId)) || {};
+    t.snapshot(
+      docs?.map(({ createdAt: _, ...d }) => d),
+      'should list context files'
+    );
+  }
+});

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

@@ -1,3 +1,5 @@
+import { randomUUID } from 'node:crypto';
+
 import type { TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';
@@ -6,6 +8,7 @@ import { ConfigModule } from '../base/config';
 import { AuthService } from '../core/auth';
 import { QuotaModule } from '../core/quota';
 import { CopilotModule } from '../plugins/copilot';
+import { CopilotContextService } from '../plugins/copilot/context';
 import { prompts, PromptService } from '../plugins/copilot/prompt';
 import {
   CopilotProviderService,
@@ -44,6 +47,7 @@ import { MockCopilotTestProvider, WorkflowTestCases } from './utils/copilot';
 const test = ava as TestFn<{
   auth: AuthService;
   module: TestingModule;
+  context: CopilotContextService;
   prompt: PromptService;
   provider: CopilotProviderService;
   session: ChatSessionService;
@@ -81,6 +85,7 @@ test.before(async t => {
   });
 
   const auth = module.get(AuthService);
+  const context = module.get(CopilotContextService);
   const prompt = module.get(PromptService);
   const provider = module.get(CopilotProviderService);
   const session = module.get(ChatSessionService);
@@ -88,6 +93,7 @@ test.before(async t => {
 
   t.context.module = module;
   t.context.auth = auth;
+  t.context.context = context;
   t.context.prompt = prompt;
   t.context.provider = provider;
   t.context.session = session;
@@ -1247,3 +1253,52 @@ test('CitationParser should not replace chunks of citation already with URLs', t
   ].join('\n');
   t.is(result, expected);
 });
+
+// ==================== context ====================
+test('should be able to manage context', async t => {
+  const { context, prompt, session } = t.context;
+
+  await prompt.set('prompt', 'model', [
+    { role: 'system', content: 'hello {{word}}' },
+  ]);
+  const chatSession = await session.create({
+    docId: 'test',
+    workspaceId: 'test',
+    userId,
+    promptName: 'prompt',
+  });
+
+  {
+    await t.throwsAsync(
+      context.create(randomUUID()),
+      { instanceOf: Error },
+      'should throw error if create context with invalid session id'
+    );
+
+    const session = context.create(chatSession);
+    await t.notThrowsAsync(session, 'should create context with chat session');
+
+    await t.notThrowsAsync(
+      context.get((await session).id),
+      'should get context after create'
+    );
+
+    await t.throwsAsync(
+      context.get(randomUUID()),
+      { instanceOf: Error },
+      'should throw error if get context with invalid id'
+    );
+  }
+
+  {
+    const session = await context.create(chatSession);
+
+    const docId = randomUUID();
+    await session.addDocRecord(docId);
+    const docs = session.listDocs().map(d => d.id);
+    t.deepEqual(docs, [docId], 'should list doc id');
+
+    await session.removeDocRecord(docId);
+    t.deepEqual(session.listDocs(), [], 'should remove doc id');
+  }
+});

packages/backend/server/src/__tests__/utils/copilot.ts

@@ -23,6 +23,7 @@ import {
   WorkflowNodeType,
   WorkflowParams,
 } from '../../plugins/copilot/workflow/types';
+import { gql } from './common';
 import { TestingApp } from './testing-app';
 import { sleep } from './utils';
 
@@ -209,6 +210,216 @@ export async function forkCopilotSession(
   return res.forkCopilotSession;
 }
 
+export async function createCopilotContext(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string
+): Promise<string> {
+  const res = await app.gql(`
+        mutation {
+          createCopilotContext(workspaceId: "${workspaceId}", sessionId: "${sessionId}")
+        }
+      `);
+
+  return res.createCopilotContext;
+}
+
+export async function matchContext(
+  app: TestingApp,
+  contextId: string,
+  content: string,
+  limit: number
+): Promise<
+  | {
+      fileId: string;
+      chunk: number;
+      content: string;
+      distance: number | null;
+    }[]
+  | undefined
+> {
+  const res = await app.gql(
+    `
+        mutation matchContext($content: String!, $contextId: String!, $limit: SafeInt) {
+          matchContext(content: $content, contextId: $contextId, limit: $limit) {
+            fileId
+            chunk
+            content
+            distance
+          }
+        }
+      `,
+    { contextId, content, limit }
+  );
+
+  return res.matchContext;
+}
+
+export async function listContext(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string
+): Promise<
+  {
+    id: string;
+    workspaceId: string;
+  }[]
+> {
+  const res = await app.gql(`
+        query {
+          currentUser {
+            copilot(workspaceId: "${workspaceId}") {
+              contexts(sessionId: "${sessionId}") {
+                id
+                workspaceId
+              }
+            }
+          }
+        }
+      `);
+
+  return res.currentUser?.copilot?.contexts;
+}
+
+export async function addContextFile(
+  app: TestingApp,
+  contextId: string,
+  blobId: string,
+  fileName: string,
+  content: Buffer
+): Promise<{ id: string }[]> {
+  const res = await app
+    .POST(gql)
+    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .field(
+      'operations',
+      JSON.stringify({
+        query: `
+          mutation addContextFile($options: AddContextFileInput!, $content: Upload!) {
+            addContextFile(content: $content, options: $options) {
+              id
+            }
+          }
+        `,
+        variables: {
+          content: null,
+          options: { contextId, blobId, fileName },
+        },
+      })
+    )
+    .field('map', JSON.stringify({ '0': ['variables.content'] }))
+    .attach('0', content, {
+      filename: fileName,
+      contentType: 'application/octet-stream',
+    })
+    .expect(200);
+
+  return res.body.data.addContextFile;
+}
+
+export async function removeContextFile(
+  app: TestingApp,
+  contextId: string,
+  fileId: string
+): Promise<string> {
+  const res = await app.gql(
+    `
+        mutation removeContextFile($options: RemoveContextFileInput!) {
+          removeContextFile(options: $options)
+        }
+    `,
+    { options: { contextId, fileId } }
+  );
+
+  return res.removeContextFile;
+}
+
+export async function addContextDoc(
+  app: TestingApp,
+  contextId: string,
+  docId: string
+): Promise<{ id: string }[]> {
+  const res = await app.gql(
+    `
+          mutation addContextDoc($options: AddContextDocInput!) {
+            addContextDoc(options: $options) {
+              id
+            }
+          }
+        `,
+    { options: { contextId, docId } }
+  );
+
+  return res.addContextDoc;
+}
+
+export async function removeContextDoc(
+  app: TestingApp,
+  contextId: string,
+  docId: string
+): Promise<string> {
+  const res = await app.gql(
+    `
+      mutation removeContextDoc($options: RemoveContextFileInput!) {
+        removeContextDoc(options: $options)
+      }
+    `,
+    { options: { contextId, docId } }
+  );
+
+  return res.removeContextDoc;
+}
+
+export async function listContextFiles(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string,
+  contextId: string
+): Promise<
+  | {
+      docs: {
+        id: string;
+        createdAt: number;
+      }[];
+      files: {
+        id: string;
+        name: string;
+        blobId: string;
+        chunkSize: number;
+        status: string;
+        createdAt: number;
+      }[];
+    }
+  | undefined
+> {
+  const res = await app.gql(`
+        query {
+          currentUser {
+            copilot(workspaceId: "${workspaceId}") {
+              contexts(sessionId: "${sessionId}", contextId: "${contextId}") {
+                docs {
+                  id
+                  createdAt
+                }
+                files {
+                  id
+                  name
+                  blobId
+                  chunkSize
+                  status
+                  createdAt
+                }
+              }
+            }
+          }
+        }
+      `);
+
+  const { docs, files } = res.currentUser?.copilot?.contexts?.[0] || {};
+
+  return { docs, files };
+}
+
 export async function createCopilotMessage(
   app: TestingApp,
   sessionId: string,