UNPKG

@axflow/models

Version:

Zero-dependency, modular SDK for building robust natural language applications

102 lines (99 loc) 2.94 kB
// src/cohere/generation.ts import { POST, HttpError } from "@axflow/models/shared"; // src/cohere/shared.ts function headers(apiKey, customHeaders) { const headers2 = { accept: "application/json", "content-type": "application/json", ...customHeaders }; if (typeof apiKey === "string") { headers2.authorization = `Bearer ${apiKey}`; } return headers2; } // src/cohere/generation.ts var COHERE_API_URL = "https://api.cohere.ai/v1/generate"; async function run(request, options) { const url = options.apiUrl || COHERE_API_URL; const response = await POST(url, { headers: headers(options.apiKey, options.headers), body: JSON.stringify({ ...request, stream: false }), fetch: options.fetch, signal: options.signal }); return response.json(); } async function streamBytes(request, options) { const url = options.apiUrl || COHERE_API_URL; const response = await POST(url, { headers: headers(options.apiKey, options.headers), body: JSON.stringify({ ...request, stream: true }), fetch: options.fetch, signal: options.signal }); if (!response.body) { throw new HttpError("Expected response body to be a ReadableStream", response); } return response.body; } function noop(chunk) { return chunk; } async function stream(request, options) { const byteStream = await streamBytes(request, options); return byteStream.pipeThrough(new CohereGenerationDecoderStream(noop)); } function chunkToToken(chunk) { return chunk.text || ""; } async function streamTokens(request, options) { const byteStream = await streamBytes(request, options); return byteStream.pipeThrough(new CohereGenerationDecoderStream(chunkToToken)); } var CohereGeneration = class { static run = run; static stream = stream; static streamBytes = streamBytes; static streamTokens = streamTokens; }; var CohereGenerationDecoderStream = class _CohereGenerationDecoderStream extends TransformStream { static parse(line) { line = line.trim(); if (line.length === 0) { return null; } try { return JSON.parse(line); } catch (error) { throw new Error( `Invalid event: expected well-formed event lines but got ${JSON.stringify(line)}` ); } } static transformer(map) { let buffer = []; const decoder = new TextDecoder(); return (bytes, controller) => { const chunk = decoder.decode(bytes); for (let i = 0, len = chunk.length; i < len; ++i) { const isEventSeparator = chunk[i] === "\n"; if (!isEventSeparator) { buffer.push(chunk[i]); continue; } const event = _CohereGenerationDecoderStream.parse(buffer.join("")); if (event) { controller.enqueue(map(event)); } buffer = []; } }; } constructor(map) { super({ transform: _CohereGenerationDecoderStream.transformer(map) }); } }; export { CohereGeneration };