spb/polyllm

1// Shared helpers for Mistral probes. Key comes from process.env.MISTRAL_API_KEY (never logged).2import { mkdirSync, writeFileSync } from "node:fs";3import { deflateSync } from "node:zlib";4import OpenAI from "openai";5import { Mistral } from "@mistralai/mistralai";67export const BASE = "https://api.mistral.ai/v1";8export const KEY = process.env.MISTRAL_API_KEY ?? "";9if (!KEY) {10  console.error("MISTRAL_API_KEY missing (run: set -a; . ./.env; set +a)");11  process.exit(1);12}1314export const OUT = new URL("./out/", import.meta.url).pathname;15mkdirSync(OUT, { recursive: true });1617export const openai = new OpenAI({ apiKey: KEY, baseURL: BASE, timeout: 120_000, maxRetries: 0 });18export const mistral = new Mistral({ apiKey: KEY, retryConfig: { strategy: "none" } });1920export const CHAT_MODELS = [21  "mistral-large-latest",22  "mistral-medium-latest",23  "mistral-small-latest",24  "ministral-8b-latest",25  "magistral-medium-latest",26  "magistral-small-latest",27  "codestral-latest",28];2930export function save(name: string, data: unknown) {31  writeFileSync(`${OUT}${name}`, JSON.stringify(data, null, 2));32  console.log(`saved out/${name}`);33}3435const HDR = /ratelimit|request-id|x-request|retry-after|content-type|server|date|cf-ray|x-kong|x-envoy|x-mistral|x-ms/i;3637/** Raw request; returns status, selected headers, and parsed body (never echoes the key). */38export async function raw(path: string, init: RequestInit & { key?: string } = {}) {39  const { key, ...rest } = init;40  const res = await fetch(`${BASE}${path}`, {41    ...rest,42    headers: {43      "Content-Type": "application/json",44      Accept: "application/json",45      Authorization: `Bearer ${key ?? KEY}`,46      ...(rest.headers ?? {}),47    },48  });49  const text = await res.text();50  let body: unknown = text;51  try {52    body = JSON.parse(text);53  } catch {}54  const headers: Record<string, string> = {};55  res.headers.forEach((v, k) => {56    if (HDR.test(k)) headers[k] = v;57  });58  return { status: res.status, headers, body };59}6061export function post(path: string, payload: unknown, key?: string) {62  return raw(path, { method: "POST", body: JSON.stringify(payload), key });63}6465/** Raw SSE POST; returns list of parsed events (data lines). */66export async function rawSSE(path: string, payload: unknown, key?: string) {67  const res = await fetch(`${BASE}${path}`, {68    method: "POST",69    headers: { "Content-Type": "application/json", Accept: "text/event-stream", Authorization: `Bearer ${key ?? KEY}` },70    body: JSON.stringify(payload),71  });72  const headers: Record<string, string> = {};73  res.headers.forEach((v, k) => {74    if (HDR.test(k)) headers[k] = v;75  });76  if (!res.ok || !res.body) {77    const text = await res.text();78    let body: unknown = text;79    try {80      body = JSON.parse(text);81    } catch {}82    return { status: res.status, headers, error: body, events: [] as any[] };83  }84  const reader = res.body.getReader();85  const dec = new TextDecoder();86  let buf = "";87  const events: { event?: string; data: any; raw?: string }[] = [];88  for (;;) {89    const { value, done } = await reader.read();90    if (done) break;91    buf += dec.decode(value, { stream: true });92    let idx;93    while ((idx = buf.indexOf("\n\n")) >= 0) {94      const block = buf.slice(0, idx);95      buf = buf.slice(idx + 2);96      let ev: string | undefined;97      const datas: string[] = [];98      for (const line of block.split("\n")) {99        if (line.startsWith("event:")) ev = line.slice(6).trim();100        else if (line.startsWith("data:")) datas.push(line.slice(5).trim());101      }102      if (!datas.length) continue;103      const d = datas.join("\n");104      if (d === "[DONE]") {105        events.push({ event: ev, data: "[DONE]" });106        continue;107      }108      try {109        events.push({ event: ev, data: JSON.parse(d) });110      } catch {111        events.push({ event: ev, data: null, raw: d });112      }113    }114  }115  return { status: res.status, headers, events };116}117118export function short(s: unknown, n = 300) {119  const t = typeof s === "string" ? s : JSON.stringify(s) ?? String(s);120  return t.length > n ? t.slice(0, n) + "…" : t;121}122123/** Minimal valid PNG (RGB, solid colour with a diagonal), size x size, base64 data URL. */124export function pngDataUrl(size = 32): string {125  const crcTable = new Int32Array(256);126  for (let n = 0; n < 256; n++) {127    let c = n;128    for (let k = 0; k < 8; k++) c = c & 1 ? 0xedb88320 ^ (c >>> 1) : c >>> 1;129    crcTable[n] = c;130  }131  const crc32 = (buf: Buffer) => {132    let c = -1;133    for (const b of buf) c = crcTable[(c ^ b) & 0xff] ^ (c >>> 8);134    return (c ^ -1) >>> 0;135  };136  const chunk = (type: string, data: Buffer) => {137    const len = Buffer.alloc(4);138    len.writeUInt32BE(data.length);139    const td = Buffer.concat([Buffer.from(type, "ascii"), data]);140    const crc = Buffer.alloc(4);141    crc.writeUInt32BE(crc32(td));142    return Buffer.concat([len, td, crc]);143  };144  const ihdr = Buffer.alloc(13);145  ihdr.writeUInt32BE(size, 0);146  ihdr.writeUInt32BE(size, 4);147  ihdr[8] = 8; // bit depth148  ihdr[9] = 2; // RGB149  const rows: Buffer[] = [];150  for (let y = 0; y < size; y++) {151    const row = Buffer.alloc(1 + size * 3);152    for (let x = 0; x < size; x++) {153      const onDiag = Math.abs(x - y) < 2;154      row[1 + x * 3] = onDiag ? 255 : 30; // R155      row[2 + x * 3] = onDiag ? 255 : 60; // G156      row[3 + x * 3] = onDiag ? 255 : 200; // B157    }158    rows.push(row);159  }160  const idat = deflateSync(Buffer.concat(rows));161  const png = Buffer.concat([162    Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]),163    chunk("IHDR", ihdr),164    chunk("IDAT", idat),165    chunk("IEND", Buffer.alloc(0)),166  ]);167  return `data:image/png;base64,${png.toString("base64")}`;168}169