Skip to content

Commit 244c2e1

Browse files
LucienShuiLucienShui
committed
feat(chat): Add open chat API endpoint with permissions and service
1 parent c0c2f5e commit 244c2e1

File tree

9 files changed

+308
-4
lines changed

9 files changed

+308
-4
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: 97 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,106 @@
11
import { TestClient } from 'test/test-client';
2-
import { APIKeyPermissionType } from './api-key.entity';
2+
import { APIKeyPermissionTarget, APIKeyPermissionType } from './api-key.entity';
33

4-
describe('APIKeyController (e2e)', () => {
4+
process.env.OBB_WIZARD_BASE_URL = 'http://localhost:8000';
5+
6+
// Mock the WizardAPIService to avoid needing the actual wizard service during tests
7+
jest.mock('../wizard/api.wizard.service', () => {
8+
return {
9+
WizardAPIService: jest.fn().mockImplementation(() => ({
10+
request: jest.fn().mockResolvedValue({ success: true }),
11+
proxy: jest.fn().mockResolvedValue({ success: true }),
12+
search: jest.fn().mockResolvedValue({ results: [] }),
13+
})),
14+
};
15+
});
16+
17+
// Mock fetch for any direct fetch calls
18+
global.fetch = jest.fn().mockResolvedValue({
19+
ok: true,
20+
json: jest.fn().mockResolvedValue({ success: true }),
21+
text: jest.fn().mockResolvedValue('success'),
22+
status: 200,
23+
body: {
24+
getReader: jest.fn().mockReturnValue({
25+
read: jest
26+
.fn()
27+
.mockResolvedValueOnce({
28+
done: false,
29+
value: new TextEncoder().encode(
30+
'data: {"response_type":"bos","role":"assistant","id":"test-msg-id"}\n',
31+
),
32+
})
33+
.mockResolvedValueOnce({
34+
done: false,
35+
value: new TextEncoder().encode(
36+
'data: {"response_type":"delta","message":{"content":"Hello"}}\n',
37+
),
38+
})
39+
.mockResolvedValueOnce({
40+
done: false,
41+
value: new TextEncoder().encode('data: {"response_type":"eos"}\n'),
42+
})
43+
.mockResolvedValueOnce({
44+
done: false,
45+
value: new TextEncoder().encode('data: {"response_type":"done"}\n'),
46+
})
47+
.mockResolvedValueOnce({ done: true, value: undefined }),
48+
cancel: jest.fn(),
49+
}),
50+
},
51+
}) as jest.MockedFunction<typeof fetch>;
52+
53+
describe('OpenAPIKeyController (e2e)', () => {
554
let client: TestClient;
55+
let chatApiKeyValue: string;
56+
let noChatApiKeyValue: string;
657

758
beforeAll(async () => {
859
client = await TestClient.create();
60+
61+
// Create an API key with CHAT CREATE permissions for testing
62+
const chatApiKeyData = {
63+
user_id: client.user.id,
64+
namespace_id: client.namespace.id,
65+
attrs: {
66+
root_resource_id: client.namespace.root_resource_id,
67+
permissions: [
68+
{
69+
target: APIKeyPermissionTarget.CHAT,
70+
permissions: [APIKeyPermissionType.CREATE],
71+
},
72+
],
73+
},
74+
};
75+
76+
const chatResponse = await client
77+
.post('/api/v1/api-keys')
78+
.send(chatApiKeyData)
79+
.expect(201);
80+
81+
chatApiKeyValue = chatResponse.body.value;
82+
83+
// Create an API key without CHAT permissions for negative testing
84+
const noChatApiKeyData = {
85+
user_id: client.user.id,
86+
namespace_id: client.namespace.id,
87+
attrs: {
88+
root_resource_id: client.namespace.root_resource_id,
89+
permissions: [
90+
{
91+
target: APIKeyPermissionTarget.RESOURCES,
92+
permissions: [APIKeyPermissionType.READ],
93+
},
94+
],
95+
},
96+
};
97+
98+
const noChatResponse = await client
99+
.post('/api/v1/api-keys')
100+
.send(noChatApiKeyData)
101+
.expect(201);
102+
103+
noChatApiKeyValue = noChatResponse.body.value;
9104
});
10105

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

‎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: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
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+
}
110+
}
111+
112+
return { messages };
113+
}
114+
}

‎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) => {
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)