spb/polyllm

1import type { ContentPart, ToolCallPart, UnifiedChatResponse, UnifiedStreamEvent, Usage, ProviderId, FinishReason, Citation } from "./types";23/**4 * Accumulates a unified stream into a full response. Used by adapters to implement5 * `chat()` on top of `streamChat()` and by the server to persist assistant messages.6 */7export class StreamAccumulator {8  text = "";9  reasoning = "";10  reasoningSignature?: string;11  reasoningProviderData?: unknown;12  toolCalls = new Map<string, { name: string; args: string; final?: Record<string, unknown>; providerData?: unknown }>();13  citations: Citation[] = [];14  usage?: Usage;15  finishReason: FinishReason = "other";16  providerData: Record<string, unknown> = {};17  id?: string;18  model?: string;19  error?: UnifiedStreamEvent & { type: "error" };20  firstTokenAt?: number;2122  constructor(private readonly startedAt = Date.now()) {}2324  push(ev: UnifiedStreamEvent) {25    switch (ev.type) {26      case "start":27        this.id = ev.id ?? this.id;28        this.model = ev.model ?? this.model;29        break;30      case "text-delta":31        if (!this.firstTokenAt && ev.text) this.firstTokenAt = Date.now();32        this.text += ev.text;33        break;34      case "reasoning-delta":35        if (!this.firstTokenAt && ev.text) this.firstTokenAt = Date.now();36        this.reasoning += ev.text;37        break;38      case "reasoning-signature":39        this.reasoningSignature = ev.signature;40        this.reasoningProviderData = ev.providerData;41        break;42      case "tool-start":43        this.toolCalls.set(ev.id, { name: ev.name, args: "" });44        break;45      case "tool-delta": {46        const tc = this.toolCalls.get(ev.id);47        if (tc) tc.args += ev.argumentsDelta;48        break;49      }50      case "tool-end": {51        const prev = this.toolCalls.get(ev.id);52        // Adapters that stream deltas emit a partial tool-end (empty argumentsText) — keep the accumulated text.53        const args = ev.argumentsText || prev?.args || "";54        this.toolCalls.set(ev.id, { name: ev.name || prev?.name || "", args, final: Object.keys(ev.arguments ?? {}).length ? ev.arguments : undefined, providerData: ev.providerData ?? prev?.providerData });55        break;56      }57      case "citation":58        this.citations.push(ev.citation);59        break;60      case "usage":61        this.usage = { ...(this.usage ?? {}), ...ev.usage } as Usage;62        break;63      case "provider-data":64        Object.assign(this.providerData, ev.data);65        break;66      case "finish":67        this.finishReason = ev.reason;68        break;69      case "error":70        this.error = ev;71        this.finishReason = "error";72        break;73    }74  }7576  get ttftMs(): number | undefined {77    return this.firstTokenAt ? this.firstTokenAt - this.startedAt : undefined;78  }7980  toolCallParts(): ToolCallPart[] {81    return [...this.toolCalls.entries()].map(([id, tc]) => ({82      type: "tool-call",83      id,84      name: tc.name,85      arguments: tc.final ?? safeJson(tc.args),86      argumentsText: tc.args,87      providerData: tc.providerData,88    }));89  }9091  contentParts(): ContentPart[] {92    const parts: ContentPart[] = [];93    if (this.reasoning) parts.push({ type: "reasoning", text: this.reasoning, signature: this.reasoningSignature, providerData: this.reasoningProviderData });94    if (this.text) parts.push({ type: "text", text: this.text });95    parts.push(...this.toolCallParts());96    return parts;97  }9899  toResponse(provider: ProviderId, model: string): UnifiedChatResponse {100    return {101      id: this.id,102      provider,103      model: this.model ?? model,104      content: this.contentParts(),105      text: this.text,106      reasoning: this.reasoning || undefined,107      toolCalls: this.toolCallParts(),108      citations: this.citations.length ? this.citations : undefined,109      usage: this.usage,110      finishReason: this.finishReason,111      latencyMs: Date.now() - this.startedAt,112      providerData: Object.keys(this.providerData).length ? this.providerData : undefined,113    };114  }115}116117export function safeJson(text: string): Record<string, unknown> {118  if (!text || !text.trim()) return {};119  try {120    const v = JSON.parse(text);121    return v && typeof v === "object" ? (v as Record<string, unknown>) : { value: v };122  } catch {123    return { _raw: text };124  }125}126127/** Implements `chat()` from `streamChat()` for adapters (no duplicated request-building code). */128export async function collectStream(provider: ProviderId, model: string, stream: AsyncIterable<UnifiedStreamEvent>): Promise<UnifiedChatResponse> {129  const acc = new StreamAccumulator();130  for await (const ev of stream) {131    acc.push(ev);132    if (ev.type === "error") {133      const { PolyProviderError } = await import("./types");134      throw new PolyProviderError(ev.error);135    }136  }137  return acc.toResponse(provider, model);138}139