SPB Git forge

spb/polyllm

Public
15commits 1branches 0releases
2.2 MBsize
maindefault branch
13 days agolast push
TypeScript 97.4% SQL 1% JavaScript 0.9% CSS 0.6%
6.0 KB · 140 lines typescript
Raw Blame History
1import OpenAI from "openai";2import { writeFileSync, mkdirSync } from "node:fs";3import { deflateSync } from "node:zlib";45export const BASE = "https://api.moonshot.ai/v1";6export const KEY = process.env.KIMI_API_KEY ?? "";7if (!KEY) throw new Error("KIMI_API_KEY missing (set -a; . ./.env; set +a)");89export const MODELS = ["kimi-k2.6", "kimi-k2.7-code", "kimi-k2.7-code-highspeed", "kimi-k3"] as const;10export type ModelId = (typeof MODELS)[number];1112export const client = new OpenAI({ apiKey: KEY, baseURL: BASE, timeout: 600_000, maxRetries: 0 });1314mkdirSync(new URL("./out/", import.meta.url), { recursive: true });15export function save(name: string, data: unknown) {16  const p = new URL(`./out/${name}.json`, import.meta.url);17  writeFileSync(p, JSON.stringify(data, null, 2));18}1920/** Raw POST returning status, headers (redacted) and parsed body (or text). */21export const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));22export let last429: { status: number; headers: Record<string, string>; body: unknown } | null = null;2324/** Retries 429 rate_limit/engine_overloaded up to 8× (3 s, 6 s, …) unless RETRY429=0. */25export async function rawPost(path: string, body: unknown, extraHeaders: Record<string, string> = {}, key = KEY) {26  const maxTries = process.env.RETRY429 === "0" ? 1 : 8;27  let attempt = 0;28  for (;;) {29    attempt++;30    const res = await fetch(`${BASE}${path}`, {31      method: "POST",32      headers: { "content-type": "application/json", authorization: `Bearer ${key}`, ...extraHeaders },33      body: JSON.stringify(body),34    });35    const text = await res.text();36    let json: unknown = null;37    try { json = JSON.parse(text); } catch { /* text */ }38    const out = { status: res.status, headers: pickHeaders(res.headers), body: json ?? text, attempts: attempt };39    if (res.status === 429 && attempt < maxTries) {40      last429 = out;41      const t = (json as any)?.error?.type;42      if (t === "rate_limit_reached_error" || t === "engine_overloaded_error") { await sleep(3000 * attempt); continue; }43    }44    return out;45  }46}4748export async function rawGet(path: string, key = KEY) {49  const res = await fetch(`${BASE}${path}`, { headers: { authorization: `Bearer ${key}` } });50  const text = await res.text();51  let json: unknown = null;52  try { json = JSON.parse(text); } catch { /* text */ }53  return { status: res.status, headers: pickHeaders(res.headers), body: json ?? text };54}5556export function pickHeaders(h: Headers) {57  const out: Record<string, string> = {};58  h.forEach((v, k) => {59    if (/^(set-cookie|authorization)$/i.test(k)) return;60    if (/^(msh-|x-msh-|x-ratelimit|ratelimit|retry-after|content-type|cf-ray|server|date)/i.test(k)) out[k] = v;61  });62  return out;63}6465/** Raw SSE POST: returns every parsed `data:` JSON chunk plus status/headers. */66export async function rawSSE(path: string, body: unknown, extraHeaders: Record<string, string> = {}) {67  const res = await fetch(`${BASE}${path}`, {68    method: "POST",69    headers: { "content-type": "application/json", authorization: `Bearer ${KEY}`, ...extraHeaders },70    body: JSON.stringify({ ...(body as object), stream: true }),71  });72  const headers = pickHeaders(res.headers);73  if (!res.ok || !res.body) {74    const text = await res.text();75    let json: unknown = null; try { json = JSON.parse(text); } catch {}76    return { status: res.status, headers, chunks: [] as any[], error: json ?? text, rawLines: [] as string[] };77  }78  const reader = res.body.getReader();79  const dec = new TextDecoder();80  let buf = "";81  const chunks: any[] = [];82  const rawLines: string[] = [];83  let done = false;84  while (!done) {85    const r = await reader.read();86    if (r.done) break;87    buf += dec.decode(r.value, { stream: true });88    let idx;89    while ((idx = buf.indexOf("\n\n")) >= 0) {90      const evt = buf.slice(0, idx); buf = buf.slice(idx + 2);91      for (const line of evt.split("\n")) {92        rawLines.push(line);93        if (!line.startsWith("data:")) continue;94        const payload = line.slice(5).trim();95        if (payload === "[DONE]") { done = true; chunks.push("[DONE]"); continue; }96        try { chunks.push(JSON.parse(payload)); } catch { chunks.push({ unparsed: payload }); }97      }98    }99  }100  return { status: res.status, headers, chunks, rawLines, error: null };101}102103/** Minimal valid PNG (RGBA), size x size, solid colour. */104export function pngDataUrl(size = 32, rgba: [number, number, number, number] = [200, 30, 30, 255]) {105  const crcTable = new Int32Array(256);106  for (let n = 0; n < 256; n++) { let c = n; for (let k = 0; k < 8; k++) c = c & 1 ? 0xedb88320 ^ (c >>> 1) : c >>> 1; crcTable[n] = c; }107  const crc32 = (buf: Buffer) => { let c = -1; for (const b of buf) c = crcTable[(c ^ b) & 0xff] ^ (c >>> 8); return (c ^ -1) >>> 0; };108  const chunk = (type: string, data: Buffer) => {109    const len = Buffer.alloc(4); len.writeUInt32BE(data.length);110    const td = Buffer.concat([Buffer.from(type, "ascii"), data]);111    const crc = Buffer.alloc(4); crc.writeUInt32BE(crc32(td));112    return Buffer.concat([len, td, crc]);113  };114  const ihdr = Buffer.alloc(13);115  ihdr.writeUInt32BE(size, 0); ihdr.writeUInt32BE(size, 4);116  ihdr[8] = 8; ihdr[9] = 6; ihdr[10] = 0; ihdr[11] = 0; ihdr[12] = 0;117  const row = Buffer.alloc(1 + size * 4);118  for (let x = 0; x < size; x++) row.set(rgba, 1 + x * 4);119  const raw = Buffer.concat(Array.from({ length: size }, () => row));120  const png = Buffer.concat([121    Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]),122    chunk("IHDR", ihdr), chunk("IDAT", deflateSync(raw)), chunk("IEND", Buffer.alloc(0)),123  ]);124  return `data:image/png;base64,${png.toString("base64")}`;125}126127export function short(s: unknown, n = 300) {128  const t = typeof s === "string" ? s : (JSON.stringify(s) ?? String(s));129  return t.length > n ? t.slice(0, n) + `…(${t.length})` : t;130}131132export const WEATHER_TOOL = {133  type: "function" as const,134  function: {135    name: "get_weather",136    description: "Get the current weather for a city",137    parameters: { type: "object", properties: { city: { type: "string", description: "City name" } }, required: ["city"] },138  },139};140