diff --git a/apps/cli/src/ai/agent.ts b/apps/cli/src/ai/agent.ts index 2a7a3ca7..9ed95fe1 100644 --- a/apps/cli/src/ai/agent.ts +++ b/apps/cli/src/ai/agent.ts @@ -7,6 +7,7 @@ import { DEFAULT_MAX_TOOL_ROUNDS, + createLlmRouteDecider, decideRequestRoute, runAgentCore, checkAndCompress, @@ -166,14 +167,23 @@ export async function runServerAgent(options: RunAgentOptions): Promise { try { const routeStartedAt = Date.now() - const routeDecision = await decideRequestRoute({ - userMessage, - chatType, - locale, - dataSnapshot, - availableTools: tools.map((tool) => tool.name), - skillSummary: skillDef?.name ?? (skillMenu ? 'auto_skill_menu' : undefined), - }) + const routeDecision = await decideRequestRoute( + { + userMessage, + chatType, + locale, + dataSnapshot, + availableTools: tools.map((tool) => tool.name), + skillSummary: skillDef?.name ?? (skillMenu ? 'auto_skill_menu' : undefined), + }, + { + llmRouter: createLlmRouteDecider({ + piModel, + apiKey: llmConfig.apiKey, + abortSignal, + }), + } + ) aiLogger?.info('Router', 'Shadow route decision', { ...routeDecision, elapsedMs: Date.now() - routeStartedAt, diff --git a/apps/desktop/main/ai/agent/index.ts b/apps/desktop/main/ai/agent/index.ts index c0ffc89a..35e28887 100644 --- a/apps/desktop/main/ai/agent/index.ts +++ b/apps/desktop/main/ai/agent/index.ts @@ -11,6 +11,7 @@ import { aiLogger, isDebugMode } from '../logger' import { t as i18nT } from '../../i18n' import { DEFAULT_MAX_TOOL_ROUNDS, + createLlmRouteDecider, decideRequestRoute, runAgentCore, streamSimple, @@ -129,15 +130,24 @@ export class Agent { } const routeStartedAt = Date.now() - const routeDecision = await decideRequestRoute({ - userMessage, - chatType: this.chatType, - locale: this.locale, - dataSnapshot: this.context.dataSnapshot, - availableTools: piTools.map((tool) => tool.name), - assistantSummary: this.assistantConfig?.name, - skillSummary: this.skillCtx?.skillDef?.name ?? (this.skillCtx?.skillMenu ? 'auto_skill_menu' : undefined), - }) + const routeDecision = await decideRequestRoute( + { + userMessage, + chatType: this.chatType, + locale: this.locale, + dataSnapshot: this.context.dataSnapshot, + availableTools: piTools.map((tool) => tool.name), + assistantSummary: this.assistantConfig?.name, + skillSummary: this.skillCtx?.skillDef?.name ?? (this.skillCtx?.skillMenu ? 'auto_skill_menu' : undefined), + }, + { + llmRouter: createLlmRouteDecider({ + piModel: this.piModel, + apiKey: this.apiKey, + abortSignal: this.abortSignal, + }), + } + ) aiLogger.info('Router', 'Shadow route decision', { ...routeDecision, elapsedMs: Date.now() - routeStartedAt, diff --git a/packages/node-runtime/src/ai/agent/__tests__/router.test.ts b/packages/node-runtime/src/ai/agent/__tests__/router.test.ts index d9a1ec82..39b0a075 100644 --- a/packages/node-runtime/src/ai/agent/__tests__/router.test.ts +++ b/packages/node-runtime/src/ai/agent/__tests__/router.test.ts @@ -1,6 +1,6 @@ import { describe, it } from 'node:test' import assert from 'node:assert/strict' -import { decideRequestRoute } from '../router' +import { createLlmRouteDecider, decideRequestRoute } from '../router' import type { RouteDecision, RouteDecisionSource } from '../routing-types' const baseInput = { @@ -86,4 +86,64 @@ describe('decideRequestRoute', () => { assert.ok(decision.confidence < 0.6) assert.match(decision.reason, /ambiguous/i) }) + + it('creates an LLM fallback decider that parses route JSON', async () => { + const decider = createLlmRouteDecider({ + complete: async () => '{"route":"direct_response","confidence":1.4,"reason":"No local data dependency."}', + }) + + const decision = await decider( + { + ...baseInput, + userMessage: '帮我看一下这个情况。', + }, + { + route: 'tool_assisted', + confidence: 0.45, + reason: 'Ambiguous request.', + source: 'rule', + } + ) + + assert.equal(decision.route, 'direct_response') + assert.equal(decision.source, 'llm') + assert.equal(decision.confidence, 1) + assert.match(decision.reason, /No local data/) + }) + + it('keeps the rule decision when LLM fallback output is invalid', async () => { + const decider = createLlmRouteDecider({ + complete: async () => '{"route":"unknown","confidence":0.9,"reason":"bad"}', + }) + const ruleDecision: RouteDecision = { + route: 'tool_assisted', + confidence: 0.45, + reason: 'Ambiguous request.', + source: 'rule', + } + + const decision = await decider(baseInput, ruleDecision) + + assert.equal(decision.route, 'tool_assisted') + assert.equal(decision.source, 'rule') + assert.match(decision.reason, /invalid/i) + }) + + it('does not mark invalid LLM fallback as an LLM-sourced decision', async () => { + const decision = await decideRequestRoute( + { + ...baseInput, + userMessage: '帮我看一下这个情况。', + }, + { + llmRouter: createLlmRouteDecider({ + complete: async () => '{"route":"unknown","confidence":0.9,"reason":"bad"}', + }), + } + ) + + assert.equal(decision.route, 'tool_assisted') + assert.equal(decision.source, 'rule') + assert.match(decision.reason, /invalid/i) + }) }) diff --git a/packages/node-runtime/src/ai/agent/index.ts b/packages/node-runtime/src/ai/agent/index.ts index 888084cc..0aa59456 100644 --- a/packages/node-runtime/src/ai/agent/index.ts +++ b/packages/node-runtime/src/ai/agent/index.ts @@ -1,5 +1,5 @@ export type { AgentCoreOptions, AgentCoreEvent, AgentCoreResult, AgentTokenUsage, SimpleHistoryMessage } from './types' export { runAgentCore } from './core' export { DEFAULT_MAX_TOOL_ROUNDS } from './constants' -export { decideRequestRoute } from './router' +export { createLlmRouteDecider, decideRequestRoute } from './router' export type { LlmRouteDecider, RequestRoute, RouteDecision, RouteDecisionSource, RouterInput } from './routing-types' diff --git a/packages/node-runtime/src/ai/agent/router.ts b/packages/node-runtime/src/ai/agent/router.ts index 6241ce1c..c1013598 100644 --- a/packages/node-runtime/src/ai/agent/router.ts +++ b/packages/node-runtime/src/ai/agent/router.ts @@ -1,10 +1,32 @@ import type { LlmRouteDecider, RouteDecision, RouterInput } from './routing-types' +import { + completeSimple, + type Api as PiApi, + type Model as PiModel, + type TextContent as PiTextContent, +} from '@earendil-works/pi-ai' export interface DecideRequestRouteOptions { llmRouter?: LlmRouteDecider llmFallbackConfidenceThreshold?: number } +export type LlmRouterCompletionResult = + | string + | { + text: string + usage?: RouteDecision['usage'] + } + +export interface CreateLlmRouteDeciderOptions { + piModel?: PiModel + apiKey?: string + complete?: (prompt: string, signal?: AbortSignal) => Promise | LlmRouterCompletionResult + abortSignal?: AbortSignal + maxTokens?: number + temperature?: number +} + const DEFAULT_LLM_FALLBACK_CONFIDENCE_THRESHOLD = 0.6 function clampConfidence(value: number): number { @@ -20,6 +42,69 @@ function matchAny(text: string, patterns: RegExp[]): boolean { return patterns.some((pattern) => pattern.test(text)) } +function isRoute(value: unknown): value is RouteDecision['route'] { + return value === 'direct_response' || value === 'tool_assisted' || value === 'planned_execution' +} + +function resultToText(result: LlmRouterCompletionResult): { text: string; usage?: RouteDecision['usage'] } { + return typeof result === 'string' ? { text: result } : result +} + +function extractJsonObject(text: string): string { + const fenced = text.match(/```(?:json)?\s*([\s\S]*?)```/i) + if (fenced?.[1]) return fenced[1].trim() + const start = text.indexOf('{') + const end = text.lastIndexOf('}') + if (start >= 0 && end > start) return text.slice(start, end + 1) + return text.trim() +} + +function parseLlmDecision(rawText: string, usage?: RouteDecision['usage']): RouteDecision | null { + try { + const parsed = JSON.parse(extractJsonObject(rawText)) as Record + if (!isRoute(parsed.route)) return null + return { + route: parsed.route, + confidence: clampConfidence(typeof parsed.confidence === 'number' ? parsed.confidence : 0.5), + reason: typeof parsed.reason === 'string' && parsed.reason.trim() ? parsed.reason.trim() : 'LLM route decision.', + source: 'llm', + usage, + } + } catch { + return null + } +} + +function buildLlmRouterPrompt(input: RouterInput, ruleDecision: RouteDecision): string { + const toolNames = input.availableTools?.length ? input.availableTools.join(', ') : '(none)' + const dataSnapshot = input.dataSnapshot + ? `${input.dataSnapshot.name}, ${input.dataSnapshot.totalMessages} messages, ${input.dataSnapshot.totalMembers} members` + : '(none)' + + return `Classify the user's ChatLab request into exactly one route. + +Routes: +- direct_response: no local chat data/tools needed. +- tool_assisted: local chat data/tools needed, but the task is simple. +- planned_execution: complex analysis, multi-step evidence gathering, retries, long-range comparison, relationship analysis, or insufficient-evidence handling. + +Return strict JSON only: +{"route":"direct_response|tool_assisted|planned_execution","confidence":0.0,"reason":"short reason"} + +Context: +- locale: ${input.locale} +- chatType: ${input.chatType} +- availableTools: ${toolNames} +- dataSnapshot: ${dataSnapshot} +- assistantSummary: ${input.assistantSummary ?? '(none)'} +- skillSummary: ${input.skillSummary ?? '(none)'} +- recentIntentSummary: ${input.recentIntentSummary ?? '(none)'} +- ruleDecision: ${ruleDecision.route}, confidence=${ruleDecision.confidence}, reason=${ruleDecision.reason} + +User request: +${input.userMessage}` +} + function directRuleDecision(text: string): RouteDecision | null { if ( matchAny(text, [/不用查聊天记录|不用查询|不用查数据|不用看本地数据/i, /润色|改写|翻译|polish|rewrite|translate/i]) @@ -156,6 +241,67 @@ export async function decideRequestRoute( return { ...llmDecision, confidence: clampConfidence(llmDecision.confidence), - source: 'llm', + } +} + +export function createLlmRouteDecider(options: CreateLlmRouteDeciderOptions): LlmRouteDecider { + return async (input, ruleDecision) => { + const prompt = buildLlmRouterPrompt(input, ruleDecision) + + try { + const rawResult = options.complete + ? await options.complete(prompt, options.abortSignal) + : await completeWithPiAi(prompt, options) + const { text, usage } = resultToText(rawResult) + const decision = parseLlmDecision(text, usage) + if (decision) return decision + return { + ...ruleDecision, + reason: `${ruleDecision.reason} LLM fallback returned invalid route JSON.`, + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + return { + ...ruleDecision, + reason: `${ruleDecision.reason} LLM fallback failed: ${message}`, + } + } + } +} + +async function completeWithPiAi( + prompt: string, + options: CreateLlmRouteDeciderOptions +): Promise { + if (!options.piModel || !options.apiKey) { + throw new Error('LLM router requires piModel and apiKey') + } + + const result = await completeSimple( + options.piModel, + { + systemPrompt: 'You are a strict JSON classifier for ChatLab AI routing. Return JSON only.', + messages: [{ role: 'user', content: [{ type: 'text', text: prompt }], timestamp: Date.now() }] as any, + }, + { + apiKey: options.apiKey, + maxTokens: options.maxTokens ?? 180, + temperature: options.temperature ?? 0, + signal: options.abortSignal, + } + ) + + const text = result.content + .filter((item): item is PiTextContent => item.type === 'text') + .map((item) => item.text) + .join('') + + return { + text, + usage: { + promptTokens: result.usage?.input, + completionTokens: result.usage?.output, + totalTokens: result.usage?.totalTokens, + }, } } diff --git a/packages/node-runtime/src/ai/agent/routing-types.ts b/packages/node-runtime/src/ai/agent/routing-types.ts index 9ae383b4..765f1758 100644 --- a/packages/node-runtime/src/ai/agent/routing-types.ts +++ b/packages/node-runtime/src/ai/agent/routing-types.ts @@ -20,6 +20,11 @@ export interface RouteDecision { confidence: number reason: string source: RouteDecisionSource + usage?: { + promptTokens?: number + completionTokens?: number + totalTokens?: number + } } export type LlmRouteDecider = ( diff --git a/packages/node-runtime/src/ai/index.ts b/packages/node-runtime/src/ai/index.ts index 045af67b..d8a47185 100644 --- a/packages/node-runtime/src/ai/index.ts +++ b/packages/node-runtime/src/ai/index.ts @@ -97,7 +97,7 @@ export type { PreprocessingPipelineOptions, PreprocessingPipelineResult } from ' // Agent Core export type { AgentCoreOptions, AgentCoreEvent, AgentCoreResult, AgentTokenUsage, SimpleHistoryMessage } from './agent' -export { DEFAULT_MAX_TOOL_ROUNDS, decideRequestRoute, runAgentCore } from './agent' +export { DEFAULT_MAX_TOOL_ROUNDS, createLlmRouteDecider, decideRequestRoute, runAgentCore } from './agent' export type { LlmRouteDecider, RequestRoute, RouteDecision, RouteDecisionSource, RouterInput } from './agent' // Agent Event Handler diff --git a/packages/node-runtime/src/index.ts b/packages/node-runtime/src/index.ts index 3021020e..36bcd36d 100644 --- a/packages/node-runtime/src/index.ts +++ b/packages/node-runtime/src/index.ts @@ -203,7 +203,7 @@ export type { AIChat, AIMessage, AIMessageRole, ContentBlock, TokenUsageData, AI // Agent Core export type { AgentCoreOptions, AgentCoreEvent, AgentCoreResult, AgentTokenUsage, SimpleHistoryMessage } from './ai' -export { DEFAULT_MAX_TOOL_ROUNDS, decideRequestRoute, runAgentCore } from './ai' +export { DEFAULT_MAX_TOOL_ROUNDS, createLlmRouteDecider, decideRequestRoute, runAgentCore } from './ai' export type { LlmRouteDecider, RequestRoute, RouteDecision, RouteDecisionSource, RouterInput } from './ai' // Summary generation