luulam/AFFiNE
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

chore(server): improve gql types (#11441)

Browse Source
This commit is contained in:
darkskygit
2025-04-03 07:30:51 +00:00
parent 70a318f1c4
commit b4c643e8bc
5 changed files with 61 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
packages/backend/server/src/__tests__/copilot.spec.ts
View File

@@ -443,13 +443,13 @@ test('should be able to process message id', async t => {
}); });
const s = (await session.get(sessionId))!; const s = (await session.get(sessionId))!;
const textMessage = (await session.createMessage({ const textMessage = await session.createMessage({
sessionId, sessionId,
content: 'hello', content: 'hello',
}))!; });
const anotherSessionMessage = (await session.createMessage({ const anotherSessionMessage = await session.createMessage({
sessionId: 'another-session-id', sessionId: 'another-session-id',
}))!; });
await t.notThrowsAsync( await t.notThrowsAsync(
s.pushByMessageId(textMessage), s.pushByMessageId(textMessage),
@@ -486,10 +486,10 @@ test('should be able to generate with message id', async t => {
}); });
const s = (await session.get(sessionId))!; const s = (await session.get(sessionId))!;
const message = (await session.createMessage({ const message = await session.createMessage({
sessionId, sessionId,
content: 'hello', content: 'hello',
}))!; });
await s.pushByMessageId(message); await s.pushByMessageId(message);
const finalMessages = s const finalMessages = s
@@ -508,10 +508,10 @@ test('should be able to generate with message id', async t => {
}); });
const s = (await session.get(sessionId))!; const s = (await session.get(sessionId))!;
const message = (await session.createMessage({ const message = await session.createMessage({
sessionId, sessionId,
attachments: ['https://affine.pro/example.jpg'], attachments: ['https://affine.pro/example.jpg'],
}))!; });
await s.pushByMessageId(message); await s.pushByMessageId(message);
const finalMessages = s const finalMessages = s
@@ -535,9 +535,9 @@ test('should be able to generate with message id', async t => {
}); });
const s = (await session.get(sessionId))!; const s = (await session.get(sessionId))!;
const message = (await session.createMessage({ const message = await session.createMessage({
sessionId, sessionId,
}))!; });
await s.pushByMessageId(message); await s.pushByMessageId(message);
const finalMessages = s const finalMessages = s
@@ -563,10 +563,10 @@ test('should save message correctly', async t => {
}); });
const s = (await session.get(sessionId))!; const s = (await session.get(sessionId))!;
const message = (await session.createMessage({ const message = await session.createMessage({
sessionId, sessionId,
content: 'hello', content: 'hello',
}))!; });
await s.pushByMessageId(message); await s.pushByMessageId(message);
t.is(s.stashMessages.length, 1, 'should get stash messages'); t.is(s.stashMessages.length, 1, 'should get stash messages');
@@ -592,10 +592,10 @@ test('should revert message correctly', async t => {
}); });
const s = (await session.get(sessionId))!; const s = (await session.get(sessionId))!;
const message = (await session.createMessage({ const message = await session.createMessage({
sessionId, sessionId,
content: '1', content: '1',
}))!; });
await s.pushByMessageId(message); await s.pushByMessageId(message);
await s.save(); await s.save();

32
packages/backend/server/src/plugins/copilot/context/resolver.ts
View File

@@ -261,12 +261,12 @@ export class CopilotContextRootResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('sessionId', { nullable: true }) sessionId?: string, @Args('sessionId', { nullable: true }) sessionId?: string,
@Args('contextId', { nullable: true }) contextId?: string @Args('contextId', { nullable: true }) contextId?: string
) { ): Promise<CopilotContextType[]> {
if (sessionId || contextId) { if (sessionId || contextId) {
const lockFlag = `${COPILOT_LOCKER}:context:${sessionId || contextId}`; const lockFlag = `${COPILOT_LOCKER}:context:${sessionId || contextId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
if (contextId) { if (contextId) {
@@ -294,11 +294,11 @@ export class CopilotContextRootResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('workspaceId') workspaceId: string, @Args('workspaceId') workspaceId: string,
@Args('sessionId') sessionId: string @Args('sessionId') sessionId: string
) { ): Promise<string> {
const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`; const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
await this.checkChatSession(user, sessionId, workspaceId); await this.checkChatSession(user, sessionId, workspaceId);
@@ -314,7 +314,7 @@ export class CopilotContextRootResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('workspaceId') workspaceId: string, @Args('workspaceId') workspaceId: string,
@Args('docId', { type: () => [String] }) docIds: string[] @Args('docId', { type: () => [String] }) docIds: string[]
) { ): Promise<boolean> {
await this.ac await this.ac
.user(user.id) .user(user.id)
.workspace(workspaceId) .workspace(workspaceId)
@@ -339,7 +339,7 @@ export class CopilotContextRootResolver {
async queryWorkspaceEmbeddingStatus( async queryWorkspaceEmbeddingStatus(
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('workspaceId') workspaceId: string @Args('workspaceId') workspaceId: string
) { ): Promise<ContextWorkspaceEmbeddingStatus> {
await this.ac await this.ac
.user(user.id) .user(user.id)
.workspace(workspaceId) .workspace(workspaceId)
@@ -386,7 +386,7 @@ export class CopilotContextResolver {
@CallMetric('ai', 'context_file_list') @CallMetric('ai', 'context_file_list')
async collections( async collections(
@Parent() context: CopilotContextType @Parent() context: CopilotContextType
): Promise<ContextCategory[]> { ): Promise<CopilotContextCategory[]> {
const session = await this.context.get(context.id); const session = await this.context.get(context.id);
const collections = session.collections; const collections = session.collections;
await this.models.copilotContext.mergeDocStatus( await this.models.copilotContext.mergeDocStatus(
@@ -403,7 +403,7 @@ export class CopilotContextResolver {
@CallMetric('ai', 'context_file_list') @CallMetric('ai', 'context_file_list')
async tags( async tags(
@Parent() context: CopilotContextType @Parent() context: CopilotContextType
): Promise<ContextCategory[]> { ): Promise<CopilotContextCategory[]> {
const session = await this.context.get(context.id); const session = await this.context.get(context.id);
const tags = session.tags; const tags = session.tags;
await this.models.copilotContext.mergeDocStatus( await this.models.copilotContext.mergeDocStatus(
@@ -485,11 +485,11 @@ export class CopilotContextResolver {
async removeContextCategory( async removeContextCategory(
@Args({ name: 'options', type: () => RemoveContextCategoryInput }) @Args({ name: 'options', type: () => RemoveContextCategoryInput })
options: RemoveContextCategoryInput options: RemoveContextCategoryInput
) { ): Promise<boolean> {
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`; const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
const session = await this.context.get(options.contextId); const session = await this.context.get(options.contextId);
@@ -545,11 +545,11 @@ export class CopilotContextResolver {
async removeContextDoc( async removeContextDoc(
@Args({ name: 'options', type: () => RemoveContextDocInput }) @Args({ name: 'options', type: () => RemoveContextDocInput })
options: RemoveContextDocInput options: RemoveContextDocInput
) { ): Promise<boolean> {
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`; const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
const session = await this.context.get(options.contextId); const session = await this.context.get(options.contextId);
@@ -574,7 +574,7 @@ export class CopilotContextResolver {
options: AddContextFileInput, options: AddContextFileInput,
@Args({ name: 'content', type: () => GraphQLUpload }) @Args({ name: 'content', type: () => GraphQLUpload })
content: FileUpload content: FileUpload
) { ): Promise<CopilotContextFile> {
if (!this.context.canEmbedding) { if (!this.context.canEmbedding) {
throw new CopilotEmbeddingUnavailable(); throw new CopilotEmbeddingUnavailable();
} }
@@ -582,7 +582,7 @@ export class CopilotContextResolver {
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`; const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
const length = Number(ctx.req.headers['content-length']); const length = Number(ctx.req.headers['content-length']);
@@ -632,7 +632,7 @@ export class CopilotContextResolver {
async removeContextFile( async removeContextFile(
@Args({ name: 'options', type: () => RemoveContextFileInput }) @Args({ name: 'options', type: () => RemoveContextFileInput })
options: RemoveContextFileInput options: RemoveContextFileInput
) { ): Promise<boolean> {
if (!this.context.canEmbedding) { if (!this.context.canEmbedding) {
throw new CopilotEmbeddingUnavailable(); throw new CopilotEmbeddingUnavailable();
} }
@@ -640,7 +640,7 @@ export class CopilotContextResolver {
const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`; const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
const session = await this.context.get(options.contextId); const session = await this.context.get(options.contextId);

24
packages/backend/server/src/plugins/copilot/message.ts
View File

@@ -1,6 +1,6 @@
import { randomUUID } from 'node:crypto'; import { randomUUID } from 'node:crypto';
import { Injectable, Logger } from '@nestjs/common'; import { Injectable } from '@nestjs/common';
import { SessionCache } from '../../base'; import { SessionCache } from '../../base';
import { SubmittedMessage, SubmittedMessageSchema } from './types'; import { SubmittedMessage, SubmittedMessageSchema } from './types';
@@ -10,26 +10,18 @@ const CHAT_MESSAGE_TTL = 3600 * 1 * 1000; // 1 hours
@Injectable() @Injectable()
export class ChatMessageCache { export class ChatMessageCache {
private readonly logger = new Logger(ChatMessageCache.name);
constructor(private readonly cache: SessionCache) {} constructor(private readonly cache: SessionCache) {}
async get(id: string): Promise<SubmittedMessage | undefined> { async get(id: string): Promise<SubmittedMessage | undefined> {
return await this.cache.get(`${CHAT_MESSAGE_KEY}:${id}`); return await this.cache.get(`${CHAT_MESSAGE_KEY}:${id}`);
} }
async set(message: SubmittedMessage): Promise<string | undefined> { async set(message: SubmittedMessage): Promise<string> {
try { const parsedMessage = SubmittedMessageSchema.parse(message);
const parsed = SubmittedMessageSchema.safeParse(message); const id = randomUUID();
if (parsed.success) { await this.cache.set(`${CHAT_MESSAGE_KEY}:${id}`, parsedMessage, {
const id = randomUUID(); ttl: CHAT_MESSAGE_TTL,
await this.cache.set(`${CHAT_MESSAGE_KEY}:${id}`, parsed.data, { });
ttl: CHAT_MESSAGE_TTL, return id;
});
return id;
}
} catch (e: any) {
this.logger.error(`Failed to get chat message from cache: ${e.message}`);
}
return undefined;
} }
} }

42
packages/backend/server/src/plugins/copilot/resolver.ts
View File

@@ -169,7 +169,7 @@ class QueryChatHistoriesInput implements Partial<ListHistoriesOptions> {
class ChatMessageType implements Partial<ChatMessage> { class ChatMessageType implements Partial<ChatMessage> {
// id will be null if message is a prompt message // id will be null if message is a prompt message
@Field(() => ID, { nullable: true }) @Field(() => ID, { nullable: true })
id!: string; id!: string | undefined;
@Field(() => String) @Field(() => String)
role!: 'system' | 'assistant' | 'user'; role!: 'system' | 'assistant' | 'user';
@@ -310,7 +310,7 @@ export class CopilotResolver {
description: 'Get the quota of the user in the workspace', description: 'Get the quota of the user in the workspace',
complexity: 2, complexity: 2,
}) })
async getQuota(@CurrentUser() user: CurrentUser) { async getQuota(@CurrentUser() user: CurrentUser): Promise<CopilotQuotaType> {
return await this.chatSession.getQuota(user.id); return await this.chatSession.getQuota(user.id);
} }
@@ -324,8 +324,8 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('docId', { nullable: true }) docId?: string, @Args('docId', { nullable: true }) docId?: string,
@Args('options', { nullable: true }) options?: QueryChatSessionsInput @Args('options', { nullable: true }) options?: QueryChatSessionsInput
) { ): Promise<string[]> {
return await this.sessions(copilot, user, docId, options); return (await this.sessions(copilot, user, docId, options)).map(s => s.id);
} }
@ResolveField(() => [CopilotSessionType], { @ResolveField(() => [CopilotSessionType], {
@@ -337,7 +337,7 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('docId', { nullable: true }) docId?: string, @Args('docId', { nullable: true }) docId?: string,
@Args('options', { nullable: true }) options?: QueryChatSessionsInput @Args('options', { nullable: true }) options?: QueryChatSessionsInput
) { ): Promise<CopilotSessionType[]> {
if (!copilot.workspaceId) return []; if (!copilot.workspaceId) return [];
await this.ac await this.ac
.user(user.id) .user(user.id)
@@ -359,7 +359,7 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('docId', { nullable: true }) docId?: string, @Args('docId', { nullable: true }) docId?: string,
@Args('options', { nullable: true }) options?: QueryChatHistoriesInput @Args('options', { nullable: true }) options?: QueryChatHistoriesInput
) { ): Promise<CopilotHistoriesType[]> {
const workspaceId = copilot.workspaceId; const workspaceId = copilot.workspaceId;
if (!workspaceId) { if (!workspaceId) {
return []; return [];
@@ -387,7 +387,9 @@ export class CopilotResolver {
return histories.map(h => ({ return histories.map(h => ({
...h, ...h,
// filter out empty messages // filter out empty messages
messages: h.messages.filter(m => m.content || m.attachments?.length), messages: h.messages.filter(
m => m.content || m.attachments?.length
) as ChatMessageType[],
})); }));
} }
@@ -399,12 +401,12 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => CreateChatSessionInput }) @Args({ name: 'options', type: () => CreateChatSessionInput })
options: CreateChatSessionInput options: CreateChatSessionInput
) { ): Promise<string> {
await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update'); await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update');
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`; const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
if (options.workspaceId === options.docId) { if (options.workspaceId === options.docId) {
@@ -428,7 +430,7 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => UpdateChatSessionInput }) @Args({ name: 'options', type: () => UpdateChatSessionInput })
options: UpdateChatSessionInput options: UpdateChatSessionInput
) { ): Promise<string> {
const session = await this.chatSession.get(options.sessionId); const session = await this.chatSession.get(options.sessionId);
if (!session) { if (!session) {
throw new CopilotSessionNotFound(); throw new CopilotSessionNotFound();
@@ -442,7 +444,7 @@ export class CopilotResolver {
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${workspaceId}`; const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${workspaceId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
await this.chatSession.checkQuota(user.id); await this.chatSession.checkQuota(user.id);
@@ -460,12 +462,12 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => ForkChatSessionInput }) @Args({ name: 'options', type: () => ForkChatSessionInput })
options: ForkChatSessionInput options: ForkChatSessionInput
) { ): Promise<string> {
await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update'); await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update');
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`; const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
if (options.workspaceId === options.docId) { if (options.workspaceId === options.docId) {
@@ -489,15 +491,15 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => DeleteSessionInput }) @Args({ name: 'options', type: () => DeleteSessionInput })
options: DeleteSessionInput options: DeleteSessionInput
) { ): Promise<string[]> {
await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update'); await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update');
if (!options.sessionIds.length) { if (!options.sessionIds.length) {
return new NotFoundException('Session not found'); throw new NotFoundException('Session not found');
} }
const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`; const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
return await this.chatSession.cleanup({ return await this.chatSession.cleanup({
@@ -514,15 +516,15 @@ export class CopilotResolver {
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args({ name: 'options', type: () => CreateChatMessageInput }) @Args({ name: 'options', type: () => CreateChatMessageInput })
options: CreateChatMessageInput options: CreateChatMessageInput
) { ): Promise<string> {
const lockFlag = `${COPILOT_LOCKER}:message:${user?.id}:${options.sessionId}`; const lockFlag = `${COPILOT_LOCKER}:message:${user?.id}:${options.sessionId}`;
await using lock = await this.mutex.acquire(lockFlag); await using lock = await this.mutex.acquire(lockFlag);
if (!lock) { if (!lock) {
return new TooManyRequest('Server is busy'); throw new TooManyRequest('Server is busy');
} }
const session = await this.chatSession.get(options.sessionId); const session = await this.chatSession.get(options.sessionId);
if (!session || session.config.userId !== user.id) { if (!session || session.config.userId !== user.id) {
return new BadRequestException('Session not found'); throw new BadRequestException('Session not found');
} }
if (options.blobs) { if (options.blobs) {
@@ -564,7 +566,7 @@ export class UserCopilotResolver {
async copilot( async copilot(
@CurrentUser() user: CurrentUser, @CurrentUser() user: CurrentUser,
@Args('workspaceId', { nullable: true }) workspaceId?: string @Args('workspaceId', { nullable: true }) workspaceId?: string
) { ): Promise<CopilotType> {
if (workspaceId) { if (workspaceId) {
await this.ac await this.ac
.user(user.id) .user(user.id)

2
packages/backend/server/src/plugins/copilot/session.ts
View File

@@ -713,7 +713,7 @@ export class ChatSessionService {
}); });
} }
async createMessage(message: SubmittedMessage): Promise<string | undefined> { async createMessage(message: SubmittedMessage): Promise<string> {
return await this.messageCache.set(message); return await this.messageCache.set(message);
} }