101 lines · 4.2 KB
import { env } from "cloudflare:workers";
import type { ClefRequest, ClefResponse } from "./decide";
import { mockProvider } from "./mock";
import type { ChatMessage, MergeFixProvider, TokenUsage } from "./plan";
import { ModelCallError } from "./plan";
import type { ByokKey } from "./store";

const MAX_OUTPUT_TOKENS = 2048;

type ChatReply = {
  response?: unknown;
  choices?: Array<{ message?: { content?: unknown } }>;
  usage?: { prompt_tokens?: number; completion_tokens?: number; input_tokens?: number; output_tokens?: number };
};

function readReply(reply: ChatReply): { text: string; usage: TokenUsage } {
  const raw = reply.choices?.[0]?.message?.content ?? reply.response;
  if (typeof raw !== "string") throw new ModelCallError("The model returned no text");
  return {
    text: raw,
    usage: {
      inputTokens: reply.usage?.prompt_tokens ?? reply.usage?.input_tokens ?? 0,
      outputTokens: reply.usage?.completion_tokens ?? reply.usage?.output_tokens ?? 0,
    },
  };
}

type Unchecked = { run(model: string, inputs: unknown, options?: AiOptions): Promise<unknown> };

export type GatewayMetadata = { account_id: string; feature: string } & Record<string, string>;

export type ChatTarget = {
  /** Workers AI model used when the account has no bring-your-own key. */
  model: string;
  byok: ByokKey | null;
  metadata: GatewayMetadata;
  maxTokens?: number;
};

/** The model actually billed for a chat call: the account's BYOK model, else the platform one. */