pi-ai generic stream

pi-ai-stream.ts

#!/usr/bin/env -S deno run --allow-env=GEMINI_API_KEY,ANTHROPIC_*,OPENAI_*,GOOGLE_*,AWS_*,WS_* --allow-net=api.anthropic.com,bedrock-runtime.us-east-1.amazonaws.com,bedrock-runtime.us-west-2.amazonaws.com,bedrock-runtime.eu-west-1.amazonaws.com,bedrock-runtime.ap-southeast-1.amazonaws.com,bedrock-runtime.ap-northeast-1.amazonaws.com,generativelanguage.googleapis.com,api.openai.com --allow-read=/dev/stdin --allow-write=/dev/stdout

import { getModel, streamSimple } from "npm:@mariozechner/pi-ai";
import type {
  Context,
  KnownProvider,
  SimpleStreamOptions,
} from "npm:@mariozechner/pi-ai";

interface Request {
  provider: KnownProvider;
  modelId: string;
  context: Context;
  options?: SimpleStreamOptions;
}

async function readStdin(): Promise<string> {
  const chunks: Buffer[] = [];
  for await (const chunk of process.stdin) {
    chunks.push(chunk);
  }
  return Buffer.concat(chunks).toString("utf-8");
}

async function main(): Promise<number> {
  const controller = new AbortController();

  process.on("SIGTERM", () => {
    controller.abort();
  });

  process.on("SIGINT", () => {
    controller.abort();
  });

  const input = await readStdin();
  const request: Request = JSON.parse(input);

  const model = getModel(request.provider, request.modelId as any);
  if (!model) {
    const errorEvent = {
      type: "error",
      reason: "error",
      error: {
        role: "assistant",
        content: [],
        api: "unknown",
        provider: request.provider,
        model: request.modelId,
        usage: {
          input: 0,
          output: 0,
          cacheRead: 0,
          cacheWrite: 0,
          totalTokens: 0,
          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
        },
        stopReason: "error",
        errorMessage: `Unknown model: ${request.provider}/${request.modelId}`,
        timestamp: Date.now(),
      },
    };
    process.stdout.write(JSON.stringify(errorEvent) + "\n");
    return 1;
  }

  const options: SimpleStreamOptions = {
    // Default reasoning to "high" if not specified and model supports it
    reasoning: model.reasoning ? "high" : undefined,
    ...request.options,
    signal: controller.signal,
  };

  const stream = streamSimple(model, request.context, options);

  for await (const event of stream) {
    process.stdout.write(JSON.stringify(event) + "\n");
  }

  return 0;
}

try {
  const code = await main();
  process.exit(code);
} catch (err) {
  const errorEvent = {
    type: "error",
    reason: "error",
    error: {
      role: "assistant",
      content: [],
      api: "unknown",
      provider: "unknown",
      model: "unknown",
      usage: {
        input: 0,
        output: 0,
        cacheRead: 0,
        cacheWrite: 0,
        totalTokens: 0,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      },
      stopReason: "error",
      errorMessage: err instanceof Error ? err.message : String(err),
      timestamp: Date.now(),
    },
  };
  process.stdout.write(JSON.stringify(errorEvent) + "\n");
  process.exit(1);
}