axiom
Version:
Axiom AI SDK provides an API to wrap your AI calls with observability instrumentation.
408 lines (398 loc) • 14.1 kB
JavaScript
#!/usr/bin/env node
import {
AxiomReporter,
flush
} from "./chunk-MNOTFSB6.js";
import "./chunk-KEXKKQVW.js";
// src/bin.ts
import pkg from "@next/env";
import { Command as Command4 } from "commander";
// src/commands/push.command.ts
import { Command } from "commander";
// src/transpiler.ts
import { build } from "esbuild";
import os from "os";
import fs from "fs/promises";
import path from "path";
import { basename } from "path";
async function loadPromptModule(filePath) {
const result = await build({
entryPoints: [filePath],
bundle: true,
write: false,
platform: "node",
format: "esm",
target: ["node18"],
sourcemap: false,
external: [
// Only Node.js built-ins should be external
"fs",
"fs/promises",
"node:fs",
"node:fs/promises",
"readline",
"node:readline",
"path",
"node:path",
"os",
"node:os",
"url",
"node:url",
"util",
"node:util",
"crypto",
"node:crypto",
"events",
"node:events",
"stream",
"node:stream",
"buffer",
"node:buffer",
"process",
"node:process"
]
});
const code = result.outputFiles[0].text;
const tempDir = os.tmpdir();
const tempFileName = `axiom-ai-prompt-${Date.now()}-${Math.random().toString(36).substring(2)}.mjs`;
const tempFilePath = path.join(tempDir, tempFileName);
try {
await fs.writeFile(tempFilePath, code, "utf-8");
const moduleUrl = `file://${tempFilePath}`;
const module = await import(moduleUrl);
return module.default || module;
} finally {
try {
await fs.unlink(tempFilePath);
} catch (error) {
console.warn(`Failed to clean up temporary file ${tempFilePath}:`, error);
}
}
}
function convertTypeBoxArgumentsToJsonSchema(arguments_) {
if (!arguments_ || typeof arguments_ !== "object") {
return {
type: "object",
properties: {},
required: [],
additionalProperties: false
};
}
const properties = {};
const required = [];
for (const [key, value] of Object.entries(arguments_)) {
if (value && typeof value === "object" && value.type) {
properties[key] = {
type: value.type,
...value.description && { description: value.description },
...value.enum && { enum: value.enum },
...value.items && { items: value.items },
...value.properties && { properties: value.properties },
...value.required && { required: value.required }
};
required.push(key);
}
}
return {
type: "object",
properties,
required,
additionalProperties: false
};
}
function extractPromptFromModule(moduleContent, filePath) {
const fileBaseName = basename(filePath, ".ts");
const defaultId = fileBaseName.toLowerCase().replace(/[^a-z0-9]/g, "-");
const convertedArguments = convertTypeBoxArgumentsToJsonSchema(moduleContent.arguments);
const prompt = {
name: moduleContent.name || "Untitled Prompt",
slug: moduleContent.slug || defaultId,
messages: moduleContent.messages || [],
model: moduleContent.model,
options: moduleContent.options,
arguments: convertedArguments,
id: moduleContent.id || defaultId,
version: moduleContent.version || "1.0.0",
// Optional fields from API response
...moduleContent.promptId && { promptId: moduleContent.promptId },
...moduleContent.description && { description: moduleContent.description }
};
if (!prompt.name) {
throw new Error("Prompt must have a name");
}
if (!prompt.slug) {
throw new Error("Prompt must have a slug");
}
if (!Array.isArray(prompt.messages)) {
throw new Error("Prompt messages must be an array");
}
if (!prompt.model) {
throw new Error("Prompt must have a model");
}
return prompt;
}
function transformJsonSchemaToTypeBox(schema) {
if (schema.type === "string") {
if (schema.enum && Array.isArray(schema.enum)) {
const literals = schema.enum.map((value) => `Type.Literal('${value}')`).join(", ");
const options = schema.description ? `, { description: '${schema.description}' }` : "";
return `Type.Union([${literals}]${options})`;
} else {
const options = schema.description ? `{ description: '${schema.description}' }` : "";
return `Type.String(${options})`;
}
}
if (schema.type === "number" || schema.type === "integer") {
const typeMethod = schema.type === "integer" ? "Integer" : "Number";
const options = schema.description ? `{ description: '${schema.description}' }` : "";
return `Type.${typeMethod}(${options})`;
}
if (schema.type === "boolean") {
const options = schema.description ? `{ description: '${schema.description}' }` : "";
return `Type.Boolean(${options})`;
}
if (schema.type === "array") {
const itemsType = schema.items ? transformJsonSchemaToTypeBox(schema.items) : "Type.String()";
const options = schema.description ? `, { description: '${schema.description}' }` : "";
return `Type.Array(${itemsType}${options})`;
}
if (schema.type === "object") {
if (schema.properties) {
const props = Object.entries(schema.properties).map(([key, value]) => {
const isRequired = schema.required && schema.required.includes(key);
const propType = transformJsonSchemaToTypeBox(value);
return ` ${key}: ${isRequired ? propType : `Type.Optional(${propType})`}`;
}).join(",\n");
const options = schema.description ? `, { description: '${schema.description}' }` : "";
return `Type.Object({
${props}
}${options})`;
} else {
const options = schema.description ? `{ description: '${schema.description}' }` : "";
return `Type.Object({}${options ? `, ${options}` : ""})`;
}
}
return "Type.String()";
}
function generatePromptFileFromApiResponse(apiResponse) {
const { prompt, version } = apiResponse;
const { data, options } = version;
let argumentsCode = "{}";
if (data.arguments && data.arguments.properties) {
const argEntries = Object.entries(data.arguments.properties).map(([key, schema]) => {
const isRequired = data.arguments.required && data.arguments.required.includes(key);
const typeCode = transformJsonSchemaToTypeBox(schema);
return ` ${key}: ${isRequired ? typeCode : `Type.Optional(${typeCode})`}`;
}).join(",\n");
if (argEntries) {
argumentsCode = `{
${argEntries}
}`;
}
}
return `import { Type } from 'axiom/ai';
export default {
name: '${prompt.name}',
slug: '${prompt.slug}',
description: '${prompt.description || ""}',
messages: [${data.messages.map(
(msg) => `
{
role: '${msg.role}',
content: '${msg.content.replace(/'/g, "\\'")}',
}`
).join(",")}
],
model: '${data.model || "gpt-4"}',
options: {
${options ? Object.entries(options).map(([key, value]) => ` ${key}: ${value}`).join(",\n") : ""}
},
arguments: ${argumentsCode},
version: '${version.version}',
promptId: '${prompt.promptId}',
};
`;
}
// src/commands/push.command.ts
import fs2 from "fs/promises";
import readline from "readline";
async function askConfirmation(message) {
const rl = readline.createInterface({
input: process.stdin,
output: process.stdout
});
return new Promise((resolve2) => {
rl.question(`${message} (y/N): `, (answer) => {
rl.close();
resolve2(answer.toLowerCase() === "y" || answer.toLowerCase() === "yes");
});
});
}
var loadPushCommand = (program2) => {
const push = new Command("push").description("Push a new version of an object").argument(
"<object>",
"The object to push, could be a prompt, en eval, a monitor, a dashboard, etc."
).option("--prod", "Adds the production tag to the prompt").option("--yes", "Automatically confirm overwriting the file with server response").action(async (filePath, options) => {
let content = null;
if (!filePath.endsWith(".prompt.ts")) {
console.error("Prompt files must end with .prompt.ts");
process.exit(1);
}
try {
const moduleContent = await loadPromptModule(filePath);
const promptData = extractPromptFromModule(moduleContent, filePath);
content = promptData;
console.log(`Transpiled prompt: ${promptData.name} (${promptData.slug})`);
} catch (error) {
console.error("Failed to transpile prompt file:", error);
process.exit(1);
}
if (!content) {
console.error("No content found");
process.exit(1);
}
let shouldProceed = options.yes;
if (!shouldProceed) {
shouldProceed = await askConfirmation(
`This will push "${content.name}" to Axiom and overwrite ${filePath}, are you sure you want to continue?`
);
}
if (!shouldProceed) {
console.log("Push operation cancelled.");
process.exit(0);
}
try {
const response = await fetch(`${process.env.AXIOM_URL}/v1/prompts`, {
method: "POST",
headers: {
Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
"Content-Type": "application/json",
"x-axiom-client": "axiom-ai-cli",
"x-axiom-check": "good"
},
body: JSON.stringify({
...content,
tags: options.yes ? ["production"] : []
})
});
if (!response.ok) {
try {
const errorText = await response.clone().json();
console.error(`Failed to fetch prompt: ${response.status} ${response.statusText}`);
console.error(JSON.stringify(errorText, null, 2));
process.exit(1);
} catch (_error) {
const errorText = await response.clone().text();
console.error(`Failed to fetch prompt: ${response.status} ${response.statusText}`);
console.error(errorText);
process.exit(1);
}
}
const apiResponse = await response.json();
console.log(
`Successfully pushed prompt: ${apiResponse.prompt.name} (${apiResponse.prompt.slug})`
);
console.log(`Version: ${apiResponse.version.version}`);
const updatedTsContent = generatePromptFileFromApiResponse(apiResponse);
await fs2.writeFile(filePath, updatedTsContent, "utf-8");
console.log(`Successfully updated ${filePath}`);
} catch (error) {
console.error("Failed to push prompt:", error);
process.exit(1);
}
});
program2.addCommand(push);
};
// src/commands/pull.command.ts
import { Command as Command2 } from "commander";
import * as fs3 from "fs/promises";
import * as path2 from "path";
var loadPullCommand = (program2) => {
const pull = new Command2("pull").description("Pull a version of an object").argument(
"<slug>",
"The object to pull, could be a prompt, en eval, a monitor, a dashboard, etc."
).option("--version <version>", "The version to pull, default: latest", "latest").option("--output <path>", "Output file path (optional, defaults to <slug>.prompt.ts)").action(async (slug, options) => {
try {
console.log(`Pulling prompt: ${slug} (version: ${options.version})`);
const url = `${process.env.AXIOM_URL}/v1/prompts/${slug}`;
const response = await fetch(url, {
method: "GET",
headers: {
Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
"Content-Type": "application/json",
"x-axiom-client": "axiom-ai-cli",
"x-axiom-check": "good"
}
});
if (!response.ok) {
try {
const errorText = await response.clone().json();
console.error(`Failed to fetch prompt: ${response.status} ${response.statusText}`);
console.error(JSON.stringify(errorText, null, 2));
process.exit(1);
} catch (_error) {
const errorText = await response.clone().text();
console.error(`Failed to fetch prompt: ${response.status} ${response.statusText}`);
console.error(errorText);
process.exit(1);
}
}
const apiResponse = await response.json();
const tsContent = generatePromptFileFromApiResponse(apiResponse);
const outputPath = options.output || `${slug}.prompt.ts`;
const fullPath = path2.resolve(outputPath);
await fs3.writeFile(fullPath, tsContent, "utf-8");
console.log(`Successfully generated prompt file: ${fullPath}`);
console.log(`Prompt: ${apiResponse.prompt.name} (${apiResponse.prompt.slug})`);
console.log(`Version: ${apiResponse.version.version}`);
} catch (error) {
console.error("Failed to pull prompt:", error);
process.exit(1);
}
});
program2.addCommand(pull);
};
// src/commands/eval.command.ts
import { Command as Command3 } from "commander";
// src/evals/run-vitest.ts
import { createVitest, registerConsoleShortcuts } from "vitest/node";
var runVitest = async (file) => {
const vi = await createVitest("test", {
// root: process.cwd(),
mode: "test",
include: [file ? file : "**/*.eval.ts"],
reporters: ["verbose", new AxiomReporter()],
environment: "node",
browser: void 0
});
await vi.start();
const dispose = registerConsoleShortcuts(vi, process.stdin, process.stdout);
if (!vi.shouldKeepServer()) {
dispose();
await flush();
await vi.close();
process.exit(0);
}
await flush();
};
// src/commands/eval.command.ts
var loadRunCommand = (program2) => {
return program2.addCommand(
new Command3("eval").description("run evals locally").argument("<path>", "Path to an eval test file, should be in the form of *.eval.ts").action(async (file) => {
if (!process.env.AXIOM_URL || !process.env.AXIOM_TOKEN || !process.env.AXIOM_DATASET) {
throw new Error("AXIOM_URL, AXIOM_TOKEN, and AXIOM_DATASET must be set");
}
await runVitest(file);
})
);
};
// src/bin.ts
var { loadEnvConfig } = pkg;
loadEnvConfig(process.cwd());
var program = new Command4();
program.name("axiom").description("Axiom's CLI to manage your objects and run evals").version("0.13.0");
loadPushCommand(program);
loadPullCommand(program);
loadRunCommand(program);
program.parse();
//# sourceMappingURL=bin.js.map