spb/polyllm

1// Shared helpers for Cerebras probes. Key comes from process.env.CEREBRAS_API_KEY (never logged).2import { mkdirSync, writeFileSync } from "node:fs";3import OpenAI from "openai";4import Cerebras from "@cerebras/cerebras_cloud_sdk";56export const BASE = "https://api.cerebras.ai/v1";7export const KEY = process.env.CEREBRAS_API_KEY ?? "";8if (!KEY) {9  console.error("CEREBRAS_API_KEY missing (run: set -a; . ./.env; set +a)");10  process.exit(1);11}1213export const OUT = new URL("./out/", import.meta.url).pathname;14mkdirSync(OUT, { recursive: true });1516export const openai = new OpenAI({ apiKey: KEY, baseURL: BASE, timeout: 120_000, maxRetries: 0 });17export const cerebras = new Cerebras({ apiKey: KEY, timeout: 120_000, maxRetries: 0, warmTCPConnection: false });1819export const MODELS = ["gemma-4-31b", "qwen-3.8-27b", "gpt-oss-120b"];2021export function save(name: string, data: unknown) {22  writeFileSync(`${OUT}${name}`, JSON.stringify(data, null, 2));23  console.log(`saved out/${name}`);24}2526export const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));2728export function pickHeaders(h: Headers) {29  const headers: Record<string, string> = {};30  h.forEach((v, k) => {31    if (/ratelimit|request-id|x-request|retry-after|content-type|server|date|cf-ray|x-cerebras|via|cache/i.test(k)) headers[k] = v;32  });33  return headers;34}3536/** Raw request; returns status, selected headers, and parsed body (never echoes the key). */37export async function raw(path: string, init: RequestInit & { key?: string } = {}) {38  const { key, ...rest } = init;39  const t0 = Date.now();40  const res = await fetch(`${BASE}${path}`, {41    ...rest,42    headers: {43      "Content-Type": "application/json",44      Authorization: `Bearer ${key ?? KEY}`,45      ...(rest.headers ?? {}),46    },47  });48  const text = await res.text();49  let body: unknown = text;50  try {51    body = JSON.parse(text);52  } catch {}53  return { status: res.status, ms: Date.now() - t0, headers: pickHeaders(res.headers), body };54}5556export async function chat(payload: Record<string, unknown>, key?: string) {57  return raw("/chat/completions", { method: "POST", body: JSON.stringify(payload), key });58}5960/** Raw SSE POST; returns list of parsed events (data lines). */61export async function rawSSE(path: string, payload: unknown, key?: string) {62  const t0 = Date.now();63  const res = await fetch(`${BASE}${path}`, {64    method: "POST",65    headers: { "Content-Type": "application/json", Authorization: `Bearer ${key ?? KEY}` },66    body: JSON.stringify(payload),67  });68  const headers = pickHeaders(res.headers);69  if (!res.ok || !res.body) {70    const text = await res.text();71    let body: unknown = text;72    try {73      body = JSON.parse(text);74    } catch {}75    return { status: res.status, headers, error: body, events: [] as any[], ttfbMs: Date.now() - t0, totalMs: Date.now() - t0 };76  }77  const reader = res.body.getReader();78  const dec = new TextDecoder();79  let buf = "";80  let ttfbMs = -1;81  const events: { event?: string; data: any; raw?: string; t: number }[] = [];82  for (;;) {83    const { value, done } = await reader.read();84    if (done) break;85    if (ttfbMs < 0) ttfbMs = Date.now() - t0;86    buf += dec.decode(value, { stream: true });87    let idx;88    while ((idx = buf.indexOf("\n\n")) >= 0) {89      const block = buf.slice(0, idx);90      buf = buf.slice(idx + 2);91      let ev: string | undefined;92      const datas: string[] = [];93      for (const line of block.split("\n")) {94        if (line.startsWith("event:")) ev = line.slice(6).trim();95        else if (line.startsWith("data:")) datas.push(line.slice(5).trim());96      }97      if (!datas.length) continue;98      const d = datas.join("\n");99      if (d === "[DONE]") {100        events.push({ event: ev, data: "[DONE]", t: Date.now() - t0 });101        continue;102      }103      try {104        events.push({ event: ev, data: JSON.parse(d), t: Date.now() - t0 });105      } catch {106        events.push({ event: ev, data: null, raw: d, t: Date.now() - t0 });107      }108    }109  }110  return { status: res.status, headers, events, ttfbMs, totalMs: Date.now() - t0 };111}112113export function short(s: unknown, n = 300) {114  const t = typeof s === "string" ? s : JSON.stringify(s) ?? String(s);115  return t.length > n ? t.slice(0, n) + "…" : t;116}117