fix(llm-pi-ai): classify usage-based context overflow
Pass each resolved catalog model capacity into pi-ai stream conversion so the upstream full-message classifier can recognize provider-specific, silent, and length-stop overflow signals. Retain the harness text fallback for legacy provider wording and cover the catalog-resolution path with a mock-provider regression.
This commit is contained in:
@@ -43,7 +43,7 @@ If a listener rewrites assembled assistant content, the loop drops replay state
|
||||
## Vocabulary differences
|
||||
|
||||
- pi-ai tool-call arguments are parsed objects; the harness stores raw JSON strings. The adapter parses input and re-stringifies output.
|
||||
- pi-ai reports failures as in-stream error events; these map to `finish {kind:'error'|'aborted'}` chunks, with recognized context overflow normalized to `CONTEXT_WINDOW_EXCEEDED`.
|
||||
- pi-ai reports failures as in-stream error events; these map to `finish {kind:'error'|'aborted'}` chunks. Provider-specific error text and usage signals evaluated against the resolved model's context window normalize overflow to `CONTEXT_WINDOW_EXCEEDED`.
|
||||
- pi-ai folds reasoning tokens into output usage; there is no separate reasoning count to map.
|
||||
- `GenerateOptions.stop` is rejected with `UNSUPPORTED_OPTION` because pi-ai's common streaming surface cannot guarantee it across providers.
|
||||
|
||||
|
||||
@@ -115,7 +115,7 @@ export class PiAiAdapter extends LlmAdapter {
|
||||
// Harness-owned and therefore win collisions.
|
||||
headers: requestHeaders(profile.headers),
|
||||
})
|
||||
yield* toStreamChunks(events)
|
||||
yield* toStreamChunks(events, model.contextWindow)
|
||||
} finally {
|
||||
options.signal?.removeEventListener('abort', onCallerAbort)
|
||||
controller.abort('consumer stopped streaming')
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
|
||||
import { CallId, CONTEXT_WINDOW_EXCEEDED_CODE, isContextWindowExceededError, LlmError } from '@deepseek-ai/dsh-llm'
|
||||
import type { FinishReason, StreamChunk, TokenUsage } from '@deepseek-ai/dsh-llm'
|
||||
import { isContextOverflow } from '@earendil-works/pi-ai'
|
||||
import type { AssistantMessage, AssistantMessageEvent, Usage as PiUsage } from '@earendil-works/pi-ai'
|
||||
import { toPiReplayState } from './replay.ts'
|
||||
|
||||
@@ -30,10 +31,6 @@ export function mapUsage(usage: PiUsage): TokenUsage {
|
||||
function classifyPiAiError(message: string): string {
|
||||
if (/\b(?:401|403)\b/.test(message)) return 'AUTH'
|
||||
if (/\b429\b|rate.?limit/i.test(message)) return 'RATE_LIMIT'
|
||||
// TODO: Classify the full message with pi-ai's isContextOverflow() and the
|
||||
// resolved model's contextWindow so provider-specific and usage-based overflows
|
||||
// reach automatic compaction.
|
||||
if (isContextWindowExceededError(message)) return CONTEXT_WINDOW_EXCEEDED_CODE
|
||||
if (/\b400\b|invalid.?request/i.test(message)) return 'INVALID_REQUEST'
|
||||
if (/\b5\d\d\b/.test(message)) return 'SERVER'
|
||||
return 'PI_AI_ERROR'
|
||||
@@ -42,9 +39,22 @@ function classifyPiAiError(message: string): string {
|
||||
/**
|
||||
* Map a terminal pi-ai event to the harness finish reason.
|
||||
* @param message - the assistant message carried by the `done` or `error` event.
|
||||
* @param contextWindow - resolved catalog capacity for usage-based overflow detection.
|
||||
* @returns the harness reason; `error` yields `{kind: 'error'}` with a code classified from the error text.
|
||||
*/
|
||||
export function mapStopReason(message: AssistantMessage): FinishReason {
|
||||
export function mapStopReason(message: AssistantMessage, contextWindow?: number): FinishReason {
|
||||
const piAiOverflow = isContextOverflow(message, contextWindow)
|
||||
const harnessOverflow = message.stopReason === 'error'
|
||||
&& message.errorMessage !== undefined
|
||||
&& isContextWindowExceededError(message.errorMessage)
|
||||
if (piAiOverflow || harnessOverflow) {
|
||||
return {
|
||||
kind: 'error',
|
||||
message: message.errorMessage ?? `pi-ai detected context overflow for model "${message.model}"`,
|
||||
code: CONTEXT_WINDOW_EXCEEDED_CODE,
|
||||
}
|
||||
}
|
||||
|
||||
switch (message.stopReason) {
|
||||
case 'stop': return { kind: 'stop' }
|
||||
case 'length': return { kind: 'max-tokens' }
|
||||
@@ -62,10 +72,14 @@ export function mapStopReason(message: AssistantMessage): FinishReason {
|
||||
* mid-stream — failures arrive as `error` events, which become error/aborted
|
||||
* `finish` chunks (the harness protocol's other error-delivery style).
|
||||
* @param events - one assistant turn's pi-ai event stream.
|
||||
* @param contextWindow - resolved catalog capacity for usage-based overflow detection.
|
||||
* @returns the harness chunks, ending with `usage` then `finish`; throws
|
||||
* `LlmError` (`STREAM_CLOSED`) if the source ends without a terminal event.
|
||||
*/
|
||||
export async function* toStreamChunks(events: AsyncIterable<AssistantMessageEvent>): AsyncGenerator<StreamChunk> {
|
||||
export async function* toStreamChunks(
|
||||
events: AsyncIterable<AssistantMessageEvent>,
|
||||
contextWindow?: number,
|
||||
): AsyncGenerator<StreamChunk> {
|
||||
// pi-ai contentIndex ↔ our block index map 1:1 (both count blocks from 0
|
||||
// in stream order), but we track ids per index for tool calls.
|
||||
const toolIds = new Map<number, { id: string; name: string }>()
|
||||
@@ -128,13 +142,17 @@ export async function* toStreamChunks(events: AsyncIterable<AssistantMessageEven
|
||||
break
|
||||
case 'done':
|
||||
yield { type: 'usage', usage: mapUsage(event.message.usage) }
|
||||
yield { type: 'finish', reason: mapStopReason(event.message), replayState: toPiReplayState(event.message) }
|
||||
yield {
|
||||
type: 'finish',
|
||||
reason: mapStopReason(event.message, contextWindow),
|
||||
replayState: toPiReplayState(event.message),
|
||||
}
|
||||
return
|
||||
case 'error':
|
||||
// In-stream error delivery (pi-ai's style) → error finish chunk
|
||||
// (the harness's other sanctioned error path besides throwing).
|
||||
yield { type: 'usage', usage: mapUsage(event.error.usage) }
|
||||
yield { type: 'finish', reason: mapStopReason(event.error) }
|
||||
yield { type: 'finish', reason: mapStopReason(event.error, contextWindow) }
|
||||
return
|
||||
// no default: AssistantMessageEvent is pi-ai's closed union; a new
|
||||
// event type should fail compilation here via tsc's exhaustiveness
|
||||
|
||||
@@ -2,9 +2,10 @@ import { createServer } from 'node:http'
|
||||
import type { IncomingMessage, Server, ServerResponse } from 'node:http'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import { Context } from 'cordis'
|
||||
import LlmService, { LlmError, userAgent } from '@deepseek-ai/dsh-llm'
|
||||
import LlmService, { CONTEXT_WINDOW_EXCEEDED_CODE, LlmError, userAgent } from '@deepseek-ai/dsh-llm'
|
||||
import * as LlmPiAi from '@deepseek-ai/dsh-llm-pi-ai'
|
||||
import { PiAiAdapter } from '@deepseek-ai/dsh-llm-pi-ai'
|
||||
import { getModels } from '@earendil-works/pi-ai'
|
||||
import { resolveProfiles } from '../src/config.ts'
|
||||
import { assemble } from './assemble.ts'
|
||||
|
||||
@@ -178,6 +179,29 @@ describe('PiAiAdapter provider routing', () => {
|
||||
const result = await assemble(ctx, { model: 'deepseek-v4-flash', messages: [] })
|
||||
expect(result.finish).toMatchObject({ kind: 'error', code })
|
||||
})
|
||||
|
||||
it('uses the resolved catalog context window for usage-based overflow detection', async () => {
|
||||
const model = getModels('deepseek').find(candidate => candidate.id === 'deepseek-v4-flash')
|
||||
if (model === undefined) throw new Error('deepseek-v4-flash missing from pi-ai test catalog')
|
||||
const events = [
|
||||
'{"choices":[{"delta":{"role":"assistant","content":""},"index":0,"finish_reason":null}]}',
|
||||
JSON.stringify({
|
||||
choices: [{ delta: {}, index: 0, finish_reason: 'stop' }],
|
||||
usage: { prompt_tokens: model.contextWindow + 1, completion_tokens: 0 },
|
||||
}),
|
||||
'[DONE]',
|
||||
]
|
||||
const server = await mockServer([{ events }])
|
||||
const ctx = await harness(server.url)
|
||||
|
||||
const result = await assemble(ctx, { model: model.id, messages: [] })
|
||||
|
||||
expect(result.finish).toEqual({
|
||||
kind: 'error',
|
||||
message: `pi-ai detected context overflow for model "${model.id}"`,
|
||||
code: CONTEXT_WINDOW_EXCEEDED_CODE,
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('provider profile lifecycle', () => {
|
||||
|
||||
@@ -537,6 +537,34 @@ describe('mapStopReason / mapUsage', () => {
|
||||
}))).toMatchObject({ kind: 'error', code: 'INVALID_REQUEST' })
|
||||
})
|
||||
|
||||
it('uses pi-ai provider-specific overflow classification without losing rate-limit exclusions', () => {
|
||||
expect(mapStopReason(assistant({
|
||||
stopReason: 'error',
|
||||
errorMessage: 'prompt is too long: 213462 tokens > 200000 maximum',
|
||||
}))).toMatchObject({ kind: 'error', code: CONTEXT_WINDOW_EXCEEDED_CODE })
|
||||
expect(mapStopReason(assistant({
|
||||
stopReason: 'error',
|
||||
errorMessage: 'ThrottlingException: Too many tokens, rate limit reached',
|
||||
}))).toMatchObject({ kind: 'error', code: 'RATE_LIMIT' })
|
||||
})
|
||||
|
||||
it('uses the resolved context window for silent and length-stop overflows', () => {
|
||||
const silent = assistant({ stopReason: 'stop', usage: usage(101, 0) })
|
||||
expect(mapStopReason(silent)).toEqual({ kind: 'stop' })
|
||||
expect(mapStopReason(silent, 100)).toEqual({
|
||||
kind: 'error',
|
||||
message: 'pi-ai detected context overflow for model "deepseek-v4-flash"',
|
||||
code: CONTEXT_WINDOW_EXCEEDED_CODE,
|
||||
})
|
||||
|
||||
const truncated = assistant({ stopReason: 'length', usage: usage(80, 0, 19) })
|
||||
expect(mapStopReason(truncated)).toEqual({ kind: 'max-tokens' })
|
||||
expect(mapStopReason(truncated, 100)).toMatchObject({
|
||||
kind: 'error',
|
||||
code: CONTEXT_WINDOW_EXCEEDED_CODE,
|
||||
})
|
||||
})
|
||||
|
||||
it('maps cache fields only when nonzero', () => {
|
||||
expect(mapUsage(usage(10, 5, 8, 2))).toEqual({
|
||||
inputTokens: 10,
|
||||
|
||||
Reference in New Issue
Block a user