From fbcc0f299f8bfe4d0ab026426db196ff8db09aa1 Mon Sep 17 00:00:00 2001 From: digua Date: Sat, 22 Aug 2026 13:02:16 +0800 Subject: [PATCH] fix(ai): budget cross-chat entity results Keep full entity resolution details available to the application while bounding the model-visible payload with the runtime tokenizer. Preserve coverage and truncation metadata, prioritize usable scopes fairly, and cap ambiguity hints to avoid duplicate unbounded source lists. --- .../src/definitions/cross-chat-tools.test.ts | 105 +++++++++- .../tools/src/definitions/cross-chat-tools.ts | 194 +++++++++++++++++- 2 files changed, 297 insertions(+), 2 deletions(-) diff --git a/packages/tools/src/definitions/cross-chat-tools.test.ts b/packages/tools/src/definitions/cross-chat-tools.test.ts index 63045845..9768828a 100644 --- a/packages/tools/src/definitions/cross-chat-tools.test.ts +++ b/packages/tools/src/definitions/cross-chat-tools.test.ts @@ -199,7 +199,11 @@ describe('cross-chat agent registry', () => { displayName: '小红', platform: 'test', aliases: ['小红 A'], - sourceSessions: [{ id: 'private-1', name: '小红 A', type: ChatType.PRIVATE }], + sourceSessions: Array.from({ length: 5 }, (_, index) => ({ + id: `group-${index}`, + name: `小红所在群 ${index}`, + type: ChatType.GROUP, + })), }, { contactKey: 'test:xiaohong-2', @@ -226,12 +230,111 @@ describe('cross-chat agent registry', () => { } }, } as Partial) + context.maxToolResultTokens = 2_000 + context.countTokens = (text) => Math.ceil(text.length / 4) const tool = CROSS_CHAT_AGENT_TOOL_REGISTRY.find((item) => item.name === 'resolve_chat_entities') assert.ok(tool) const result = await tool.handler({ entities: [{ type: 'contact', displayName: '小红' }] }, context) assert.deepEqual(resolvedRefs, []) assert.equal((result.data as { contactLookups: Array<{ status: string }> }).contactLookups[0]?.status, 'ambiguous') + const modelData = JSON.parse(result.content) as { + contactLookups: Array<{ + candidates: Array<{ + sourceSessionCount: number + sourceSessionHints: unknown[] + sourceSessionHintsTruncated: boolean + }> + }> + } + assert.equal(modelData.contactLookups[0]?.candidates[0]?.sourceSessionCount, 5) + assert.equal(modelData.contactLookups[0]?.candidates[0]?.sourceSessionHints.length, 3) + assert.equal(modelData.contactLookups[0]?.candidates[0]?.sourceSessionHintsTruncated, true) + }) + + it('budgets model-visible entity scopes without dropping full tool details', async () => { + const resolvedSessions = Array.from({ length: 100 }, (_, index) => ({ + sessionId: `group-${index}`, + sessionName: `工作交流群 ${index}`, + sessionType: ChatType.GROUP, + platform: 'test' as const, + lastMessageTs: 1000 - index, + memberId: index + 1, + memberPlatformId: 'xiaohong', + memberName: '小红', + })) + const sourceSessions = resolvedSessions.map((session) => ({ + id: session.sessionId, + name: session.sessionName, + type: session.sessionType, + })) + const context = createContext({ + lookupContact: () => ({ + query: '小红', + status: 'resolved', + cacheStatus: 'fresh', + totalCandidates: 1, + candidates: [ + { + contactKey: 'test:xiaohong', + displayName: '小红', + platform: 'test', + aliases: [], + sourceSessions, + }, + ], + }), + resolveEntities: async () => ({ + contacts: [ + { + ref: { type: 'contact', contactKey: 'test:xiaohong', displayName: '小红' }, + status: 'resolved', + cacheStatus: 'fresh', + sessions: resolvedSessions, + unresolvedSessionIds: [], + failedSessionIds: [], + }, + ], + sessions: [], + unresolved: [], + coverage: { + requestedEntities: 1, + resolvedEntities: 1, + candidateSessions: resolvedSessions.length, + resolvedSessions: resolvedSessions.length, + failedSessions: 0, + }, + }), + }) + const countTokens = (text: string) => Math.ceil(text.length / 4) + context.maxToolResultTokens = 800 + context.countTokens = countTokens + const tool = CROSS_CHAT_AGENT_TOOL_REGISTRY.find((item) => item.name === 'resolve_chat_entities') + assert.ok(tool) + + const result = await tool.handler({ entities: [{ type: 'contact', displayName: '小红' }] }, context) + const modelData = JSON.parse(result.content) as { + contacts: Array<{ sessionCount: number; returnedSessions: number; sessions: unknown[] }> + coverage: { returnedSourceScopes: number; truncated: boolean; truncatedReasons: string[] } + contactLookups: Array<{ + candidates: Array<{ sourceSessionCount: number; sourceSessions?: unknown[] }> + }> + } + const details = result.data as { + contacts: Array<{ sessions: unknown[] }> + contactLookups: Array<{ candidates: Array<{ sourceSessions: unknown[] }> }> + } + + assert.ok(countTokens(result.content) <= 800) + assert.equal(modelData.coverage.truncated, true) + assert.ok(modelData.coverage.truncatedReasons.includes('tool_result_budget')) + assert.ok(modelData.coverage.returnedSourceScopes < resolvedSessions.length) + assert.equal(modelData.contacts[0]?.sessionCount, resolvedSessions.length) + assert.equal(modelData.contacts[0]?.returnedSessions, modelData.contacts[0]?.sessions.length) + assert.equal(modelData.contactLookups[0]?.candidates[0]?.sourceSessionCount, sourceSessions.length) + assert.equal(modelData.contactLookups[0]?.candidates[0]?.sourceSessions, undefined) + assert.equal(details.contacts[0]?.sessions.length, resolvedSessions.length) + assert.equal(details.contactLookups[0]?.candidates[0]?.sourceSessions.length, sourceSessions.length) }) it('sanitizes each session before returning global search evidence', async () => { diff --git a/packages/tools/src/definitions/cross-chat-tools.ts b/packages/tools/src/definitions/cross-chat-tools.ts index 37f09703..9ae821eb 100644 --- a/packages/tools/src/definitions/cross-chat-tools.ts +++ b/packages/tools/src/definitions/cross-chat-tools.ts @@ -1,8 +1,11 @@ import type { AIEntityRef, CrossChatContactLookupResult, + CrossChatEntityResolution, CrossChatEvidencePayload, CrossChatMessageSource, + CrossChatResolvedContactSession, + CrossChatResolvedSession, CrossChatSearchScope, } from '@openchatlab/shared-types' import type { CrossChatToolExecutionContext, JsonSchema, ToolDefinition, ToolResult } from '../types' @@ -126,7 +129,13 @@ async function resolveHandler( signal: context.abortSignal, }) const data = { ...resolution, contactLookups } - return { content: JSON.stringify(data), data } + const modelData = limitEntityResolutionToBudget( + resolution, + contactLookups, + context.maxToolResultTokens, + context.countTokens + ) + return { content: JSON.stringify(modelData), data } } async function searchHandler( @@ -395,6 +404,189 @@ function limitMessagesToBudget( return { messages: limited, truncated: limited.length < prepared.length } } +type EntityResolutionScopeItem = + | { kind: 'contact'; contactIndex: number; session: CrossChatResolvedContactSession } + | { kind: 'session'; sessionIndex: number; session: CrossChatResolvedSession } + +function limitEntityResolutionToBudget( + resolution: CrossChatEntityResolution, + contactLookups: CrossChatContactLookupResult[], + maxToolResultTokens: number | undefined, + injectedCountTokens?: (text: string) => number +): unknown { + const fullData = { ...resolution, contactLookups } + if (!maxToolResultTokens || maxToolResultTokens <= 0) return fullData + + const countTokens = injectedCountTokens ?? estimatePayloadTokens + const scopeItems = buildEntityResolutionScopePriority(resolution) + let includeLookupHints = true + let emptyPayload = buildEntityResolutionPayload(resolution, contactLookups, scopeItems, 0, includeLookupHints) + // 先舍弃只用于姓名消歧的来源提示，把预算优先留给后续工具真正需要的精确 scope。 + if (countTokens(JSON.stringify(emptyPayload)) > maxToolResultTokens) { + includeLookupHints = false + emptyPayload = buildEntityResolutionPayload(resolution, contactLookups, scopeItems, 0, includeLookupHints) + } + if (countTokens(JSON.stringify(emptyPayload)) > maxToolResultTokens) { + return buildMinimalEntityResolutionPayload(resolution, contactLookups, maxToolResultTokens, countTokens) + } + + let low = 0 + let high = scopeItems.length + // 最终 JSON 每次都用平台 tokenizer 复核；二分前缀避免大量来源时反复编码完整增长数组。 + while (low < high) { + const middle = Math.ceil((low + high) / 2) + const candidate = buildEntityResolutionPayload(resolution, contactLookups, scopeItems, middle, includeLookupHints) + if (countTokens(JSON.stringify(candidate)) <= maxToolResultTokens) low = middle + else high = middle - 1 + } + return buildEntityResolutionPayload(resolution, contactLookups, scopeItems, low, includeLookupHints) +} + +function buildEntityResolutionScopePriority(resolution: CrossChatEntityResolution): EntityResolutionScopeItem[] { + const items: EntityResolutionScopeItem[] = [] + // 用户显式选择的会话优先；联系人来源按最近时间轮询，避免第一个联系人耗尽全部预算。 + for (const [sessionIndex, session] of resolution.sessions.entries()) { + if (session.status === 'resolved' && session.session) items.push({ kind: 'session', sessionIndex, session }) + } + + const contactSessions = resolution.contacts.map((contact) => + [...contact.sessions].sort( + (left, right) => + (right.lastMessageTs ?? Number.NEGATIVE_INFINITY) - (left.lastMessageTs ?? Number.NEGATIVE_INFINITY) + ) + ) + for (let depth = 0; ; depth++) { + let added = false + for (const [contactIndex, sessions] of contactSessions.entries()) { + const session = sessions[depth] + if (!session) continue + items.push({ kind: 'contact', contactIndex, session }) + added = true + } + if (!added) break + } + return items +} + +function buildEntityResolutionPayload( + resolution: CrossChatEntityResolution, + contactLookups: CrossChatContactLookupResult[], + scopeItems: EntityResolutionScopeItem[], + selectedScopeCount: number, + includeLookupHints: boolean +): Record { + const selectedContactSessions = new Map() + const selectedSessionIndexes = new Set() + for (const item of scopeItems.slice(0, selectedScopeCount)) { + if (item.kind === 'session') { + selectedSessionIndexes.add(item.sessionIndex) + continue + } + const sessions = selectedContactSessions.get(item.contactIndex) ?? [] + sessions.push(item.session) + selectedContactSessions.set(item.contactIndex, sessions) + } + + const truncated = selectedScopeCount < scopeItems.length + const contacts = resolution.contacts.map((contact, index) => { + const sessions = selectedContactSessions.get(index) ?? [] + return { + ref: contact.ref, + status: contact.status, + cacheStatus: contact.cacheStatus, + sessionCount: contact.sessions.length, + returnedSessions: sessions.length, + sessions, + unresolvedSessionCount: contact.unresolvedSessionIds.length, + failedSessionCount: contact.failedSessionIds.length, + } + }) + const sessions = resolution.sessions.filter( + (session, index) => session.status === 'unresolved' || selectedSessionIndexes.has(index) + ) + + return { + contacts, + sessions, + unresolved: resolution.unresolved, + coverage: { + ...resolution.coverage, + returnedContactEntities: contacts.length, + returnedSessionEntities: sessions.length, + returnedSourceScopes: selectedScopeCount, + truncated, + truncatedReasons: truncated ? ['tool_result_budget'] : [], + }, + contactLookups: buildModelContactLookups(contactLookups, includeLookupHints), + } +} + +function buildModelContactLookups( + contactLookups: CrossChatContactLookupResult[], + includeHints: boolean +): Array> { + return contactLookups.map((lookup) => ({ + query: lookup.query, + status: lookup.status, + cacheStatus: lookup.cacheStatus, + totalCandidates: lookup.totalCandidates, + candidates: lookup.candidates.map((candidate) => { + const hints = includeHints && lookup.status === 'ambiguous' ? candidate.sourceSessions.slice(0, 3) : [] + return { + contactKey: candidate.contactKey, + displayName: candidate.displayName, + platform: candidate.platform, + aliases: candidate.aliases.slice(0, 8), + aliasCount: candidate.aliases.length, + sourceSessionCount: candidate.sourceSessions.length, + sourceSessionHintsReturned: hints.length, + sourceSessionHintsTruncated: lookup.status === 'ambiguous' && hints.length < candidate.sourceSessions.length, + ...(hints.length > 0 + ? { + sourceSessionHints: hints, + } + : {}), + } + }), + })) +} + +function buildMinimalEntityResolutionPayload( + resolution: CrossChatEntityResolution, + contactLookups: CrossChatContactLookupResult[], + maxToolResultTokens: number, + countTokens: (text: string) => number +): Record { + const coverage = { + ...resolution.coverage, + returnedContactEntities: 0, + returnedSessionEntities: 0, + returnedSourceScopes: 0, + truncated: true, + truncatedReasons: ['tool_result_budget'], + } + const candidates: Array> = [ + { + contacts: [], + sessions: [], + unresolved: [], + coverage, + contactLookups: contactLookups.map((lookup) => ({ + query: lookup.query, + status: lookup.status, + totalCandidates: lookup.totalCandidates, + })), + }, + { coverage }, + { truncated: true, truncatedReasons: ['tool_result_budget'] }, + ] + return ( + candidates.find((candidate) => countTokens(JSON.stringify(candidate)) <= maxToolResultTokens) ?? { + truncated: true, + } + ) +} + function estimatePayloadTokens(text: string): number { // Browser/MCP 等未注入 Node tokenizer 的调用方仍需避免按 ASCII 密度低估中文结果。 let cjkChars = 0