ai-xray 1.2.0 → 2.0.0

package/src/commands/tokenize.ts

@@ -0,0 +1,96 @@
+import * as fs from 'fs';
+
+export interface TokenResult {
+    characters: number;
+    words: number;
+    estimated_tokens: number;
+    cost_estimate?: {
+        model: string;
+        input_cost_usd: number;
+    };
+}
+
+// Simple BPE-style token estimation
+// Average: 4 characters per token for English, but varies by content
+function estimateTokens(text: string): number {
+    // Count tokens more accurately by considering:
+    // - Word boundaries (~1 token per word on average)
+    // - Punctuation (~1 token per ~4 chars)
+    // - Numbers (~1 token per number sequence)
+    
+    const words = text.split(/\s+/).filter(w => w.length > 0);
+    const wordCount = words.length;
+    
+    // Base estimation: words + chars/4
+    const charBased = Math.ceil(text.length / 4);
+    const wordBased = Math.ceil(wordCount * 1.3); // ~1.3 tokens per word
+    
+    // Use average of both approaches
+    const estimated = Math.round((charBased + wordBased) / 2);
+    
+    return Math.max(1, estimated);
+}
+
+// Pricing constants (USD per 1M tokens) - approximate
+const MODEL_PRICING: Record<string, number> = {
+    'gpt-4o': 2.50,      // $2.50 per 1M input tokens
+    'gpt-4o-mini': 0.15,
+    'gpt-4-turbo': 10.00,
+    'gpt-4': 30.00,
+    'gpt-3.5-turbo': 0.50,
+    'claude-3-5-sonnet-2024': 3.00,
+    'claude-3-opus-2024': 15.00,
+    'claude-3-haiku-2024': 0.25,
+    'gemini-1.5-pro': 1.25,
+    'gemini-1.5-flash': 0.075,
+};
+
+function getModelPrice(modelName: string): number | null {
+    const lowerModel = modelName.toLowerCase();
+    for (const [key, price] of Object.entries(MODEL_PRICING)) {
+        if (lowerModel.includes(key)) {
+            return price;
+        }
+    }
+    return null;
+}
+
+export async function countTokens(
+    input: string,
+    options?: { model?: string }
+): Promise<TokenResult> {
+    let text = input;
+
+    // If input looks like a file path, read the file
+    if (!input.includes('\n') && fs.existsSync(input)) {
+        try {
+            text = fs.readFileSync(input, 'utf-8');
+        } catch {
+            // Keep original input if file read fails
+        }
+    }
+
+    const characters = text.length;
+    const words = text.split(/\s+/).filter(w => w.length > 0).length;
+    const estimatedTokens = estimateTokens(text);
+
+    const result: TokenResult = {
+        characters,
+        words,
+        estimated_tokens: estimatedTokens,
+    };
+
+    // Add cost estimate if model specified
+    if (options?.model) {
+        const price = getModelPrice(options.model);
+        if (price !== null) {
+            const cost = (estimatedTokens / 1000000) * price;
+            result.cost_estimate = {
+                model: options.model,
+                input_cost_usd: parseFloat(cost.toFixed(5)),
+            };
+        }
+    }
+
+    return result;
+}

package/src/utils/http.ts

@@ -0,0 +1,86 @@
+import * as http from 'http';
+import * as https from 'https';
+import { URL } from 'url';
+
+export interface HttpRequestOptions {
+    method?: string;
+    headers?: Record<string, string>;
+    body?: string;
+    timeout?: number;
+}
+
+export interface HttpResponse {
+    statusCode: number;
+    headers: Record<string, string>;
+    body: string;
+}
+
+export function request(
+    urlString: string,
+    options: HttpRequestOptions = {}
+): Promise<HttpResponse> {
+    return new Promise((resolve, reject) => {
+        const url = new URL(urlString);
+        const isHttps = url.protocol === 'https:';
+        const client = isHttps ? https : http;
+
+        const requestOptions: http.RequestOptions = {
+            hostname: url.hostname,
+            port: url.port || (isHttps ? 443 : 80),
+            path: url.pathname + url.search,
+            method: options.method || 'GET',
+            headers: {
+                'Content-Type': 'application/json',
+                'Accept': 'application/json',
+                ...options.headers,
+            },
+            timeout: options.timeout || 30000,
+        };
+
+        const req = client.request(requestOptions, (res) => {
+            let body = '';
+            res.on('data', (chunk) => { body += chunk; });
+            res.on('end', () => {
+                const headers: Record<string, string> = {};
+                for (const [key, value] of Object.entries(res.headers)) {
+                    if (typeof value === 'string') {
+                        headers[key] = value;
+                    } else if (Array.isArray(value)) {
+                        headers[key] = value.join(', ');
+                    }
+                }
+                resolve({
+                    statusCode: res.statusCode || 0,
+                    headers,
+                    body,
+                });
+            });
+        });
+
+        req.on('error', reject);
+        req.on('timeout', () => {
+            req.destroy();
+            reject(new Error('Request timeout'));
+        });
+
+        if (options.body) {
+            req.write(options.body);
+        }
+        req.end();
+    });
+}
+
+export function requestJson<T = unknown>(
+    urlString: string,
+    options: HttpRequestOptions = {}
+): Promise<{ response: HttpResponse; data: T }> {
+    return request(urlString, options).then((response) => {
+        let data: T;
+        try {
+            data = JSON.parse(response.body) as T;
+        } catch {
+            throw new Error(`Invalid JSON response: ${response.body.slice(0, 100)}`);
+        }
+        return { response, data };
+    });
+}

package/src/utils/output.ts

@@ -1,123 +1,36 @@
-/**
- * Core output management for ai-xray.
- * Guarantees Machine-First protocol: Pure JSON on stdout, structured error on stderr.
- */
-
-// A very rough token estimator (approximately 4 characters per token for English/code)
-export function estimateTokens(text: string): number {
-    return Math.max(1, Math.ceil(text.length / 4));
-}
-
-// Ensure the final object doesn't exceed the token budget if one is provided
-export function applyBudget(data: any, budget?: number): any {
-    if (!budget || budget <= 0) return data;
-
-    const serialized = JSON.stringify(data);
-    const estimated = estimateTokens(serialized);
-
-    if (estimated > budget) {
-        // Actually truncate: serialize to JSON, cut to budget char limit, close cleanly
-        const budgetChars = budget * 4;
-
-        if (typeof data === 'object' && data !== null && !Array.isArray(data)) {
-            // Try pruning array values first (largest arrays get trimmed)
-            const pruned = { ...data };
-            const keys = Object.keys(pruned);
-
-            // Sort keys by serialized size descending to prune largest first
-            const keySizes = keys
-                .filter(k => k !== '_meta')
-                .map(k => ({ key: k, size: JSON.stringify(pruned[k]).length }))
-                .sort((a, b) => b.size - a.size);
-
-            for (const { key } of keySizes) {
-                const currentSize = JSON.stringify(pruned).length;
-                if (currentSize <= budgetChars) break;
-
-                if (Array.isArray(pruned[key]) && pruned[key].length > 0) {
-                    // Trim array items from the end until under budget
-                    const arr = [...pruned[key]];
-                    while (arr.length > 0 && JSON.stringify({ ...pruned, [key]: arr }).length > budgetChars) {
-                        arr.pop();
-                    }
-                    pruned[key] = arr.length > 0 ? arr : `[TRUNCATED: ${pruned[key].length} items removed]`;
-                } else if (typeof pruned[key] === 'string' && pruned[key].length > 200) {
-                    // Truncate long string values
-                    const allowedLen = Math.max(100, pruned[key].length - (currentSize - budgetChars));
-                    pruned[key] = pruned[key].slice(0, allowedLen) + '...[TRUNCATED]';
-                }
-            }
-
-            // Final fallback: if still over budget, hard-truncate the JSON string
-            const prunedJson = JSON.stringify(pruned);
-            if (prunedJson.length > budgetChars) {
-                const truncatedStr = prunedJson.slice(0, budgetChars);
-                // Try to parse back; if not valid, wrap as raw truncated output
-                return {
-                    _truncated_raw: truncatedStr,
-                    _meta: {
-                        tokensEstimate: estimated,
-                        budget,
-                        truncated: true,
-                        hint: `Output exceeded budget of ${budget} tokens. Data was hard-truncated.`
-                    }
-                };
-            }
-
-            return {
-                ...pruned,
-                _meta: {
-                    ...(pruned._meta || {}),
-                    tokensEstimate: estimateTokens(JSON.stringify(pruned)),
-                    budget,
-                    truncated: true,
-                    hint: `Output exceeded budget of ${budget} tokens. Some fields were pruned.`
-                }
-            };
-        }
-
-        // Non-object data: hard truncate the serialized string
-        const truncated = serialized.slice(0, budgetChars);
-        return {
-            _truncated_raw: truncated,
-            _meta: {
-                tokensEstimate: estimated,
-                budget,
-                truncated: true,
-                hint: `Output exceeded budget of ${budget} tokens. Data was hard-truncated.`
-            }
-        };
-    }
-
-    // Under budget — attach meta info
-    if (typeof data === 'object' && data !== null && !Array.isArray(data)) {
-        return {
-            ...data,
-            _meta: {
-                ...(data._meta || {}),
-                tokensEstimate: estimated,
-                budget,
-                truncated: false
-            }
-        };
-    }
-
-    return data;
-}
-
-export function outputSuccess(data: any, budget?: number) {
-    const finalData = applyBudget(data, budget);
-    const output = JSON.stringify(finalData, null, process.argv.includes('--pretty') ? 2 : undefined) + '\n';
-    process.stdout.write(output, () => process.exit(0));
-}
-
-export function outputError(error: unknown, code = 1) {
-    const message = error instanceof Error ? error.message : String(error);
-    const errorObj = {
-        error: message,
-        code: code
-    };
-
-    process.stderr.write(JSON.stringify(errorObj, null, process.argv.includes('--pretty') ? 2 : undefined) + '\n');
-    process.exit(code);
-}
+/**
+ * Core output management for ai-xray v2.0.0
+ * Guarantees Machine-First protocol: Pure JSON on stdout, structured error on stderr.
+ */
+
+let prettyMode = false;
+
+export function setPrettyMode(enabled: boolean): void {
+    prettyMode = enabled;
+}
+
+export function isPrettyMode(): boolean {
+    return prettyMode;
+}
+
+export function outputSuccess(data: unknown): void {
+    const output = JSON.stringify(data, null, prettyMode ? 2 : undefined) + '\n';
+    process.stdout.write(output);
+    process.exit(0);
+}
+
+export function outputError(error: unknown, code = 1): void {
+    const message = error instanceof Error ? error.message : String(error);
+    const errorObj = {
+        error: message,
+        code
+    };
+    const output = JSON.stringify(errorObj, null, prettyMode ? 2 : undefined) + '\n';
+    process.stderr.write(output);
+    process.exit(code);
+}
+
+export function outputJson(data: unknown): void {
+    const output = JSON.stringify(data, null, prettyMode ? 2 : undefined) + '\n';
+    process.stdout.write(output);
+}

package/src/utils/timer.ts

@@ -0,0 +1,75 @@
+export class Timer {
+    private startTime: number = 0;
+    private endTime: number = 0;
+    private running: boolean = false;
+
+    start(): void {
+        this.startTime = performance.now();
+        this.running = true;
+    }
+
+    stop(): number {
+        if (!this.running) {
+            throw new Error('Timer is not running');
+        }
+        this.endTime = performance.now();
+        this.running = false;
+        return this.elapsed();
+    }
+
+    elapsed(): number {
+        if (this.running) {
+            return performance.now() - this.startTime;
+        }
+        return this.endTime - this.startTime;
+    }
+
+    reset(): void {
+        this.startTime = 0;
+        this.endTime = 0;
+        this.running = false;
+    }
+
+    static measure<T>(fn: () => T | Promise<T>): Promise<{ result: T; elapsed_ms: number }> {
+        return (async () => {
+            const timer = new Timer();
+            timer.start();
+            const result = await fn();
+            const elapsed = timer.stop();
+            return { result, elapsed_ms: elapsed };
+        })();
+    }
+
+    static async measureAsync<T>(fn: () => Promise<T>): Promise<{ result: T; elapsed_ms: number }> {
+        const timer = new Timer();
+        timer.start();
+        const result = await fn();
+        const elapsed = timer.stop();
+        return { result, elapsed_ms: elapsed };
+    }
+}
+
+export function median(values: number[]): number {
+    if (values.length === 0) return 0;
+    const sorted = [...values].sort((a, b) => a - b);
+    const mid = Math.floor(sorted.length / 2);
+    return sorted.length % 2 === 0
+        ? (sorted[mid - 1] + sorted[mid]) / 2
+        : sorted[mid];
+}
+
+export function mean(values: number[]): number {
+    if (values.length === 0) return 0;
+    return values.reduce((a, b) => a + b, 0) / values.length;
+}
+
+export function p95(values: number[]): number {
+    if (values.length === 0) return 0;
+    const sorted = [...values].sort((a, b) => a - b);
+    const index = Math.ceil(sorted.length * 0.95) - 1;
+    return sorted[index];
+}
+
+export function sum(values: number[]): number {
+    return values.reduce((a, b) => a + b, 0);
+}

package/tests/bench.test.ts

@@ -0,0 +1,13 @@
+import { describe, it, expect } from 'vitest';
+import { bench, BenchResult } from '../src/commands/bench';
+
+describe('bench', () => {
+    it('should be a function', () => {
+        expect(typeof bench).toBe('function');
+    });
+
+    it('should accept options parameter', () => {
+        // Just verify the function accepts the expected parameters
+        expect(bench).toBeDefined();
+    });
+});

package/tests/client.test.ts

@@ -0,0 +1,37 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { loadConfig, ProviderConfig } from '../src/client';
+
+describe('client', () => {
+    beforeEach(() => {
+        delete process.env.AI_XRAY_API_KEY;
+        delete process.env.AI_XRAY_BASE_URL;
+        delete process.env.AI_XRAY_MODEL;
+    });
+
+    it('should use default values when no env vars set', () => {
+        const config = loadConfig();
+        expect(config.baseUrl).toBe('https://api.openai.com/v1');
+        expect(config.model).toBe('gpt-4o');
+        expect(config.apiKey).toBeUndefined();
+    });
+
+    it('should use environment variables when set', () => {
+        process.env.AI_XRAY_API_KEY = 'test-key';
+        process.env.AI_XRAY_BASE_URL = 'https://custom.api.com/v1';
+        process.env.AI_XRAY_MODEL = 'gpt-5';
+
+        const config = loadConfig();
+        expect(config.apiKey).toBe('test-key');
+        expect(config.baseUrl).toBe('https://custom.api.com/v1');
+        expect(config.model).toBe('gpt-5');
+    });
+
+    it('should use provider config when specified', () => {
+        // This test requires a config file, so we just test the function exists
+        expect(loadConfig).toBeDefined();
+    });
+
+    it('should loadConfig be a function', () => {
+        expect(typeof loadConfig).toBe('function');
+    });
+});

package/tests/compare.test.ts

@@ -0,0 +1,24 @@
+import { describe, it, expect } from 'vitest';
+import { compare, CompareResult } from '../src/commands/compare';
+
+describe('compare', () => {
+    it('should be a function', () => {
+        expect(typeof compare).toBe('function');
+    });
+
+    it('should return result with expected structure', async () => {
+        const result = await compare(['openai'], { prompt: 'Say hi' });
+        
+        expect(result).toHaveProperty('prompt');
+        expect(result).toHaveProperty('results');
+        
+        expect(result.prompt).toBe('Say hi');
+        expect(Array.isArray(result.results)).toBe(true);
+    });
+
+    it('should accept custom prompt', async () => {
+        const customPrompt = 'Write a haiku';
+        const result = await compare(['openai'], { prompt: customPrompt });
+        expect(result.prompt).toBe(customPrompt);
+    });
+});

package/tests/http.test.ts

@@ -0,0 +1,12 @@
+import { describe, it, expect } from 'vitest';
+import { request } from '../src/utils/http';
+
+describe('http', () => {
+    it('should be defined', () => {
+        expect(request).toBeDefined();
+    });
+
+    it('should be a function', () => {
+        expect(typeof request).toBe('function');
+    });
+});

package/tests/id.test.ts

@@ -0,0 +1,13 @@
+import { describe, it, expect } from 'vitest';
+import { identify } from '../src/commands/id';
+
+describe('id', () => {
+    it('should be a function', () => {
+        expect(typeof identify).toBe('function');
+    });
+
+    it('should accept config parameter', () => {
+        // Just verify the function exists and accepts expected params
+        expect(identify).toBeDefined();
+    });
+});

package/tests/ping.test.ts

@@ -0,0 +1,12 @@
+import { describe, it, expect } from 'vitest';
+import { ping } from '../src/commands/ping';
+
+describe('ping', () => {
+    it('should be a function', () => {
+        expect(typeof ping).toBe('function');
+    });
+
+    it('should accept config parameter', () => {
+        expect(ping).toBeDefined();
+    });
+});

package/tests/probe.test.ts

@@ -0,0 +1,13 @@
+import { describe, it, expect } from 'vitest';
+import { probe, ProbeResult } from '../src/commands/probe';
+
+describe('probe', () => {
+    it('should be a function', () => {
+        expect(typeof probe).toBe('function');
+    });
+
+    it('should accept config parameter', () => {
+        // Just verify the function exists and accepts expected params
+        expect(probe).toBeDefined();
+    });
+});

package/tests/tokenize.test.ts

@@ -0,0 +1,32 @@
+import { describe, it, expect } from 'vitest';
+import { countTokens } from '../src/commands/tokenize';
+
+describe('tokenize', () => {
+    it('should be a function', () => {
+        expect(typeof countTokens).toBe('function');
+    });
+
+    it('should count characters accurately', async () => {
+        const result = await countTokens('Hello world');
+        expect(result.characters).toBe(11);
+    });
+
+    it('should count words accurately', async () => {
+        const result = await countTokens('Hello world from AI');
+        expect(result.words).toBe(4);
+    });
+
+    it('should estimate tokens within reasonable range', async () => {
+        const text = 'The quick brown fox jumps over the lazy dog';
+        const result = await countTokens(text);
+        expect(result.estimated_tokens).toBeGreaterThan(0);
+        expect(result.estimated_tokens).toBeLessThan(50);
+    });
+
+    it('should compute cost estimate for known models', async () => {
+        const result = await countTokens('Hello world', { model: 'gpt-4o' });
+        expect(result.cost_estimate).toBeDefined();
+        expect(result.cost_estimate?.model).toBe('gpt-4o');
+        expect(result.cost_estimate?.input_cost_usd).toBeGreaterThan(0);
+    });
+});

package/tsup.config.ts

@@ -1,11 +1,11 @@
-import { defineConfig } from 'tsup';
-
-export default defineConfig({
-    entry: ['src/cli.ts'],
-    format: ['cjs'],
-    target: 'node18',
-    clean: true,
-    minify: true,
-    // We bundle everything so there are zero dependencies at runtime
-    noExternal: [/(.*)/],
-});
+import { defineConfig } from 'tsup';
+
+export default defineConfig({
+    entry: ['src/cli.ts'],
+    format: ['cjs'],
+    target: 'node18',
+    clean: true,
+    sourcemap: true,
+    // We bundle everything so there are zero dependencies at runtime
+    noExternal: [/(.*)/],
+});

package/vitest.config.ts

@@ -0,0 +1,13 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+    test: {
+        globals: true,
+        environment: 'node',
+        include: ['tests/**/*.test.ts'],
+        coverage: {
+            provider: 'v8',
+            reporter: ['text', 'json', 'html'],
+        },
+    },
+});