Explorar o código

fix(llm-pi-ai): preserve harness adapter contract

Tianyi Cui hai 3 meses
pai
achega
922e2f913e

+ 7 - 1
packages/llm-deepseek/tests/adapter.e2e.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from 'vitest'
+import { afterEach, describe, expect, it } from 'vitest'
 import { Context } from 'cordis'
 import LlmService, { CallId } from '@deepseek-ai/dsh-llm'
 import type { GenerateResult, Message, ToolSchema } from '@deepseek-ai/dsh-llm'
@@ -13,14 +13,20 @@ import type { Config } from '@deepseek-ai/dsh-llm-deepseek'
 
 const FLASH = 'deepseek-v4-flash'
 const PRO = 'deepseek-v4-pro'
+const contexts: Context[] = []
 
 async function harness(model: string, config: Partial<Config> = {}) {
   const ctx = new Context()
+  contexts.push(ctx)
   await ctx.plugin(LlmService)
   await ctx.plugin(LlmDeepSeek, { models: [model], ...config })
   return ctx
 }
 
+afterEach(async () => {
+  await Promise.all(contexts.splice(0).map(ctx => ctx.fiber.dispose()))
+})
+
 function ask(text: string): Message[] {
   return [{ role: 'user', content: [{ type: 'text', text }] }]
 }

+ 2 - 2
packages/llm-pi-ai/README.md

@@ -6,10 +6,10 @@ DeepSeek adapter for the harness LLM seam backed by [`@earendil-works/pi-ai`](ht
 
 `@deepseek-ai/dsh-llm-deepseek` already talks to the same endpoint. This package is its **design-verification twin**: same models, same wire protocol, completely different internals — a unified LLM library with its own event vocabulary versus hand-rolled fetch/SSE. Anything the harness `StreamChunk` protocol cannot express for BOTH implementations is a core-vocabulary bug. The differences it exercised on purpose:
 
-- pi-ai hands back tool-call `arguments` as **parsed objects**; the harness keeps raw JSON strings (re-stringified at `block-end`).
+- pi-ai hands tool-call `arguments` around as **parsed objects**; the harness keeps raw JSON strings. The adapter patches replay payloads back to the original raw strings before sending them, and re-stringifies parsed output tool calls at `block-end`.
 - pi-ai reports failures as **in-stream error events** (it never throws mid-stream); these map to `finish {kind:'error'|'aborted'}` chunks — the protocol's other sanctioned error path besides throwing (which llm-deepseek uses).
 - pi-ai folds reasoning tokens into `usage.output`; there is no separate reasoning count to map.
-- pi-ai's options omit stop sequences; `GenerateOptions.stop` is injected via its `onPayload` hook.
+- pi-ai's options omit some DeepSeek/OpenAI-compatible details; the adapter uses its `onPayload` hook to preserve the harness contract (`stop`, per-tool `strict`, omitted reasoning effort, raw replayed tool arguments).
 
 ## Config
 

+ 66 - 12
packages/llm-pi-ai/src/adapter.ts

@@ -14,7 +14,7 @@
 import { stream as piStream } from '@earendil-works/pi-ai'
 import type { Model } from '@earendil-works/pi-ai'
 import { LlmAdapter, LlmError } from '@deepseek-ai/dsh-llm'
-import type { GenerateOptions, StreamChunk } from '@deepseek-ai/dsh-llm'
+import type { GenerateOptions, StreamChunk, ToolSchema } from '@deepseek-ai/dsh-llm'
 import { toPiContext, toStreamChunks } from './convert.ts'
 
 /** Reasoning levels surfaced by this adapter (DeepSeek wire: high|max). */
@@ -59,12 +59,71 @@ export function buildModel(modelId: string, options: PiAiAdapterOptions): Model<
   }
 }
 
+type Payload = {
+  tools?: { function?: { name?: unknown; strict?: unknown } }[]
+  messages?: {
+    role?: unknown
+    tool_calls?: { id?: unknown; function?: { arguments?: unknown } }[]
+  }[]
+  reasoning_effort?: unknown
+  stop?: unknown
+}
+
+function rawToolArguments(options: GenerateOptions): Map<string, string> {
+  const raw = new Map<string, string>()
+  for (const message of options.messages) {
+    if (message.role !== 'assistant') continue
+    for (const block of message.content) {
+      if (block.type === 'tool-call') raw.set(block.id, block.arguments)
+    }
+  }
+  return raw
+}
+
+function strictByToolName(tools: ToolSchema[] | undefined): Map<string, boolean | undefined> {
+  return new Map((tools ?? []).map(tool => [tool.name, tool.strict]))
+}
+
+function patchPayload(payload: unknown, options: GenerateOptions, reasoning: PiAiReasoning | undefined): unknown {
+  if (typeof payload !== 'object' || payload === null) return payload
+  const body = payload as Payload
+
+  if (reasoning === undefined) {
+    delete body.reasoning_effort
+  }
+  if (options.stop !== undefined) {
+    body.stop = options.stop
+  }
+
+  const strictByName = strictByToolName(options.tools)
+  for (const tool of body.tools ?? []) {
+    const name = tool.function?.name
+    if (typeof name !== 'string') continue
+    const strict = strictByName.get(name)
+    if (strict === undefined) delete tool.function?.strict
+    else if (tool.function !== undefined) tool.function.strict = strict
+  }
+
+  const rawById = rawToolArguments(options)
+  for (const message of body.messages ?? []) {
+    if (message.role !== 'assistant') continue
+    for (const call of message.tool_calls ?? []) {
+      if (typeof call.id !== 'string') continue
+      const raw = rawById.get(call.id)
+      if (raw !== undefined && call.function !== undefined) call.function.arguments = raw
+    }
+  }
+
+  return body
+}
+
 /**
  * pi-ai-backed adapter. One instance serves every registered model name.
  *
  * Implementation notes:
- * - `GenerateOptions.stop` is injected via pi-ai's `onPayload` hook (its
- *   public options omit stop sequences).
+ * - `onPayload` patches provider payload details pi-ai cannot express directly:
+ *   stop sequences, per-tool strict, omitted reasoning effort, and raw replayed
+ *   tool-call arguments.
  * - `prefill` throws UNSUPPORTED (same contract as dsh-llm-deepseek).
  * - pi-ai reports request failures as in-stream error events; convert.ts
  *   maps them to `finish {kind:'error'|'aborted'}` chunks rather than
@@ -86,8 +145,9 @@ export class PiAiAdapter extends LlmAdapter {
     const model = buildModel(options.model, this.options)
     // Undefined config means "provider default" (DeepSeek: thinking ENABLED),
     // matching llm-deepseek's omission semantics. pi-ai derives the wire
-    // thinking toggle from whether reasoningEffort is passed, so undefined
-    // maps to 'high' here; only an explicit 'off' disables thinking.
+    // thinking toggle from whether reasoningEffort is passed, so undefined maps
+    // internally to 'high' to get `thinking: enabled`; patchPayload then removes
+    // `reasoning_effort` so the provider chooses its default effort.
     const reasoning = this.options.reasoning ?? 'high'
 
     // pi-ai's event stream has no iterator-return cancellation hook: if our
@@ -106,13 +166,7 @@ export class PiAiAdapter extends LlmAdapter {
         ...options.maxTokens !== undefined ? { maxTokens: options.maxTokens } : {},
         signal: controller.signal,
         ...reasoning !== 'off' ? { reasoningEffort: reasoning } : {},
-        ...options.stop !== undefined ? {
-          // pi-ai's options omit stop sequences; inject them into the raw body.
-          onPayload: (payload: unknown) => {
-            (payload as Record<string, unknown>).stop = options.stop
-            return payload
-          },
-        } : {},
+        onPayload: payload => patchPayload(payload, options, this.options.reasoning),
         maxRetries: 0,
       })
 

+ 15 - 6
packages/llm-pi-ai/src/convert.ts

@@ -7,7 +7,8 @@
  * exists — an independent implementation stress-tests the StreamChunk
  * protocol):
  * - pi-ai tool-call `arguments` are PARSED OBJECTS; the harness keeps the
- *   raw JSON string. We parse on the way in and re-stringify on the way out.
+ *   raw JSON string. We parse on the way into pi-ai, patch provider payloads
+ *   back to the original raw string in the adapter, and re-stringify on output.
  * - pi-ai reports errors as in-stream `error` events (it never throws
  *   mid-stream); the harness expresses those as `finish {kind:'error'}` /
  *   `{kind:'aborted'}` chunks.
@@ -17,7 +18,7 @@
  * @module dsh-llm-pi-ai/convert
  */
 
-import { CallId } from '@deepseek-ai/dsh-llm'
+import { CallId, LlmError } from '@deepseek-ai/dsh-llm'
 import type { FinishReason, GenerateOptions, Message, StreamChunk, TokenUsage } from '@deepseek-ai/dsh-llm'
 import type {
   AssistantMessage,
@@ -168,6 +169,14 @@ 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'
+  if (/\b400\b|invalid.?request/i.test(message)) return 'INVALID_REQUEST'
+  if (/\b5\d\d\b/.test(message)) return 'SERVER'
+  return 'PI_AI_ERROR'
+}
+
 /** Map a terminal pi-ai event to the harness finish reason. */
 export function mapStopReason(message: AssistantMessage): FinishReason {
   switch (message.stopReason) {
@@ -175,10 +184,9 @@ export function mapStopReason(message: AssistantMessage): FinishReason {
     case 'length': return { kind: 'max-tokens' }
     case 'toolUse': return { kind: 'tool-calls' }
     case 'aborted': return { kind: 'aborted' }
-    case 'error': return {
-      kind: 'error',
-      message: message.errorMessage ?? 'pi-ai stream error',
-      code: 'PI_AI_ERROR',
+    case 'error': {
+      const text = message.errorMessage ?? 'pi-ai stream error'
+      return { kind: 'error', message: text, code: classifyPiAiError(text) }
     }
   }
 }
@@ -264,4 +272,5 @@ export async function* toStreamChunks(events: AsyncIterable<AssistantMessageEven
       // when one is added (switch covers all current variants).
     }
   }
+  throw new LlmError('pi-ai event stream ended without done/error', 'STREAM_CLOSED')
 }

+ 8 - 1
packages/llm-pi-ai/tests/adapter.e2e.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from 'vitest'
+import { afterEach, describe, expect, it } from 'vitest'
 import { Context } from 'cordis'
 import LlmService, { CallId } from '@deepseek-ai/dsh-llm'
 import type { GenerateResult, Message, ToolSchema } from '@deepseek-ai/dsh-llm'
@@ -15,14 +15,20 @@ import * as LlmDeepSeek from '@deepseek-ai/dsh-llm-deepseek'
 
 const FLASH = 'deepseek-v4-flash'
 const PRO = 'deepseek-v4-pro'
+const contexts: Context[] = []
 
 async function harness(model: string, config: Partial<Config> = {}) {
   const ctx = new Context()
+  contexts.push(ctx)
   await ctx.plugin(LlmService)
   await ctx.plugin(LlmPiAi, { models: [model], ...config })
   return ctx
 }
 
+afterEach(async () => {
+  await Promise.all(contexts.splice(0).map(ctx => ctx.fiber.dispose()))
+})
+
 function ask(text: string): Message[] {
   return [{ role: 'user', content: [{ type: 'text', text }] }]
 }
@@ -114,6 +120,7 @@ describe.skipIf(!process.env.DEEPSEEK_API_KEY)('llm-pi-ai e2e (real API)', () =>
     // same block KINDS in the same order for a deterministic prompt — the
     // cross-implementation check that the StreamChunk design holds.
     const deepseekCtx = new Context()
+    contexts.push(deepseekCtx)
     await deepseekCtx.plugin(LlmService)
     await deepseekCtx.plugin(LlmDeepSeek, { models: [FLASH], thinking: 'disabled' })
 

+ 52 - 5
packages/llm-pi-ai/tests/adapter.spec.ts

@@ -148,6 +148,44 @@ describe('PiAiAdapter against a mock server', () => {
     expect(server.requests[0]).toMatchObject({ stop: ['END'] })
   })
 
+  it('preserves per-tool strict exactly through onPayload', async () => {
+    const server = await mockServer([{ events: textEvents }])
+    const ctx = await harness(server.url)
+    await ctx.llm.generate({
+      model: 'deepseek-v4-flash',
+      messages: [],
+      tools: [
+        { name: 'strict_true', description: 'true', parameters: {}, strict: true },
+        { name: 'strict_false', description: 'false', parameters: {}, strict: false },
+        { name: 'strict_omitted', description: 'omitted', parameters: {} },
+      ],
+    })
+
+    const request = server.requests[0] as { tools: { function: { name: string; strict?: boolean } }[] }
+    expect(request.tools.map(tool => [tool.function.name, tool.function.strict])).toEqual([
+      ['strict_true', true],
+      ['strict_false', false],
+      ['strict_omitted', undefined],
+    ])
+    expect('strict' in request.tools[2]!.function).toBe(false)
+  })
+
+  it('preserves raw replayed tool-call arguments in the provider payload', async () => {
+    const server = await mockServer([{ events: textEvents }])
+    const ctx = await harness(server.url)
+    await ctx.llm.generate({
+      model: 'deepseek-v4-flash',
+      messages: [{
+        role: 'assistant',
+        content: [{ type: 'tool-call', id: CallId('broken'), name: 'f', arguments: '{broken' }],
+      }],
+    })
+
+    const request = server.requests[0] as { messages: { role: string; tool_calls?: { id: string; function: { arguments: string } }[] }[] }
+    const assistant = request.messages.find(message => message.role === 'assistant')
+    expect(assistant?.tool_calls?.[0]?.function.arguments).toBe('{broken')
+  })
+
   it('maps HTTP errors to error finish chunks (pi-ai in-stream style)', async () => {
     const server = await mockServer([{
       status: 401,
@@ -155,10 +193,20 @@ describe('PiAiAdapter against a mock server', () => {
     }])
     const ctx = await harness(server.url)
     const result = await ctx.llm.generate({ model: 'deepseek-v4-flash', messages: [] })
-    expect(result.finish.kind).toBe('error')
+    expect(result.finish).toMatchObject({ kind: 'error', code: 'AUTH' })
     expect((result.finish as { message: string }).message).toMatch(/bad key|401/)
   })
 
+  it.each([
+    [429, 'RATE_LIMIT'],
+    [500, 'SERVER'],
+  ] as const)('maps HTTP %s to stable error code %s', async (status, code) => {
+    const server = await mockServer([{ status, body: JSON.stringify({ error: { message: `provider ${status}` } }) }])
+    const ctx = await harness(server.url)
+    const result = await ctx.llm.generate({ model: 'deepseek-v4-flash', messages: [] })
+    expect(result.finish).toMatchObject({ kind: 'error', code })
+  })
+
   it('rejects prefill with UNSUPPORTED', async () => {
     const ctx = await harness('http://127.0.0.1:1')
     await expect(ctx.llm.generate({
@@ -263,10 +311,9 @@ describe('review fixes', () => {
     const server = await mockServer([{ events: textEvents }])
     const ctx = await harness(server.url) // no reasoning key at all
     await ctx.llm.generate({ model: 'deepseek-v4-flash', messages: [] })
-    expect(server.requests[0]).toMatchObject({
-      thinking: { type: 'enabled' },
-      reasoning_effort: 'high',
-    })
+    const request = server.requests[0] as Record<string, unknown>
+    expect(request.thinking).toEqual({ type: 'enabled' })
+    expect('reasoning_effort' in request).toBe(false)
   })
 
   it('replays reasoning_content on assistant tool-call turns (passback rule)', async () => {

+ 14 - 0
packages/llm-pi-ai/tests/convert.spec.ts

@@ -271,6 +271,11 @@ describe('toStreamChunks', () => {
     const chunks = await collect(toStreamChunks(feed({ type: 'error', reason: 'aborted', error })))
     expect(chunks.at(-1)).toEqual({ type: 'finish', reason: { kind: 'aborted' } })
   })
+
+  it('rejects a stream that ends without done or error', async () => {
+    await expect(collect(toStreamChunks(feed({ type: 'start', partial: assistant() }))))
+      .rejects.toThrow(/without done\/error/)
+  })
 })
 
 describe('mapStopReason / mapUsage', () => {
@@ -288,6 +293,15 @@ describe('mapStopReason / mapUsage', () => {
       .toEqual({ kind: 'error', message: 'pi-ai stream error', code: 'PI_AI_ERROR' })
   })
 
+  it('maps routable HTTP-ish error messages to stable codes', () => {
+    expect(mapStopReason(assistant({ stopReason: 'error', errorMessage: 'HTTP 401: bad key' })))
+      .toMatchObject({ kind: 'error', code: 'AUTH' })
+    expect(mapStopReason(assistant({ stopReason: 'error', errorMessage: 'HTTP 429: rate limit' })))
+      .toMatchObject({ kind: 'error', code: 'RATE_LIMIT' })
+    expect(mapStopReason(assistant({ stopReason: 'error', errorMessage: 'HTTP 500: backend down' })))
+      .toMatchObject({ kind: 'error', code: 'SERVER' })
+  })
+
   it('maps cache fields only when nonzero', () => {
     expect(mapUsage(usage(10, 5, 8, 2))).toEqual({
       inputTokens: 10,