@braingrid/sdk

Drop-in wrapper for Vercel's ai SDK with optional integrations for json-guard, retry logic, cost tracking, and custom providers.

Features

• 100% Compatible: Preserves identical function signatures to Vercel's ai SDK
• Auto-imports: Peer dependency management ensures ai package is automatically installed
• Version Resilient: New exports from ai are automatically available via proxy
• Optional Integrations: All integrations are optional and no-op if not configured
  • JSON Guard for tool argument validation/repair
  • Retry logic with configurable backoff (shim for future @braingrid/llm-retry)
  • Cost tracking (shim for future @braingrid/cost-tracker)
  • Custom provider injection (fetch, baseURL, headers)
• Lifecycle Hooks: Monitor and debug with beforeCall, afterCall, onToken, onRetry, onCost, onGuard

Installation

pnpm add @braingrid/sdk
# The ai package is automatically installed as a peer dependency

Basic Usage

import { createAiSdk } from "@braingrid/sdk";

// Create SDK with default settings (identical to raw ai SDK)
const ai = createAiSdk();

// Use exactly like the Vercel ai SDK
const response = await ai.generateText({
  model: "gpt-4o",
  prompt: "Hello, world!",
});

// Stream text
const stream = await ai.streamText({
  model: "gpt-4o",
  prompt: "Tell me a story",
});

for await (const chunk of stream.textStream) {
  console.log(chunk.text);
}

Advanced Usage

With JSON Guard Integration

import { createAiSdk } from "@braingrid/sdk";
import { createGuard, ajvAdapter } from "@braingrid/json-guard";

const ai = createAiSdk({
  jsonGuard: {
    guard: createGuard(),
    resolveSchema: (event) => {
      // Return schema for tool validation
      if (event.toolName === "search") {
        return ajvAdapter().fromJsonSchema({
          type: "object",
          required: ["query"],
          properties: {
            query: { type: "string" },
            limit: { type: "integer", default: 10 },
          },
          additionalProperties: false,
        });
      }
      return undefined;
    },
  },
});

With Lifecycle Hooks

const ai = createAiSdk({
  hooks: {
    beforeCall: async ({ id, fn, args }) => {
      console.log(`[${id}] Calling ${fn}`, args);
    },
    afterCall: async ({ id, ok, latencyMs, usage }) => {
      console.log(`[${id}] Completed in ${latencyMs}ms`, { ok, usage });
    },
    onToken: ({ id, text, delta }) => {
      console.log(`[${id}] Token:`, delta);
    },
    onRetry: ({ id, attempt, waitMs, reason }) => {
      console.log(`[${id}] Retry attempt ${attempt} in ${waitMs}ms: ${reason}`);
    },
    onCost: ({ id, costUSD }) => {
      console.log(`[${id}] Cost: $${costUSD.toFixed(4)}`);
    },
    onGuard: ({ id, tool, ok, diagnostics }) => {
      console.log(`[${id}] Guard ${tool}:`, ok ? "passed" : "failed", diagnostics);
    },
  },
});

With Provider Configuration

const ai = createAiSdk({
  provider: {
    // Custom fetch implementation
    fetch: customFetch,
    // Base URL for API calls
    baseURL: "https://api.custom-provider.com",
    // Additional headers
    headers: {
      "X-API-Key": "your-api-key",
      "X-Organization": "your-org",
    },
  },
});

Complete Example

import { createAiSdk } from "@braingrid/sdk";
import { createGuard, ajvAdapter } from "@braingrid/json-guard";

const ai = createAiSdk({
  // JSON validation for tool calls
  jsonGuard: {
    guard: createGuard(),
    resolveSchema: (event) => {
      // Return appropriate schema based on tool
    },
  },

  // Custom provider settings
  provider: {
    baseURL: process.env.AI_API_URL,
    headers: {
      "X-API-Key": process.env.AI_API_KEY,
    },
  },

  // Retry configuration (requires @braingrid/llm-retry - OUT OF SCOPE)
  retry: {
    maxAttempts: 3,
    backoffMs: 1000,
  },

  // Cost tracking (requires @braingrid/cost-tracker - OUT OF SCOPE)
  cost: {
    pricing: {
      "gpt-4o": { promptCost: 0.01, completionCost: 0.03 },
    },
  },

  // Lifecycle hooks for monitoring
  hooks: {
    afterCall: async ({ latencyMs, usage }) => {
      metrics.record("ai.call", { latencyMs, ...usage });
    },
    onCost: ({ costUSD }) => {
      metrics.increment("ai.cost", costUSD);
    },
  },
});

// Use the SDK normally
const response = await ai.generateText({
  model: "gpt-4o",
  prompt: "Explain quantum computing",
});

API Reference

createAiSdk(options?)

Creates an AI SDK instance with optional integrations.

Options

• jsonGuard - Fully functional JSON Guard integration for tool validation
  • guard - Guard instance from @braingrid/json-guard
  • resolveSchema - Function to resolve schemas for tool validation
• retry - Retry configuration (planned for future @braingrid/llm-retry)
• logger - Logger configuration (planned for future @braingrid/logger)
• cost - Cost tracking configuration (planned for future @braingrid/cost-tracker)
• provider - Custom provider configuration
  • fetch - Custom fetch implementation
  • baseURL - Base URL for API calls
  • headers - Additional headers to inject
• hooks - Lifecycle hooks
  • beforeCall - Called before making an SDK call
  • afterCall - Called after an SDK call completes
  • onRetry - Called when a retry occurs
  • onToken - Called for each token in a stream
  • onCost - Called when cost is calculated
  • onGuard - Called when json-guard processes tool arguments

Returns

An AiSdk instance with the same API as Vercel's ai package, including:

• generateText - Generate text using an LLM
• streamText - Stream text from an LLM
• All other exports from the ai package (via proxy)

Feature Status

Fully Functional

• ✅ JSON Guard Integration - Validates and repairs tool arguments using @braingrid/json-guard
• ✅ Provider Configuration - Custom fetch, baseURL, and headers
• ✅ Lifecycle Hooks - Complete monitoring via beforeCall, afterCall, onToken, onGuard

Planned (Future Packages)

The following features are planned but not yet implemented:

• ⏳ @braingrid/llm-retry - Retry logic with configurable backoff
• ⏳ @braingrid/logger - Structured logging
• ⏳ @braingrid/cost-tracker - Token usage and cost tracking

Configuration for these planned features is accepted but will no-op gracefully until the packages are available.

License

Apache-2.0