Agent/clients/omlx-client.ts
2026-08-06 22:17:39 +08:00

316 lines
11 KiB
TypeScript

// @ts-nocheck
import { omlxUrl, zhongtaiOmlxProxyUrl, loadOmlxConfig, loadThinkingEnabled } from '../utils/config.js';
import { reportTokens, reportStreamDelta, reportStreamEnd, reportPromptCache } from '../utils/metrics-client.js';
import { AGENT_IDENTITY_CORE } from '../core/agent-identity-core.js';
import { DOMAIN_RULES } from '../core/domain-rules.js';
import { scanToolCallBuffer, finalizeToolCallBuffer, cleanToolCallArtifacts } from '../utils/tool-call-parse.js';
import { createThinkSplitter, stripThinkTags } from '../core/think-parse.js';
import { friendlyLocalFetchError } from '../utils/model-manager.js';
const IDENTITY = `
${AGENT_IDENTITY_CORE}
你是"ccSparkle Agent",地铁乘务专业的智能助手,基于本地部署的 MLX 模型(Apple Silicon 原生优化),通过 MCP 协议操作本地地铁运营图 Electron 应用。
【回答风格 - 严格遵守】
- **只用简体中文**回复用户,禁止整句英文(见上方【语言】规则)。
- 用自然口语化的中文,专业、简洁、有温度。
- **禁止使用任何 emoji 表情符号**,全部用纯文字表达。
- 不要机械复述身份描述。
【工具调用规则】
- **仅当本轮 messages 里提供了 tools、且用户有明确操作指令时**才可调用工具。
- 自我介绍、向评委讲解、寒暄、概念说明:**禁止**输出 \`<tool_call>\`,直接写中文正文。
- 有工具可调时格式: \`<tool_call>{"name":"metro_xxx","arguments":{...}}</tool_call>\`
- metro_load_timetable_full 的 arguments **只有** query(字符串,如"工作日"),不要臆造 dayType 等 schema 里没有的字段
- 工具失败时如实告知错误。
- 多候选时列给用户选,不自己代决定。
${DOMAIN_RULES}
`;
export async function chat({ messages, tools, signal, onTextDelta, onThinkingDelta, maxTokens: maxTokensOpt, temperature }) {
const cfg = loadOmlxConfig();
if (!cfg.apiKey) {
throw new Error('未配置 oMLX API Key,请点配置按钮(配)设置');
}
const thinkOn = loadThinkingEnabled();
const leadIdx = messages.findIndex(m => m.role !== 'system');
const leadEnd = leadIdx === -1 ? messages.length : leadIdx;
const systemExtra = messages.slice(0, leadEnd)
.map(m => m.content).filter(Boolean).join('\n\n');
const omlxMessages = [{
role: 'system',
content: systemExtra ? `${IDENTITY}\n\n${systemExtra}` : IDENTITY,
}];
for (const m of messages.slice(leadEnd)) {
if (m.role === 'assistant' && m.toolCalls && m.toolCalls.length > 0) {
omlxMessages.push({
role: 'assistant',
content: m.content || null,
tool_calls: m.toolCalls.map(tc => ({
id: tc.id,
type: 'function',
function: {
name: tc.name,
arguments: typeof tc.arguments === 'string'
? tc.arguments
: JSON.stringify(tc.arguments || {}),
},
})),
});
} else if (m.role === 'tool') {
omlxMessages.push({
role: 'tool',
tool_call_id: m.toolCallId,
content: m.content,
});
} else {
let content = m.content || '';
if (m.role === 'user' && content) {
if (!thinkOn && !content.includes('/no_think') && !content.includes('/think')) {
content = `${content} /no_think`;
} else if (thinkOn && !content.includes('/think') && !content.includes('/no_think')) {
content = `${content} /think`;
}
}
omlxMessages.push({ role: m.role, content });
}
}
const payload = {
model: cfg.model,
messages: omlxMessages,
stream: true,
stream_options: { include_usage: true },
chat_template_kwargs: { enable_thinking: !!thinkOn },
enable_thinking: !!thinkOn,
};
if (maxTokensOpt != null) {
payload.max_tokens = maxTokensOpt;
}
if (temperature != null) {
payload.temperature = temperature;
}
if (tools && tools.length > 0) {
payload.tools = tools.map(t => ({
type: 'function',
function: {
name: t.name,
description: t.description || '',
parameters: t.inputSchema || { type: 'object', properties: {} },
},
}));
}
let r;
try {
r = await fetch(zhongtaiOmlxProxyUrl(), {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'X-Cloud-Url': omlxUrl(),
'X-Cloud-Token': cfg.apiKey,
},
body: JSON.stringify(payload),
signal,
});
} catch (e) {
throw friendlyLocalFetchError('omlx', e);
}
if (!r.ok) {
const text = await r.text();
if (r.status === 502 || /ECONNREFUSED|fetch failed|connect/i.test(text)) {
throw friendlyLocalFetchError('omlx', new Error(text.slice(0, 120)));
}
throw new Error(`oMLX HTTP ${r.status}: ${text.slice(0, 300)}`);
}
return consumeStream(r.body, { onTextDelta, onThinkingDelta, thinkOn });
}
let _warmedUp = false;
export async function warmup() {
if (_warmedUp) return true;
const cfg = loadOmlxConfig();
if (!cfg.apiKey) return false;
try {
const r = await fetch(omlxUrl(), {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Authorization': `Bearer ${cfg.apiKey}`,
},
body: JSON.stringify({
model: cfg.model,
messages: [
{ role: 'system', content: IDENTITY },
{ role: 'user', content: 'ping' },
],
stream: false,
max_tokens: 1,
temperature: 0,
}),
});
if (r.ok) {
_warmedUp = true;
return true;
}
return false;
} catch (e) {
return false;
}
}
export async function listModels() {
const cfg = loadOmlxConfig();
if (!cfg.apiKey) return [];
try {
const r = await fetch(`${cfg.baseUrl.replace(/\/+$/, '')}/v1/models`, {
headers: { 'Authorization': `Bearer ${cfg.apiKey}` },
});
if (!r.ok) return [];
const data = await r.json();
return (data.data || []).map(m => m.id);
} catch (e) {
return [];
}
}
async function consumeStream(body, { onTextDelta, onThinkingDelta, thinkOn }) {
const reader = body.getReader();
const decoder = new TextDecoder();
let lineBuf = '';
let textBuf = '';
let fullText = '';
let fullThinking = '';
let toolCalls = [];
let evalCount = 0;
const t0 = performance.now();
const thinkSplit = createThinkSplitter(thinkOn);
let rawSample = '';
let rawLen = 0;
let sawReasoningField = false;
let reasoningFieldLen = 0;
const emitContent = (raw) => {
if (!raw) return;
rawLen += raw.length;
if (rawSample.length < 400) rawSample += raw.slice(0, 400 - rawSample.length);
const split = thinkSplit.push(raw);
if (split.thinking) {
fullThinking += split.thinking;
if (thinkOn) onThinkingDelta?.(split.thinking);
}
if (split.content) {
textBuf += split.content;
fullText += split.content;
const result = scanToolCallBuffer(textBuf);
if (result.safeText && onTextDelta) onTextDelta(result.safeText);
textBuf = result.remainder;
if (result.toolCalls.length > 0) toolCalls.push(...result.toolCalls);
reportStreamDelta(split.content.length);
}
};
while (true) {
const { done, value } = await reader.read();
if (done) break;
lineBuf += decoder.decode(value, { stream: true });
const lines = lineBuf.split('\n');
lineBuf = lines.pop();
for (const line of lines) {
if (!line.startsWith('data:')) continue;
const data = line.slice(5).trim();
if (!data || data === '[DONE]') continue;
let chunk;
try {
chunk = JSON.parse(data);
} catch (e) {
continue;
}
const choice = chunk.choices?.[0];
const delta = choice?.delta;
const thinkPiece = delta?.reasoning_content || delta?.reasoning || '';
if (thinkPiece) {
sawReasoningField = true;
reasoningFieldLen += thinkPiece.length;
fullThinking += thinkPiece;
if (thinkOn) onThinkingDelta?.(thinkPiece);
reportStreamDelta(thinkPiece.length);
}
if (delta?.content) emitContent(delta.content);
if (delta?.tool_calls) {
for (const tc of delta.tool_calls) {
if (tc.function?.name) {
toolCalls.push({
id: tc.id || `toolu_openai_${Date.now()}_${toolCalls.length}`,
name: tc.function.name,
arguments: typeof tc.function.arguments === 'string'
? safeParse(tc.function.arguments)
: (tc.function.arguments || {}),
});
}
}
}
if (chunk.usage) {
evalCount = chunk.usage.completion_tokens || evalCount;
reportPromptCache(chunk.usage);
}
}
}
const tail = thinkSplit.flush();
if (tail.thinking) {
fullThinking += tail.thinking;
if (thinkOn) onThinkingDelta?.(tail.thinking);
}
if (tail.content) {
textBuf += tail.content;
fullText += tail.content;
}
const final = finalizeToolCallBuffer(textBuf);
if (final.safeText && onTextDelta) onTextDelta(final.safeText);
if (final.toolCalls.length > 0) toolCalls.push(...final.toolCalls);
const cleanedText = stripThinkTags(cleanToolCallArtifacts(fullText));
if (evalCount > 0) {
const durationMs = performance.now() - t0;
reportTokens(evalCount, durationMs);
reportStreamEnd(evalCount);
}
return { text: cleanedText, thinking: fullThinking, toolCalls, evalCount };
}
function safeParse(s) {
try {
let cleaned = s.replace(/\{\{/g, '{').replace(/\}\}/g, '}');
return JSON.parse(cleaned);
} catch (e) {
return { _raw: s };
}
}