Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-07-13 08:06:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: fork session support (#7367)

Browse Source
This commit is contained in:
darkskygit
2024-07-03 03:10:09 +00:00
parent 10df1fb4b7
commit 61870c04d0
12 changed files with 292 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/migrations/20240625052649_add_fork_session/migration.sql
+2
View File
@@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "ai_sessions_metadata" ADD COLUMN "parent_session_id" VARCHAR(36);
packages/backend/server/schema.prisma
+11 -9
View File
@@ -481,15 +481,17 @@ model AiSessionMessage {
}
model AiSession {
id String @id @default(uuid()) @db.VarChar(36)
userId String @map("user_id") @db.VarChar(36)
workspaceId String @map("workspace_id") @db.VarChar(36)
docId String @map("doc_id") @db.VarChar(36)
promptName String @map("prompt_name") @db.VarChar(32)
messageCost Int @default(0)
tokenCost Int @default(0)
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
deletedAt DateTime? @map("deleted_at") @db.Timestamptz(6)
id String @id @default(uuid()) @db.VarChar(36)
userId String @map("user_id") @db.VarChar(36)
workspaceId String @map("workspace_id") @db.VarChar(36)
docId String @map("doc_id") @db.VarChar(36)
promptName String @map("prompt_name") @db.VarChar(32)
// the session id of the parent session if this session is a forked session
parentSessionId String? @map("parent_session_id") @db.VarChar(36)
messageCost Int @default(0)
tokenCost Int @default(0)
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
deletedAt DateTime? @map("deleted_at") @db.Timestamptz(6)
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
prompt AiPrompt @relation(fields: [promptName], references: [name], onDelete: Cascade)
packages/backend/server/src/fundamentals/error/def.ts
+2 -1
View File
@@ -440,7 +440,8 @@ export const USER_FRIENDLY_ERRORS = {
},
copilot_message_not_found: {
type: 'resource_not_found',
message: `Copilot message not found.`,
args: { messageId: 'string' },
message: ({ messageId }) => `Copilot message ${messageId} not found.`,
},
copilot_prompt_not_found: {
type: 'resource_not_found',
packages/backend/server/src/fundamentals/error/errors.gen.ts
+7 -3
View File
@@ -391,10 +391,14 @@ export class CopilotActionTaken extends UserFriendlyError {
super('action_forbidden', 'copilot_action_taken', message);
}
}
@ObjectType()
class CopilotMessageNotFoundDataType {
@Field() messageId!: string
}
export class CopilotMessageNotFound extends UserFriendlyError {
constructor(message?: string) {
super('resource_not_found', 'copilot_message_not_found', message);
constructor(args: CopilotMessageNotFoundDataType, message?: string | ((args: CopilotMessageNotFoundDataType) => string)) {
super('resource_not_found', 'copilot_message_not_found', message, args);
}
}
@ObjectType()
@@ -542,5 +546,5 @@ registerEnumType(ErrorNames, {
export const ErrorDataUnionType = createUnionType({
name: 'ErrorDataUnion',
types: () =>
[UnknownOauthProviderDataType, MissingOauthQueryParameterDataType, InvalidPasswordLengthDataType, WorkspaceNotFoundDataType, NotInWorkspaceDataType, WorkspaceAccessDeniedDataType, WorkspaceOwnerNotFoundDataType, DocNotFoundDataType, DocAccessDeniedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderSideErrorDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType] as const,
[UnknownOauthProviderDataType, MissingOauthQueryParameterDataType, InvalidPasswordLengthDataType, WorkspaceNotFoundDataType, NotInWorkspaceDataType, WorkspaceAccessDeniedDataType, WorkspaceOwnerNotFoundDataType, DocNotFoundDataType, DocAccessDeniedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderSideErrorDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType] as const,
});
packages/backend/server/src/plugins/copilot/resolver.ts
+54 -4
View File
@@ -27,7 +27,7 @@ import {
FileUpload,
MutexService,
Throttle,
TooManyRequestsException,
TooManyRequest,
} from '../../fundamentals';
import { PromptService } from './prompt';
import { ChatSessionService } from './session';
@@ -60,6 +60,24 @@ class CreateChatSessionInput {
promptName!: string;
}
@InputType()
class ForkChatSessionInput {
@Field(() => String)
workspaceId!: string;
@Field(() => String)
docId!: string;
@Field(() => String)
sessionId!: string;
@Field(() => String, {
description:
'Identify a message in the array and keep it with all previous messages into a forked session.',
})
latestMessageId!: string;
}
@InputType()
class DeleteSessionInput {
@Field(() => String)
@@ -109,6 +127,10 @@ class QueryChatHistoriesInput implements Partial<ListHistoriesOptions> {
@ObjectType('ChatMessage')
class ChatMessageType implements Partial<ChatMessage> {
// id will be null if message is a prompt message
@Field(() => ID, { nullable: true })
id!: string;
@Field(() => String)
role!: 'system' | 'assistant' | 'user';
@@ -301,7 +323,7 @@ export class CopilotResolver {
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
await using lock = await this.mutex.lock(lockFlag);
if (!lock) {
return new TooManyRequestsException('Server is busy');
return new TooManyRequest('Server is busy');
}
await this.chatSession.checkQuota(user.id);
@@ -313,6 +335,34 @@ export class CopilotResolver {
return session;
}
@Mutation(() => String, {
description: 'Create a chat session',
})
async forkCopilotSession(
@CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => ForkChatSessionInput })
options: ForkChatSessionInput
) {
await this.permissions.checkCloudPagePermission(
options.workspaceId,
options.docId,
user.id
);
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
await using lock = await this.mutex.lock(lockFlag);
if (!lock) {
return new TooManyRequest('Server is busy');
}
await this.chatSession.checkQuota(user.id);
const session = await this.chatSession.fork({
...options,
userId: user.id,
});
return session;
}
@Mutation(() => [String], {
description: 'Cleanup sessions',
})
@@ -332,7 +382,7 @@ export class CopilotResolver {
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
await using lock = await this.mutex.lock(lockFlag);
if (!lock) {
return new TooManyRequestsException('Server is busy');
return new TooManyRequest('Server is busy');
}
return await this.chatSession.cleanup({
@@ -352,7 +402,7 @@ export class CopilotResolver {
const lockFlag = `${COPILOT_LOCKER}:message:${user?.id}:${options.sessionId}`;
await using lock = await this.mutex.lock(lockFlag);
if (!lock) {
return new TooManyRequestsException('Server is busy');
return new TooManyRequest('Server is busy');
}
const session = await this.chatSession.get(options.sessionId);
if (!session || session.config.userId !== user.id) {
packages/backend/server/src/plugins/copilot/session.ts
+40 -4
View File
@@ -20,6 +20,7 @@ import {
ChatHistory,
ChatMessage,
ChatMessageSchema,
ChatSessionForkOptions,
ChatSessionOptions,
ChatSessionState,
getTokenEncoder,
@@ -81,7 +82,7 @@ export class ChatSession implements AsyncDisposable {
async getMessageById(messageId: string) {
const message = await this.messageCache.get(messageId);
if (!message || message.sessionId !== this.state.sessionId) {
throw new CopilotMessageNotFound();
throw new CopilotMessageNotFound({ messageId });
}
return message;
}
@@ -89,7 +90,7 @@ export class ChatSession implements AsyncDisposable {
async pushByMessageId(messageId: string) {
const message = await this.messageCache.get(messageId);
if (!message || message.sessionId !== this.state.sessionId) {
throw new CopilotMessageNotFound();
throw new CopilotMessageNotFound({ messageId });
}
this.push({
@@ -200,6 +201,7 @@ export class ChatSessionService {
workspaceId: state.workspaceId,
docId: state.docId,
prompt: { action: { equals: null } },
parentSessionId: state.parentSessionId,
},
select: { id: true, deletedAt: true },
})) || {};
@@ -271,8 +273,9 @@ export class ChatSessionService {
userId: true,
workspaceId: true,
docId: true,
parentSessionId: true,
messages: {
select: { role: true, content: true, createdAt: true },
select: { id: true, role: true, content: true, createdAt: true },
orderBy: { createdAt: 'asc' },
},
promptName: true,
@@ -291,6 +294,7 @@ export class ChatSessionService {
userId: session.userId,
workspaceId: session.workspaceId,
docId: session.docId,
parentSessionId: session.parentSessionId,
prompt,
messages: messages.success ? messages.data : [],
};
@@ -396,6 +400,7 @@ export class ChatSessionService {
createdAt: true,
messages: {
select: {
id: true,
role: true,
content: true,
attachments: true,
@@ -430,7 +435,8 @@ export class ChatSessionService {
.filter(({ role }) => role !== 'system')
: [];
// `createdAt` is required for history sorting in frontend, let's fake the creating time of prompt messages
// `createdAt` is required for history sorting in frontend
// let's fake the creating time of prompt messages
(preload as ChatMessage[]).forEach((msg, i) => {
msg.createdAt = new Date(
createdAt.getTime() - preload.length - i - 1
@@ -495,9 +501,39 @@ export class ChatSessionService {
sessionId,
prompt,
messages: [],
// when client create chat session, we always find root session
parentSessionId: null,
});
}
async fork(options: ChatSessionForkOptions): Promise<string> {
const state = await this.getSession(options.sessionId);
if (!state) {
throw new CopilotSessionNotFound();
}
const lastMessageIdx = state.messages.findLastIndex(
({ id, role }) =>
role === AiPromptRole.assistant && id === options.latestMessageId
);
if (lastMessageIdx < 0) {
throw new CopilotMessageNotFound({ messageId: options.latestMessageId });
}
const messages = state.messages
.slice(0, lastMessageIdx + 1)
.map(m => ({ ...m, id: undefined }));
const forkedState = {
...state,
sessionId: randomUUID(),
messages: [],
parentSessionId: options.sessionId,
};
// create session
await this.setSession(forkedState);
// save message
return await this.setSession({ ...forkedState, messages });
}
async cleanup(
options: Omit<ChatSessionOptions, 'promptName'> & { sessionIds: string[] }
) {
packages/backend/server/src/plugins/copilot/types.ts
+8
View File
@@ -64,6 +64,7 @@ export type PromptMessage = z.infer<typeof PromptMessageSchema>;
export type PromptParams = NonNullable<PromptMessage['params']>;
export const ChatMessageSchema = PromptMessageSchema.extend({
id: z.string().optional(),
createdAt: z.date(),
}).strict();
@@ -98,10 +99,17 @@ export interface ChatSessionOptions {
promptName: string;
}
export interface ChatSessionForkOptions
extends Omit<ChatSessionOptions, 'promptName'> {
sessionId: string;
latestMessageId: string;
}
export interface ChatSessionState
extends Omit<ChatSessionOptions, 'promptName'> {
// connect ids
sessionId: string;
parentSessionId: string | null;
// states
prompt: ChatPrompt;
messages: ChatMessage[];
packages/backend/server/src/schema.gql
+20 -1
View File
@@ -11,6 +11,7 @@ type ChatMessage {
attachments: [String!]
content: String!
createdAt: DateTime!
id: ID
params: JSON
role: String!
}
@@ -39,6 +40,10 @@ type CopilotHistories {
tokens: Int!
}
type CopilotMessageNotFoundDataType {
messageId: String!
}
enum CopilotModels {
DallE3
Gpt4Omni
@@ -175,7 +180,7 @@ enum EarlyAccessType {
App
}
union ErrorDataUnion = BlobNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderSideErrorDataType | DocAccessDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | InvalidHistoryTimestampDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | MissingOauthQueryParameterDataType | NotInWorkspaceDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | VersionRejectedDataType | WorkspaceAccessDeniedDataType | WorkspaceNotFoundDataType | WorkspaceOwnerNotFoundDataType
union ErrorDataUnion = BlobNotFoundDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderSideErrorDataType | DocAccessDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | InvalidHistoryTimestampDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | MissingOauthQueryParameterDataType | NotInWorkspaceDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | VersionRejectedDataType | WorkspaceAccessDeniedDataType | WorkspaceNotFoundDataType | WorkspaceOwnerNotFoundDataType
enum ErrorNames {
ACCESS_DENIED
@@ -252,6 +257,17 @@ enum FeatureType {
UnlimitedWorkspace
}
input ForkChatSessionInput {
docId: String!
"""
Identify a message in the array and keep it with all previous messages into a forked session.
"""
latestMessageId: String!
sessionId: String!
workspaceId: String!
}
type HumanReadableQuotaType {
blobLimit: String!
copilotActionLimit: String
@@ -399,6 +415,9 @@ type Mutation {
"""Delete a user account"""
deleteUser(id: String!): DeleteAccount!
deleteWorkspace(id: String!): Boolean!
"""Create a chat session"""
forkCopilotSession(options: ForkChatSessionInput!): String!
invite(email: String!, permission: Permission!, sendInviteMail: Boolean, workspaceId: String!): String!
leaveWorkspace(sendLeaveMail: Boolean, workspaceId: String!, workspaceName: String!): Boolean!
publishPage(mode: PublicPageMode = Page, pageId: String!, workspaceId: String!): WorkspacePage!
packages/backend/server/tests/copilot.spec.ts
+99 -6
View File
@@ -208,11 +208,13 @@ test('should be able to manage chat session', async t => {
{ role: 'system', content: 'hello {{word}}' },
]);
const params = { word: 'world' };
const commonParams = { docId: 'test', workspaceId: 'test' };
const sessionId = await session.create({
docId: 'test',
workspaceId: 'test',
userId,
promptName: 'prompt',
...commonParams,
});
t.truthy(sessionId, 'should create session');
@@ -221,8 +223,6 @@ test('should be able to manage chat session', async t => {
t.is(s.config.promptName, 'prompt', 'should have prompt name');
t.is(s.model, 'model', 'should have model');
const params = { word: 'world' };
s.push({ role: 'user', content: 'hello', createdAt: new Date() });
// @ts-expect-error
const finalMessages = s.finish(params).map(({ createdAt: _, ...m }) => m);
@@ -239,19 +239,112 @@ test('should be able to manage chat session', async t => {
const s1 = (await session.get(sessionId))!;
t.deepEqual(
// @ts-expect-error
s1.finish(params).map(({ createdAt: _, ...m }) => m),
s1.finish(params).map(({ id: _, createdAt: __, ...m }) => m),
finalMessages,
'should same as before message'
);
t.deepEqual(
// @ts-expect-error
s1.finish({}).map(({ createdAt: _, ...m }) => m),
s1.finish({}).map(({ id: _, createdAt: __, ...m }) => m),
[
{ content: 'hello ', params: {}, role: 'system' },
{ content: 'hello', role: 'user' },
],
'should generate different message with another params'
);
// should get main session after fork if re-create a chat session for same docId and workspaceId
{
const newSessionId = await session.create({
userId,
promptName: 'prompt',
...commonParams,
});
t.is(newSessionId, sessionId, 'should get same session id');
}
});
test('should be able to fork chat session', async t => {
const { prompt, session } = t.context;
await prompt.set('prompt', 'model', [
{ role: 'system', content: 'hello {{word}}' },
]);
const params = { word: 'world' };
const commonParams = { docId: 'test', workspaceId: 'test' };
// create session
const sessionId = await session.create({
userId,
promptName: 'prompt',
...commonParams,
});
const s = (await session.get(sessionId))!;
s.push({ role: 'user', content: 'hello', createdAt: new Date() });
s.push({ role: 'assistant', content: 'world', createdAt: new Date() });
s.push({ role: 'user', content: 'aaa', createdAt: new Date() });
s.push({ role: 'assistant', content: 'bbb', createdAt: new Date() });
await s.save();
// fork session
const s1 = (await session.get(sessionId))!;
// @ts-expect-error
const latestMessageId = s1.finish({}).find(m => m.role === 'assistant')!.id;
const forkedSessionId = await session.fork({
userId,
sessionId,
latestMessageId,
...commonParams,
});
t.not(sessionId, forkedSessionId, 'should fork a new session');
// check forked session messages
{
const s2 = (await session.get(forkedSessionId))!;
const finalMessages = s2
.finish(params) // @ts-expect-error
.map(({ id: _, createdAt: __, ...m }) => m);
t.deepEqual(
finalMessages,
[
{ role: 'system', content: 'hello world', params },
{ role: 'user', content: 'hello' },
{ role: 'assistant', content: 'world' },
],
'should generate the final message'
);
}
// check original session messages
{
const s3 = (await session.get(sessionId))!;
const finalMessages = s3
.finish(params) // @ts-expect-error
.map(({ id: _, createdAt: __, ...m }) => m);
t.deepEqual(
finalMessages,
[
{ role: 'system', content: 'hello world', params },
{ role: 'user', content: 'hello' },
{ role: 'assistant', content: 'world' },
{ role: 'user', content: 'aaa' },
{ role: 'assistant', content: 'bbb' },
],
'should generate the final message'
);
}
// should get main session after fork if re-create a chat session for same docId and workspaceId
{
const newSessionId = await session.create({
userId,
promptName: 'prompt',
...commonParams,
});
t.is(newSessionId, sessionId, 'should get same session id');
}
});
test('should be able to process message id', async t => {
packages/frontend/graphql/src/graphql/fork-copilot-session.gql
+3
View File
@@ -0,0 +1,3 @@
mutation forkCopilotSession($options: ForkChatSessionInput!) {
forkCopilotSession(options: $options)
}
packages/frontend/graphql/src/graphql/index.ts
+11
View File
@@ -241,6 +241,17 @@ mutation removeEarlyAccess($email: String!) {
}`,
};
export const forkCopilotSessionMutation = {
id: 'forkCopilotSessionMutation' as const,
operationName: 'forkCopilotSession',
definitionName: 'forkCopilotSession',
containsFile: false,
query: `
mutation forkCopilotSession($options: ForkChatSessionInput!) {
forkCopilotSession(options: $options)
}`,
};
export const getCopilotHistoriesQuery = {
id: 'getCopilotHistoriesQuery' as const,
operationName: 'getCopilotHistories',
packages/frontend/graphql/src/schema.ts
+35
View File
@@ -49,6 +49,7 @@ export interface ChatMessage {
attachments: Maybe<Array<Scalars['String']['output']>>;
content: Scalars['String']['output'];
createdAt: Scalars['DateTime']['output'];
id: Maybe<Scalars['ID']['output']>;
params: Maybe<Scalars['JSON']['output']>;
role: Scalars['String']['output'];
}
@@ -81,6 +82,11 @@ export interface CopilotHistories {
tokens: Scalars['Int']['output'];
}
export interface CopilotMessageNotFoundDataType {
__typename?: 'CopilotMessageNotFoundDataType';
messageId: Scalars['String']['output'];
}
export enum CopilotModels {
DallE3 = 'DallE3',
Gpt4Omni = 'Gpt4Omni',
@@ -224,6 +230,7 @@ export enum EarlyAccessType {
export type ErrorDataUnion =
| BlobNotFoundDataType
| CopilotMessageNotFoundDataType
| CopilotPromptNotFoundDataType
| CopilotProviderSideErrorDataType
| DocAccessDeniedDataType
@@ -320,6 +327,14 @@ export enum FeatureType {
UnlimitedWorkspace = 'UnlimitedWorkspace',
}
export interface ForkChatSessionInput {
docId: Scalars['String']['input'];
/** Identify a message in the array and keep it with all previous messages into a forked session. */
latestMessageId: Scalars['String']['input'];
sessionId: Scalars['String']['input'];
workspaceId: Scalars['String']['input'];
}
export interface HumanReadableQuotaType {
__typename?: 'HumanReadableQuotaType';
blobLimit: Scalars['String']['output'];
@@ -449,6 +464,8 @@ export interface Mutation {
/** Delete a user account */
deleteUser: DeleteAccount;
deleteWorkspace: Scalars['Boolean']['output'];
/** Create a chat session */
forkCopilotSession: Scalars['String']['output'];
invite: Scalars['String']['output'];
leaveWorkspace: Scalars['Boolean']['output'];
publishPage: WorkspacePage;
@@ -562,6 +579,10 @@ export interface MutationDeleteWorkspaceArgs {
id: Scalars['String']['input'];
}
export interface MutationForkCopilotSessionArgs {
options: ForkChatSessionInput;
}
export interface MutationInviteArgs {
email: Scalars['String']['input'];
permission: Permission;
@@ -1340,6 +1361,15 @@ export type RemoveEarlyAccessMutation = {
removeEarlyAccess: number;
};
export type ForkCopilotSessionMutationVariables = Exact<{
options: ForkChatSessionInput;
}>;
export type ForkCopilotSessionMutation = {
__typename?: 'Mutation';
forkCopilotSession: string;
};
export type CredentialsRequirementFragment = {
__typename?: 'CredentialsRequirementType';
password: {
@@ -2354,6 +2384,11 @@ export type Mutations =
variables: RemoveEarlyAccessMutationVariables;
response: RemoveEarlyAccessMutation;
}
| {
name: 'forkCopilotSessionMutation';
variables: ForkCopilotSessionMutationVariables;
response: ForkCopilotSessionMutation;
}
| {
name: 'leaveWorkspaceMutation';
variables: LeaveWorkspaceMutationVariables;