@tanstack-ai-browser/apfel 0.0.1 → 0.1.0

package/src/adapter.ts

@@ -1,16 +1,17 @@
 /**
  * A TanStack AI text adapter for Apple's on-device Foundation Model, served
- * by apfel's OpenAI-compatible server (`apfel --serve --cors`). No backend
- * involved: the browser streams straight from localhost. Hand it to `chat()`
+ * by apfel's OpenAI-compatible server on the user's Mac. No backend
+ * involved: the browser streams straight from it. Hand it to `chat()`
  * like any hosted model's adapter: `chat()` runs tools and the agent loop,
  * and each `chatStream` call here is one completion from apfel, streamed
  * back as AG-UI events.
  */
-import { EventType, normalizeSystemPrompts } from "@tanstack/ai"
+import { buildBaseUsage, EventType, normalizeSystemPrompts } from "@tanstack/ai"
 import type {
   AdapterYieldChunk,
   ModelMessage,
   TextOptions,
+  TokenUsage,
   Tool
 } from "@tanstack/ai"
 import { BaseTextAdapter } from "@tanstack/ai/adapters"
@@ -24,46 +25,46 @@ import {
   createToolCallSplitter,
   parseJsonToolCalls
 } from "./tool-calls.js"
+import { matchToolName, repairToolArguments } from "./tool-repair.js"
+import type { GuidedRequest } from "./tool-repair.js"
 
-/** Where `apfel --serve` listens by default. */
+/** Where apfel's server listens by default. */
 export const APFEL_BASE_URL = "http://127.0.0.1:11434/v1"
-/** The model id apfel serves Apple's Foundation Model under. */
-export const APFEL_MODEL = "apple-foundationmodel"
+
+/** The models apfel serves: Apple's on-device Foundation Model. */
+export const APFEL_TEXT_MODELS = ["apple-foundationmodel"] as const
+
+/** A model apfel serves. */
+export type ApfelTextModel = (typeof APFEL_TEXT_MODELS)[number] | (string & {})
+
+export type ApfelTextConfig = {
+  /** apfel's OpenAI-compatible API, up to `/v1`. `APFEL_BASE_URL` by default. */
+  baseURL?: string
+  /** Headers sent with every request, such as apfel's bearer token. */
+  defaultHeaders?: Record<string, string>
+}
 
 /**
- * A model call's token counts and timings. apfel reports token counts only
- * in the usage at the end of the stream; counts it doesn't report are left
- * out.
+ * Timings of one model call, in milliseconds, in `RUN_FINISHED`'s
+ * `usage.providerUsageDetails`. The adapter measures them as the stream
+ * arrives; apfel itself reports token counts only.
  */
-export type ApfelStats = {
-  /** Tokens of the prompt. */
-  inputTokens?: number
-  /** Prompt tokens served from a cache, if apfel reports them. */
-  cachedTokens?: number
-  /** Tokens the model generated. */
-  outputTokens?: number
-  /** From sending the request to the first generated token, in ms. */
+export type ApfelProviderUsageDetails = {
+  /** From sending the request to the first generated text. */
   firstTokenMs?: number
-  /** Tokens generated after the first chunk (estimated, see below). */
-  decodeTokens?: number
-  /** From the first generated chunk to the last, in ms. */
+  /** From the first generated text to the last. */
   decodeMs?: number
-}
-
-export type ApfelOptions = {
-  /** apfel's OpenAI-compatible API, up to `/v1`. */
-  baseUrl?: string
-  /** The model id to ask apfel for. */
-  model?: string
   /**
-   * The call's stats at its first token (`done: false`, time to first token
-   * only) and when it ends (`done: true`). `callId` is unique per model
-   * call; `chat()` makes one per step of its agent loop.
+   * Tokens generated after the first chunk, estimated: apfel streams a few
+   * tokens per chunk, and the first chunk's share of the characters stands
+   * in for its share of the tokens.
+   */
+  decodeTokens?: number
+  /**
+   * Time spent rewriting the tool calls' arguments under their schemas,
+   * after the reply ended (see the README's "Tool calls").
    */
-  onStats?: (
-    stats: ApfelStats,
-    run: { threadId: string; runId: string; callId: string; done: boolean }
-  ) => void
+  toolRepairMs?: number
 }
 
 /**
@@ -283,27 +284,34 @@ async function errorBody(response: Response) {
   return raw.slice(0, 200)
 }
 
-function usageOf(usage: OpenAIUsage | null | undefined) {
+/** apfel's token counts as TanStack AI usage, with the call's timings. */
+function usageOf(
+  usage: OpenAIUsage | null | undefined,
+  details?: ApfelProviderUsageDetails
+): TokenUsage | undefined {
   if (
     usage?.prompt_tokens === undefined ||
     usage.completion_tokens === undefined
   ) {
     return undefined
   }
-  const cachedTokens = usage.prompt_tokens_details?.cached_tokens
-  return {
+  const result = buildBaseUsage<ApfelProviderUsageDetails>({
     promptTokens: usage.prompt_tokens,
     completionTokens: usage.completion_tokens,
-    totalTokens: usage.prompt_tokens + usage.completion_tokens,
-    ...(cachedTokens && { promptTokensDetails: { cachedTokens } })
+    totalTokens: usage.prompt_tokens + usage.completion_tokens
+  })
+  const cachedTokens = usage.prompt_tokens_details?.cached_tokens
+  if (cachedTokens) result.promptTokensDetails = { cachedTokens }
+  if (details && Object.keys(details).length > 0) {
+    result.providerUsageDetails = details
   }
+  return result
 }
 
 const errorText = (error: unknown) =>
   error instanceof Error ? error.message : String(error)
 
-const unreachable = (baseUrl: string) =>
-  `Couldn't reach Apple's on-device model at ${baseUrl}. Start it with \`apfel --serve --cors\`.`
+const unreachable = (baseURL: string) => `Couldn't reach apfel at ${baseURL}.`
 
 let ids = 0
 const nextId = (prefix: string) =>
@@ -311,8 +319,8 @@ const nextId = (prefix: string) =>
 
 type StreamedCall = { name: string; arguments: string }
 
-export class ApfelTextAdapter extends BaseTextAdapter<
-  string,
+export class ApfelTextAdapter<TModel extends string> extends BaseTextAdapter<
+  TModel,
   ApfelModelOptions,
   readonly ["text"],
   MetadataByModality,
@@ -321,15 +329,21 @@ export class ApfelTextAdapter extends BaseTextAdapter<
   never,
   ApfelReasoning
 > {
+  override readonly kind = "text" as const
   readonly name = "apfel" as const
-  override readonly api = "apfel" as const
-  private readonly baseUrl: string
-  private readonly onStats: ApfelOptions["onStats"]
-
-  constructor(options: ApfelOptions = {}) {
-    super({}, options.model ?? APFEL_MODEL)
-    this.baseUrl = (options.baseUrl ?? APFEL_BASE_URL).replace(/\/+$/, "")
-    this.onStats = options.onStats
+  override readonly provider = "apfel" as const
+  override readonly api = "openai-completions" as const
+  override readonly inputModalities = ["text"] as const
+  private readonly baseURL: string
+  private readonly headers: Record<string, string>
+
+  constructor(config: ApfelTextConfig, model: TModel) {
+    super({}, model)
+    this.baseURL = (config.baseURL ?? APFEL_BASE_URL).replace(/\/+$/, "")
+    this.headers = {
+      ...config.defaultHeaders,
+      "Content-Type": "application/json"
+    }
   }
 
   async *chatStream(
@@ -337,7 +351,6 @@ export class ApfelTextAdapter extends BaseTextAdapter<
   ): AsyncIterable<AdapterYieldChunk> {
     const runId = options.runId ?? nextId("run")
     const threadId = options.threadId
-    const callId = nextId("call")
     const messageId = nextId("msg")
     const reasoningId = nextId("reasoning")
     const model = this.model
@@ -367,7 +380,7 @@ export class ApfelTextAdapter extends BaseTextAdapter<
     const tools = toOpenAITools(options.tools)
     const toolNames = new Set(tools.map((tool) => tool.function.name))
     const send = options.wrapFetch ? options.wrapFetch(fetch) : fetch
-    const url = `${this.baseUrl}/chat/completions`
+    const url = `${this.baseURL}/chat/completions`
     logger.request(
       `activity=chat provider=apfel model=${model} messages=${options.messages.length} tools=${tools.length} stream=true`,
       { provider: "apfel", model }
@@ -377,7 +390,7 @@ export class ApfelTextAdapter extends BaseTextAdapter<
     try {
       response = await send(url, {
         method: "POST",
-        headers: { "Content-Type": "application/json" },
+        headers: this.headers,
         body: JSON.stringify({
           ...options.modelOptions,
           model,
@@ -401,7 +414,7 @@ export class ApfelTextAdapter extends BaseTextAdapter<
         error,
         source: "apfel.chatStream"
       })
-      yield runError(unreachable(this.baseUrl))
+      yield runError(unreachable(this.baseURL))
       return
     }
     if (!response.ok || !response.body) {
@@ -533,11 +546,6 @@ export class ApfelTextAdapter extends BaseTextAdapter<
           if (!timing.first) {
             timing.first = timing.last
             timing.firstChars = chars
-            // Token counts come only with the usage at the end.
-            this.onStats?.(
-              { firstTokenMs: timing.first - sentAt },
-              { threadId, runId, callId, done: false }
-            )
           }
           timing.chars += chars
         }
@@ -567,23 +575,17 @@ export class ApfelTextAdapter extends BaseTextAdapter<
     }
 
     const outputTokens = usage?.completion_tokens
-    this.onStats?.(
-      {
-        inputTokens: usage?.prompt_tokens,
-        cachedTokens: usage?.prompt_tokens_details?.cached_tokens,
-        outputTokens,
-        ...(timing.first && {
+    const details: ApfelProviderUsageDetails = timing.first
+      ? {
           firstTokenMs: timing.first - sentAt,
+          decodeMs: timing.last - timing.first,
           ...(outputTokens !== undefined && {
             decodeTokens:
               (outputTokens * (timing.chars - timing.firstChars)) /
-              timing.chars,
-            decodeMs: timing.last - timing.first
+              timing.chars
           })
-        })
-      },
-      { threadId, runId, callId, done: true }
-    )
+        }
+      : {}
 
     endReasoning()
     const { text: rest, markup } = splitter.flush()
@@ -604,11 +606,13 @@ export class ApfelTextAdapter extends BaseTextAdapter<
         name: call.name,
         arguments: JSON.stringify(call.args)
       }))
-    ]
+    ].map((call) => ({
+      ...call,
+      name: matchToolName(call.name, toolNames) ?? call.name
+    }))
+    let repairMs = 0
     for (const [index, call] of calls.entries()) {
       const toolCallId = nextId("tool")
-      const args = call.arguments || "{}"
-      const input = parseArguments(args)
       yield {
         type: EventType.TOOL_CALL_START,
         toolCallId,
@@ -618,6 +622,21 @@ export class ApfelTextAdapter extends BaseTextAdapter<
         model,
         timestamp: timestamp()
       }
+      const tool = options.tools?.find((t) => t.name === call.name)
+      const repairStart = performance.now()
+      const repaired = tool
+        ? await this.repairArguments(
+            tool,
+            call.arguments,
+            options,
+            send,
+            signal
+          )
+        : undefined
+      repairMs += performance.now() - repairStart
+      if (signal?.aborted) return
+      const args = repaired ? JSON.stringify(repaired) : call.arguments || "{}"
+      const input = repaired ?? parseArguments(args)
       yield {
         type: EventType.TOOL_CALL_ARGS,
         toolCallId,
@@ -640,7 +659,8 @@ export class ApfelTextAdapter extends BaseTextAdapter<
     endText()
     yield* events.splice(0)
 
-    const finishUsage = usageOf(usage)
+    if (repairMs > 0) details.toolRepairMs = repairMs
+    const finishUsage = usageOf(usage, details)
     yield {
       type: EventType.RUN_FINISHED,
       runId,
@@ -657,6 +677,75 @@ export class ApfelTextAdapter extends BaseTextAdapter<
     }
   }
 
+  /**
+   * A tool call's arguments, written again under the tool's schema with
+   * guided generation. Undefined if the tool takes no arguments or a request
+   * failed: the model's own arguments are used then.
+   */
+  private async repairArguments(
+    tool: Tool,
+    draft: string,
+    options: TextOptions<ApfelModelOptions>,
+    send: typeof fetch,
+    signal: AbortSignal | undefined
+  ) {
+    const system = normalizeSystemPrompts(options.systemPrompts)
+      .map((p) => p.content)
+      .join("\n\n")
+    const messages = toOpenAIMessages(options.messages)
+    const complete = async ({ instructions, schema, name }: GuidedRequest) => {
+      options.logger.request(
+        `activity=tool-repair provider=apfel model=${this.model} tool=${tool.name} step=${name}`,
+        { provider: "apfel", model: this.model }
+      )
+      const response = await send(`${this.baseURL}/chat/completions`, {
+        method: "POST",
+        headers: this.headers,
+        body: JSON.stringify({
+          ...options.modelOptions,
+          model: this.model,
+          // apfel takes one system message, and the conversation must end
+          // with the user or a tool result: the instructions go first.
+          messages: [
+            {
+              role: "system",
+              content: system
+                ? `${instructions}\n\nThe assistant's instructions, for context:\n${system}`
+                : instructions
+            },
+            ...messages
+          ],
+          response_format: {
+            type: "json_schema",
+            json_schema: { name, schema }
+          }
+        }),
+        signal
+      })
+      if (!response.ok) {
+        throw new Error(
+          `apfel returned HTTP ${response.status}: ${await errorBody(response)}`
+        )
+      }
+      const completion = (await response.json()) as OpenAICompletion
+      const text = (completion.choices?.[0]?.message?.content ?? "").trim()
+      return JSON.parse(
+        text.replace(/^```(?:json)?\s*|\s*```$/g, "")
+      ) as unknown
+    }
+    try {
+      return await repairToolArguments(tool, complete, draft)
+    } catch (error) {
+      if (!signal?.aborted) {
+        options.logger.errors("apfel tool-call repair failed", {
+          error,
+          source: "apfel.repairArguments"
+        })
+      }
+      return undefined
+    }
+  }
+
   /**
    * A JSON answer for `chat({ outputSchema })`, from apfel's JSON mode:
    * `response_format: json_schema` constrains Apple's model to the schema
@@ -674,9 +763,9 @@ export class ApfelTextAdapter extends BaseTextAdapter<
     )
     let response: Response
     try {
-      response = await send(`${this.baseUrl}/chat/completions`, {
+      response = await send(`${this.baseURL}/chat/completions`, {
         method: "POST",
-        headers: { "Content-Type": "application/json" },
+        headers: this.headers,
         body: JSON.stringify({
           ...chatOptions.modelOptions,
           model,
@@ -700,7 +789,7 @@ export class ApfelTextAdapter extends BaseTextAdapter<
         source: "apfel.structuredOutput"
       })
       if (chatOptions.abortController?.signal.aborted) throw error
-      throw new Error(unreachable(this.baseUrl))
+      throw new Error(unreachable(this.baseURL))
     }
     if (!response.ok) {
       throw new Error(
@@ -737,16 +826,23 @@ export class ApfelTextAdapter extends BaseTextAdapter<
 }
 
 /**
- * A TanStack AI adapter for Apple's on-device Foundation Model through
- * apfel (`apfel --serve --cors`).
+ * A TanStack AI text adapter for Apple's on-device Foundation Model, served
+ * by apfel.
  *
  * @example
  * ```ts
- * chat({ adapter: apfel(), messages, threadId })
+ * chat({
+ *   adapter: apfelText("apple-foundationmodel", { baseURL: APFEL_BASE_URL }),
+ *   messages,
+ *   threadId
+ * })
  * ```
  */
-export function apfel(options?: ApfelOptions) {
-  return new ApfelTextAdapter(options)
+export function apfelText<TModel extends ApfelTextModel>(
+  model: TModel,
+  config: ApfelTextConfig = {}
+): ApfelTextAdapter<TModel> {
+  return new ApfelTextAdapter(config, model)
 }
 
 /**
@@ -754,12 +850,13 @@ export function apfel(options?: ApfelOptions) {
  * lists it within 10 seconds.
  */
 export async function isApfelAvailable(
-  options: Pick<ApfelOptions, "baseUrl" | "model"> = {}
+  model: ApfelTextModel,
+  config: ApfelTextConfig = {}
 ) {
-  const baseUrl = (options.baseUrl ?? APFEL_BASE_URL).replace(/\/+$/, "")
-  const model = options.model ?? APFEL_MODEL
+  const baseURL = (config.baseURL ?? APFEL_BASE_URL).replace(/\/+$/, "")
   try {
-    const response = await fetch(`${baseUrl}/models`, {
+    const response = await fetch(`${baseURL}/models`, {
+      headers: { ...config.defaultHeaders },
       signal: AbortSignal.timeout(10_000)
     })
     if (!response.ok) return false

package/src/index.ts

@@ -3,9 +3,14 @@
  */
 export {
   APFEL_BASE_URL,
-  APFEL_MODEL,
+  APFEL_TEXT_MODELS,
   ApfelTextAdapter,
-  apfel,
+  apfelText,
   isApfelAvailable
 } from "./adapter.js"
-export type { ApfelModelOptions, ApfelOptions, ApfelStats } from "./adapter.js"
+export type {
+  ApfelModelOptions,
+  ApfelProviderUsageDetails,
+  ApfelTextConfig,
+  ApfelTextModel
+} from "./adapter.js"