149 lines · 6.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 { AiEndpoint } from "../ai/endpoint";
import { assertOutboundUrl, isAllowedLoopback } from "../integrations/outbound";
import { assertResolvesPublic } from "../net/resolve-guard";

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 endpoint. */
  model: string;
  /** From resolveAiEndpoint: the repo's or account's own endpoint, which pays for the call. */
  byok: AiEndpoint | null;
  metadata: GatewayMetadata;
  maxTokens?: number;