301 lines
9.3 KiB
JavaScript
301 lines
9.3 KiB
JavaScript
|
const axios = require('axios');
|
|||
|
|
|
|||
|
|
class ArkChatService {
|
|||
|
|
constructor() {
|
|||
|
|
this.baseUrl = 'https://ark.cn-beijing.volces.com/api/v3';
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_getAuth() {
|
|||
|
|
return process.env.VOLC_ARK_API_KEY || process.env.VOLC_ACCESS_KEY_ID;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_isMockMode() {
|
|||
|
|
const ep = process.env.VOLC_ARK_ENDPOINT_ID;
|
|||
|
|
return !ep || ep === 'your_ark_endpoint_id';
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* 获取方舟知识库配置（如果已配置）
|
|||
|
|
* @returns {object|null} 知识库 metadata 配置
|
|||
|
|
*/
|
|||
|
|
_getKnowledgeBaseConfig() {
|
|||
|
|
const kbIds = process.env.VOLC_ARK_KNOWLEDGE_BASE_IDS;
|
|||
|
|
if (!kbIds || kbIds === 'your_knowledge_base_dataset_id') return null;
|
|||
|
|
|
|||
|
|
const datasetIds = kbIds.split(',').map(id => id.trim()).filter(Boolean);
|
|||
|
|
if (datasetIds.length === 0) return null;
|
|||
|
|
|
|||
|
|
const topK = parseInt(process.env.VOLC_ARK_KNOWLEDGE_TOP_K) || 3;
|
|||
|
|
const threshold = parseFloat(process.env.VOLC_ARK_KNOWLEDGE_THRESHOLD) || 0.5;
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
dataset_ids: datasetIds,
|
|||
|
|
top_k: topK,
|
|||
|
|
threshold: threshold,
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* 非流式调用方舟 LLM
|
|||
|
|
*/
|
|||
|
|
async chat(messages, tools = []) {
|
|||
|
|
if (this._isMockMode()) {
|
|||
|
|
console.warn('[ArkChat] EndPointId not configured, returning mock response');
|
|||
|
|
return this._mockChat(messages);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const body = {
|
|||
|
|
model: process.env.VOLC_ARK_ENDPOINT_ID,
|
|||
|
|
messages,
|
|||
|
|
stream: false,
|
|||
|
|
};
|
|||
|
|
if (tools.length > 0) body.tools = tools;
|
|||
|
|
|
|||
|
|
// 注入方舟私域知识库配置
|
|||
|
|
const kbConfig = this._getKnowledgeBaseConfig();
|
|||
|
|
if (kbConfig) {
|
|||
|
|
body.metadata = { knowledge_base: kbConfig };
|
|||
|
|
console.log('[ArkChat] Knowledge base enabled:', kbConfig.dataset_ids);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
try {
|
|||
|
|
const response = await axios.post(`${this.baseUrl}/chat/completions`, body, {
|
|||
|
|
headers: {
|
|||
|
|
'Content-Type': 'application/json',
|
|||
|
|
'Authorization': `Bearer ${this._getAuth()}`,
|
|||
|
|
},
|
|||
|
|
timeout: 30000,
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
const choice = response.data.choices?.[0];
|
|||
|
|
if (!choice) throw new Error('No response from Ark LLM');
|
|||
|
|
|
|||
|
|
const msg = choice.message;
|
|||
|
|
return {
|
|||
|
|
content: msg.content || '',
|
|||
|
|
toolCalls: msg.tool_calls || null,
|
|||
|
|
finishReason: choice.finish_reason,
|
|||
|
|
usage: response.data.usage,
|
|||
|
|
};
|
|||
|
|
} catch (error) {
|
|||
|
|
if (error.response) {
|
|||
|
|
console.error('[ArkChat] API error:', error.response.status, error.response.data);
|
|||
|
|
}
|
|||
|
|
throw error;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* SSE 流式调用方舟 LLM，通过回调逐块输出
|
|||
|
|
* @param {Array} messages
|
|||
|
|
* @param {Array} tools
|
|||
|
|
* @param {function} onChunk - (text: string) => void
|
|||
|
|
* @param {function} onToolCall - (toolCalls: Array) => void
|
|||
|
|
* @param {function} onDone - (fullContent: string) => void
|
|||
|
|
*/
|
|||
|
|
async chatStream(messages, tools = [], { onChunk, onToolCall, onDone }) {
|
|||
|
|
if (this._isMockMode()) {
|
|||
|
|
return this._mockChatStream(messages, { onChunk, onDone });
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const body = {
|
|||
|
|
model: process.env.VOLC_ARK_ENDPOINT_ID,
|
|||
|
|
messages,
|
|||
|
|
stream: true,
|
|||
|
|
};
|
|||
|
|
if (tools.length > 0) body.tools = tools;
|
|||
|
|
|
|||
|
|
// 注入方舟私域知识库配置
|
|||
|
|
const kbConfig = this._getKnowledgeBaseConfig();
|
|||
|
|
if (kbConfig) {
|
|||
|
|
body.metadata = { knowledge_base: kbConfig };
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const response = await axios.post(`${this.baseUrl}/chat/completions`, body, {
|
|||
|
|
headers: {
|
|||
|
|
'Content-Type': 'application/json',
|
|||
|
|
'Authorization': `Bearer ${this._getAuth()}`,
|
|||
|
|
},
|
|||
|
|
timeout: 60000,
|
|||
|
|
responseType: 'stream',
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
return new Promise((resolve, reject) => {
|
|||
|
|
let fullContent = '';
|
|||
|
|
let toolCalls = [];
|
|||
|
|
let buffer = '';
|
|||
|
|
|
|||
|
|
response.data.on('data', (chunk) => {
|
|||
|
|
buffer += chunk.toString();
|
|||
|
|
const lines = buffer.split('\n');
|
|||
|
|
buffer = lines.pop() || '';
|
|||
|
|
|
|||
|
|
for (const line of lines) {
|
|||
|
|
const trimmed = line.trim();
|
|||
|
|
if (!trimmed || !trimmed.startsWith('data: ')) continue;
|
|||
|
|
const data = trimmed.slice(6);
|
|||
|
|
if (data === '[DONE]') continue;
|
|||
|
|
|
|||
|
|
try {
|
|||
|
|
const parsed = JSON.parse(data);
|
|||
|
|
const delta = parsed.choices?.[0]?.delta;
|
|||
|
|
if (!delta) continue;
|
|||
|
|
|
|||
|
|
if (delta.content) {
|
|||
|
|
fullContent += delta.content;
|
|||
|
|
onChunk?.(delta.content);
|
|||
|
|
}
|
|||
|
|
if (delta.tool_calls) {
|
|||
|
|
for (const tc of delta.tool_calls) {
|
|||
|
|
if (!toolCalls[tc.index]) {
|
|||
|
|
toolCalls[tc.index] = { id: tc.id, type: tc.type, function: { name: '', arguments: '' } };
|
|||
|
|
}
|
|||
|
|
if (tc.function?.name) toolCalls[tc.index].function.name += tc.function.name;
|
|||
|
|
if (tc.function?.arguments) toolCalls[tc.index].function.arguments += tc.function.arguments;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
} catch (e) {
|
|||
|
|
// skip malformed lines
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
response.data.on('end', () => {
|
|||
|
|
if (toolCalls.length > 0) {
|
|||
|
|
onToolCall?.(toolCalls);
|
|||
|
|
}
|
|||
|
|
onDone?.(fullContent);
|
|||
|
|
resolve({ content: fullContent, toolCalls: toolCalls.length > 0 ? toolCalls : null });
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
response.data.on('error', reject);
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* 处理包含工具调用的完整对话循环（非流式）
|
|||
|
|
*/
|
|||
|
|
async chatWithTools(messages, tools, toolExecutor) {
|
|||
|
|
const result = await this.chat(messages, tools);
|
|||
|
|
|
|||
|
|
if (!result.toolCalls || result.toolCalls.length === 0) {
|
|||
|
|
return result;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const updatedMessages = [
|
|||
|
|
...messages,
|
|||
|
|
{ role: 'assistant', content: result.content, tool_calls: result.toolCalls },
|
|||
|
|
];
|
|||
|
|
|
|||
|
|
for (const tc of result.toolCalls) {
|
|||
|
|
const args = typeof tc.function.arguments === 'string'
|
|||
|
|
? JSON.parse(tc.function.arguments)
|
|||
|
|
: tc.function.arguments;
|
|||
|
|
|
|||
|
|
const toolResult = await toolExecutor.execute(tc.function.name, args, messages);
|
|||
|
|
|
|||
|
|
updatedMessages.push({
|
|||
|
|
role: 'tool',
|
|||
|
|
tool_call_id: tc.id,
|
|||
|
|
content: JSON.stringify(toolResult),
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const finalResult = await this.chat(updatedMessages, tools);
|
|||
|
|
return {
|
|||
|
|
...finalResult,
|
|||
|
|
messages: updatedMessages,
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* 流式版工具调用循环：先流式输出，如遇工具调用则执行后再流式输出最终结果
|
|||
|
|
*/
|
|||
|
|
async chatStreamWithTools(messages, tools, toolExecutor, { onChunk, onToolCall, onDone }) {
|
|||
|
|
const result = await this.chatStream(messages, tools, {
|
|||
|
|
onChunk,
|
|||
|
|
onToolCall,
|
|||
|
|
onDone: () => {}, // don't fire onDone yet
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
if (!result.toolCalls || result.toolCalls.length === 0) {
|
|||
|
|
onDone?.(result.content);
|
|||
|
|
return result;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 通知前端正在调用工具
|
|||
|
|
onToolCall?.(result.toolCalls);
|
|||
|
|
|
|||
|
|
const updatedMessages = [
|
|||
|
|
...messages,
|
|||
|
|
{ role: 'assistant', content: result.content, tool_calls: result.toolCalls },
|
|||
|
|
];
|
|||
|
|
|
|||
|
|
for (const tc of result.toolCalls) {
|
|||
|
|
const args = typeof tc.function.arguments === 'string'
|
|||
|
|
? JSON.parse(tc.function.arguments)
|
|||
|
|
: tc.function.arguments;
|
|||
|
|
|
|||
|
|
const toolResult = await toolExecutor.execute(tc.function.name, args, messages);
|
|||
|
|
updatedMessages.push({
|
|||
|
|
role: 'tool',
|
|||
|
|
tool_call_id: tc.id,
|
|||
|
|
content: JSON.stringify(toolResult),
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 工具执行完后，流式输出最终回答
|
|||
|
|
const finalResult = await this.chatStream(updatedMessages, tools, { onChunk, onToolCall: null, onDone });
|
|||
|
|
return {
|
|||
|
|
...finalResult,
|
|||
|
|
messages: updatedMessages,
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_mockChat(messages) {
|
|||
|
|
const lastMsg = messages[messages.length - 1];
|
|||
|
|
const userText = lastMsg?.content || '';
|
|||
|
|
console.log(`[ArkChat][MOCK] User: ${userText}`);
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
content: this._getMockReply(userText),
|
|||
|
|
toolCalls: null,
|
|||
|
|
finishReason: 'stop',
|
|||
|
|
usage: { prompt_tokens: 0, completion_tokens: 0 },
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
async _mockChatStream(messages, { onChunk, onDone }) {
|
|||
|
|
const lastMsg = messages[messages.length - 1];
|
|||
|
|
const userText = lastMsg?.content || '';
|
|||
|
|
console.log(`[ArkChat][MOCK-STREAM] User: ${userText}`);
|
|||
|
|
|
|||
|
|
const reply = this._getMockReply(userText);
|
|||
|
|
// 模拟逐字输出
|
|||
|
|
for (let i = 0; i < reply.length; i++) {
|
|||
|
|
onChunk?.(reply[i]);
|
|||
|
|
await new Promise((r) => setTimeout(r, 30));
|
|||
|
|
}
|
|||
|
|
onDone?.(reply);
|
|||
|
|
return { content: reply, toolCalls: null };
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_getMockReply(userText) {
|
|||
|
|
if (userText.includes('天气')) {
|
|||
|
|
return '根据 query_weather 工具查询，北京今天晴，气温 22°C，湿度 45%，北风3级。适合外出活动！';
|
|||
|
|
}
|
|||
|
|
if (userText.includes('订单')) {
|
|||
|
|
return '通过 query_order 工具查询，您的订单（ID: 12345）当前状态为：已发货，预计明天送达。快递单号：SF1234567890。';
|
|||
|
|
}
|
|||
|
|
if (userText.includes('你好') || userText.includes('嗨') || userText.includes('hi')) {
|
|||
|
|
return '你好！我是小智，很高兴为你服务。有什么我可以帮你的吗？';
|
|||
|
|
}
|
|||
|
|
if (userText.includes('知识') || userText.includes('退货') || userText.includes('政策')) {
|
|||
|
|
return '根据知识库查询，我们的退货政策如下：自签收之日起7天内可无理由退货，15天内可换货。请保持商品及包装完好。如需退货，请在"我的订单"中提交退货申请。';
|
|||
|
|
}
|
|||
|
|
return `收到你的消息："${userText}"。当前为模拟模式，配置方舟 LLM 凭证后将接入真实 AI 模型。你可以试试问我天气、订单、退货政策等问题。`;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
module.exports = new ArkChatService();
|