@axiomhq/ai
Version:
Axiom AI SDK provides an API to wrap your AI calls with observability instrumentation.
486 lines (473 loc) • 17.6 kB
JavaScript
#!/usr/bin/env node
"use strict";
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
// If the importer is in node compatibility mode or this is not an ESM
// file that has been converted to a CommonJS file using a Babel-
// compatible transform (i.e. "__esModule" has not been set), then set
// "default" to the CommonJS "module.exports" for node compatibility.
isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
mod
));
// src/bin.ts
var import_env = __toESM(require("@next/env"), 1);
var import_commander4 = require("commander");
// src/commands/push.ts
var import_commander = require("commander");
// src/transpiler.ts
var import_esbuild = require("esbuild");
var import_node_os = __toESM(require("os"), 1);
var import_promises = __toESM(require("fs/promises"), 1);
var import_node_path = __toESM(require("path"), 1);
var import_node_path2 = require("path");
async function loadPromptModule(filePath) {
const result = await (0, import_esbuild.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 = import_node_os.default.tmpdir();
const tempFileName = `axiom-ai-prompt-${Date.now()}-${Math.random().toString(36).substring(2)}.mjs`;
const tempFilePath = import_node_path.default.join(tempDir, tempFileName);
try {
await import_promises.default.writeFile(tempFilePath, code, "utf-8");
const moduleUrl = `file://${tempFilePath}`;
const module2 = await import(moduleUrl);
return module2.default || module2;
} finally {
try {
await import_promises.default.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 = (0, import_node_path2.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 '@axiomhq/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.ts
var import_promises2 = __toESM(require("fs/promises"), 1);
var import_node_readline = __toESM(require("readline"), 1);
async function askConfirmation(message) {
const rl = import_node_readline.default.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 import_commander.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 import_promises2.default.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.ts
var import_commander2 = require("commander");
var fs3 = __toESM(require("fs/promises"), 1);
var path2 = __toESM(require("path"), 1);
var loadPullCommand = (program2) => {
const pull = new import_commander2.Command("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/run.ts
var import_commander3 = require("commander");
// src/evals/run-vitest.ts
var import_node = require("vitest/node");
// src/evals/reporter.ts
var AxiomReporter = class {
onTestSuiteReady(_testSuite) {
}
onTestSuiteResult(testSuite) {
for (const test of testSuite.children.array()) {
if (test.type !== "test") continue;
const testMeta = test.meta();
if (!testMeta.eval) {
return;
}
const scores = [];
for (const s of Object.entries(testMeta.eval.scores)) {
scores.push({ name: s[1].name, score: s[1].score });
}
}
}
async onTestRunEnd(_testModules, _errors, _reason) {
}
};
// src/evals/instrument.ts
var import_sdk_trace_node = require("@opentelemetry/sdk-trace-node");
var import_resources = require("@opentelemetry/resources");
var import_exporter_trace_otlp_http = require("@opentelemetry/exporter-trace-otlp-http");
var import_api = require("@opentelemetry/api");
var collectorOptions = {
url: process.env.AXIOM_URL ? `${process.env.AXIOM_URL}/v1/traces` : "https://api.axiom.co/v1/traces",
// Axiom API endpoint for trace data
headers: {
Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
// Replace API_TOKEN with your actual API token
"X-Axiom-Dataset": process.env.AXIOM_DATASET || ""
// Replace DATASET_NAME with your dataset
},
concurrencyLimit: 10
// an optional limit on pending requests
};
var exporter = new import_exporter_trace_otlp_http.OTLPTraceExporter(collectorOptions);
var processor = new import_sdk_trace_node.BatchSpanProcessor(exporter, {
maxQueueSize: 2048,
maxExportBatchSize: 512,
scheduledDelayMillis: 5e3,
exportTimeoutMillis: 3e4
});
var provider = new import_sdk_trace_node.NodeTracerProvider({
resource: (0, import_resources.resourceFromAttributes)({
["service.name"]: "axiom-ai",
["service.version"]: "0.8.0"
}),
spanProcessors: [processor]
});
provider.register();
var tracer = import_api.trace.getTracer("axiom-ai", "0.8.0");
var flush = async () => {
await provider.forceFlush();
};
// src/evals/run-vitest.ts
var runVitest = async (file) => {
const vi = await (0, import_node.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 = (0, import_node.registerConsoleShortcuts)(vi, process.stdin, process.stdout);
if (!vi.shouldKeepServer()) {
dispose();
await flush();
await vi.close();
process.exit(0);
}
await flush();
};
// src/commands/run.ts
var loadRunCommand = (program2) => {
return program2.addCommand(
new import_commander3.Command("run").description("run evals locally").argument("<path>", "Path to an eval test file, should be in the form of name.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 } = import_env.default;
loadEnvConfig(process.cwd());
var program = new import_commander4.Command();
program.name("axiom").description("Axiom's CLI to manage your objects and run evals").version("0.8.0");
loadPushCommand(program);
loadPullCommand(program);
loadRunCommand(program);
program.parse();
//# sourceMappingURL=bin.cjs.map