Skip to content

Commit 5ed4f0c

Browse files
LucienShuiLucienShui
authored
Merge pull request #204 from import-ai/feat/chat_api
feat(chat): Add open chat API endpoint with permissions and service
2 parents c0c2f5e + 5d63e55 commit 5ed4f0c

File tree

9 files changed

+215
-3
lines changed

9 files changed

+215
-3
lines changed

‎src/api-key/api-key.entity.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ export enum APIKeyPermissionType {
1010

1111
export enum APIKeyPermissionTarget {
1212
RESOURCES = 'resources',
13+
CHAT = 'chat',
1314
}
1415

1516
export type APIKeyPermission = {

‎src/api-key/open.api-key.e2e-spec.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { TestClient } from 'test/test-client';
22
import { APIKeyPermissionType } from './api-key.entity';
33

4-
describe('APIKeyController (e2e)', () => {
4+
describe('OpenAPIKeyController (e2e)', () => {
55
let client: TestClient;
66

77
beforeAll(async () => {

‎src/applications/apps/wechat-bot.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,10 @@ export class WechatBot extends BaseApp {
112112
APIKeyPermissionType.READ,
113113
],
114114
},
115+
{
116+
target: APIKeyPermissionTarget.CHAT,
117+
permissions: [APIKeyPermissionType.CREATE],
118+
},
115119
],
116120
},
117121
});

‎src/messages/messages.service.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,12 @@ export class MessagesService {
138138
});
139139
}
140140

141+
async findOne(id: string) {
142+
return await this.messageRepository.findOneOrFail({
143+
where: { id },
144+
});
145+
}
146+
141147
async remove(conversationId: string, messageId: string, user: User) {
142148
return await this.messageRepository.softDelete({
143149
id: messageId,

‎src/wizard/dto/open-agent-request.dto.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
import { IsBoolean, IsNotEmpty, IsOptional, IsString } from 'class-validator';
2+
import {
3+
PrivateSearchToolDto,
4+
WebSearchToolDto,
5+
} from 'omniboxd/wizard/dto/agent-request.dto';
6+
7+
export class OpenAgentRequestDto {
8+
@IsString()
9+
@IsNotEmpty()
10+
query: string;
11+
12+
tools: Array<PrivateSearchToolDto | WebSearchToolDto>;
13+
14+
@IsOptional()
15+
@IsBoolean()
16+
enable_thinking?: boolean;
17+
18+
@IsOptional()
19+
@IsString()
20+
lang?: '简体中文' | 'English';
21+
22+
@IsOptional()
23+
@IsString()
24+
parent_message_id?: string;
25+
}

‎src/wizard/open.wizard.controller.ts‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
UploadedFile,
66
UseInterceptors,
77
} from '@nestjs/common';
8+
import { RequestId } from 'omniboxd/decorators/request-id.decorators';
89
import { WizardService } from 'omniboxd/wizard/wizard.service';
910
import { CompressedCollectRequestDto } from 'omniboxd/wizard/dto/collect-request.dto';
1011
import { CollectResponseDto } from 'omniboxd/wizard/dto/collect-response.dto';
@@ -16,11 +17,16 @@ import {
1617
APIKeyPermissionType,
1718
} from 'omniboxd/api-key/api-key.entity';
1819
import { OpenCollectRequestDto } from 'omniboxd/wizard/dto/open-collect-request.dto';
20+
import { OpenAgentRequestDto } from 'omniboxd/wizard/dto/open-agent-request.dto';
1921
import { FileInterceptor } from '@nestjs/platform-express';
22+
import { OpenWizardService } from 'omniboxd/wizard/open.wizard.service';
2023

2124
@Controller('open/api/v1/wizard')
2225
export class OpenWizardController {
23-
constructor(private readonly wizardService: WizardService) {}
26+
constructor(
27+
private readonly wizardService: WizardService,
28+
private readonly openWizardService: OpenWizardService,
29+
) {}
2430

2531
@Post('collect')
2632
@APIKeyAuth({
@@ -49,4 +55,27 @@ export class OpenWizardController {
4955
compressedHtml,
5056
);
5157
}
58+
59+
@Post('ask')
60+
@APIKeyAuth({
61+
permissions: [
62+
{
63+
target: APIKeyPermissionTarget.CHAT,
64+
permissions: [APIKeyPermissionType.CREATE],
65+
},
66+
],
67+
})
68+
async ask(
69+
@APIKey() apiKey: APIKeyEntity,
70+
@UserId() userId: string,
71+
@RequestId() requestId: string,
72+
@Body() data: OpenAgentRequestDto,
73+
): Promise<any> {
74+
return await this.openWizardService.ask(
75+
userId,
76+
apiKey.namespaceId,
77+
requestId,
78+
data,
79+
);
80+
}
5281
}

‎src/wizard/open.wizard.service.ts‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import { Injectable } from '@nestjs/common';
2+
import { MessagesService } from 'omniboxd/messages/messages.service';
3+
import { ConversationsService } from 'omniboxd/conversations/conversations.service';
4+
import { WizardService } from 'omniboxd/wizard/wizard.service';
5+
import { User } from 'omniboxd/user/entities/user.entity';
6+
import { OpenAgentRequestDto } from 'omniboxd/wizard/dto/open-agent-request.dto';
7+
import { AgentRequestDto } from 'omniboxd/wizard/dto/agent-request.dto';
8+
import { ChatResponse } from 'omniboxd/wizard/dto/chat-response.dto';
9+
10+
@Injectable()
11+
export class OpenWizardService {
12+
constructor(
13+
private readonly wizardService: WizardService,
14+
private readonly messagesService: MessagesService,
15+
private readonly conversationsService: ConversationsService,
16+
) {}
17+
18+
async ask(
19+
userId: string,
20+
namespaceId: string,
21+
requestId: string,
22+
data: OpenAgentRequestDto,
23+
): Promise<any> {
24+
const conversationId = await this.resolveConversationId(
25+
userId,
26+
namespaceId,
27+
data.parent_message_id,
28+
);
29+
30+
const agentRequest: AgentRequestDto = {
31+
...data,
32+
conversation_id: conversationId,
33+
namespace_id: namespaceId,
34+
enable_thinking: data.enable_thinking ?? false,
35+
};
36+
37+
const chunks: ChatResponse[] = await this.wizardService.streamService.chat(
38+
userId,
39+
agentRequest,
40+
requestId,
41+
'ask',
42+
);
43+
44+
return this.mergeChunks(chunks);
45+
}
46+
47+
private async resolveConversationId(
48+
userId: string,
49+
namespaceId: string,
50+
parentMessageId?: string,
51+
): Promise<string> {
52+
if (parentMessageId) {
53+
// Find conversation_id from parent_message_id
54+
const parentMessage = await this.messagesService.findOne(parentMessageId);
55+
return parentMessage.conversationId;
56+
} else {
57+
// Create a new conversation
58+
const user = { id: userId } as User;
59+
const conversation = await this.conversationsService.create(
60+
namespaceId,
61+
user,
62+
);
63+
return conversation.id;
64+
}
65+
}
66+
67+
private mergeChunks(chunks: ChatResponse[]): any {
68+
const messages: any[] = [];
69+
let currentMessage: any = null;
70+
71+
for (const chunk of chunks) {
72+
if (chunk.response_type === 'bos') {
73+
const bosChunk = chunk;
74+
currentMessage = {
75+
id: bosChunk.id,
76+
role: bosChunk.role,
77+
parent_id: bosChunk.parentId,
78+
message: {
79+
role: bosChunk.role,
80+
},
81+
attrs: {},
82+
};
83+
} else if (chunk.response_type === 'delta' && currentMessage) {
84+
const deltaChunk = chunk;
85+
if (deltaChunk.message.content) {
86+
currentMessage.message.content =
87+
(currentMessage.message.content || '') + deltaChunk.message.content;
88+
}
89+
if (deltaChunk.message.reasoning_content) {
90+
currentMessage.message.reasoning_content =
91+
(currentMessage.message.reasoning_content || '') +
92+
deltaChunk.message.reasoning_content;
93+
}
94+
if (deltaChunk.message.tool_calls) {
95+
currentMessage.message.tool_calls = deltaChunk.message.tool_calls;
96+
}
97+
if (deltaChunk.message.tool_call_id) {
98+
currentMessage.message.tool_call_id = deltaChunk.message.tool_call_id;
99+
}
100+
if (deltaChunk.attrs) {
101+
currentMessage.attrs = {
102+
...currentMessage.attrs,
103+
...deltaChunk.attrs,
104+
};
105+
}
106+
} else if (chunk.response_type === 'eos' && currentMessage) {
107+
messages.push(currentMessage);
108+
currentMessage = null;
109+
} else if (chunk.response_type === 'done') {
110+
// Done
111+
} else {
112+
throw new Error(`Invalid response_type = ${chunk.response_type}`);
113+
}
114+
}
115+
116+
return { messages };
117+
}
118+
}

‎src/wizard/stream.service.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -323,4 +323,30 @@ export class StreamService {
323323
);
324324
}
325325
}
326+
327+
async chat(
328+
userId: string,
329+
body: AgentRequestDto,
330+
requestId: string,
331+
mode: 'ask' | 'write' = 'ask',
332+
): Promise<any> {
333+
const observable = await this.agentStream(userId, body, requestId, mode);
334+
335+
const chunks: ChatResponse[] = [];
336+
337+
return new Promise((resolve, reject) => {
338+
observable.subscribe({
339+
next: (event: MessageEvent) => {
340+
const chunk: ChatResponse = JSON.parse(event.data as string);
341+
chunks.push(chunk);
342+
},
343+
complete: () => {
344+
resolve(chunks);
345+
},
346+
error: (error: Error) => {
347+
reject(error);
348+
},
349+
});
350+
});
351+
}
326352
}

‎src/wizard/wizard.module.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,18 @@ import { AttachmentsModule } from 'omniboxd/attachments/attachments.module';
1414
import { TasksModule } from 'omniboxd/tasks/tasks.module';
1515
import { MinioModule } from 'omniboxd/minio/minio.module';
1616
import { OpenWizardController } from 'omniboxd/wizard/open.wizard.controller';
17+
import { ConversationsModule } from 'omniboxd/conversations/conversations.module';
18+
import { OpenWizardService } from 'omniboxd/wizard/open.wizard.service';
1719

1820
@Module({
19-
providers: [WizardService, ChunkManagerService],
21+
providers: [WizardService, ChunkManagerService, OpenWizardService],
2022
imports: [
2123
UserModule,
2224
NamespacesModule,
2325
NamespaceResourcesModule,
2426
TagModule,
2527
MessagesModule,
28+
ConversationsModule,
2629
AttachmentsModule,
2730
TasksModule,
2831
MinioModule,

0 commit comments

Comments
 (0)