@bedrockio/ai 0.3.0 → 0.4.2

package/dist/esm/xai.js

@@ -0,0 +1,10 @@
+import { OpenAiClient } from './openai.js';
+export class XAiClient extends OpenAiClient {
+    static DEFAULT_MODEL = 'grok-4-fast';
+    constructor(options) {
+        super({
+            ...options,
+            baseURL: 'https://api.x.ai/v1',
+        });
+    }
+}

package/eslint.config.js

@@ -0,0 +1,2 @@
+import { nodeImports, recommended } from '@bedrockio/eslint-plugin';
+export default [recommended, nodeImports];

package/package.json

@@ -1,20 +1,24 @@
 {
   "name": "@bedrockio/ai",
-  "version": "0.3.0",
+  "version": "0.4.2",
   "description": "Bedrock wrapper for common AI chatbots.",
   "type": "module",
   "scripts": {
-    "test": "node --no-warnings node_modules/.bin/jest",
-    "types": "tsc",
+    "test": "vitest run",
+    "test:live": "node --env-file=.env.live ./node_modules/.bin/vitest run",
     "lint": "eslint",
     "build": "scripts/build",
     "eject": "scripts/eject",
-    "prepublish": "yarn build && yarn types"
+    "build:cjs": "tsc -p tsconfig.cjs.json",
+    "build:esm": "tsc -p tsconfig.esm.json && tsc-alias -f -p tsconfig.esm.json",
+    "build:types": "tsc -p tsconfig.types.json",
+    "prepublish": "yarn build"
   },
   "types": "types/index.d.ts",
   "main": "./dist/cjs/index.js",
   "exports": {
     ".": {
+      "types": "./types/index.d.ts",
       "import": "./src/index.js",
       "require": "./dist/cjs/index.js"
     }
@@ -31,27 +35,24 @@
     "url": "https://github.com/bedrockio/router"
   },
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.36.3",
+    "@anthropic-ai/sdk": "^0.65.0",
     "@google/generative-ai": "^0.21.0",
     "glob": "^11.0.1",
     "mustache": "^4.2.0",
-    "openai": "^4.83.0"
+    "openai": "^6.1.0",
+    "partial-json": "^0.1.7"
   },
   "devDependencies": {
-    "@babel/cli": "^7.26.4",
-    "@babel/core": "^7.26.0",
-    "@babel/eslint-parser": "^7.26.5",
-    "@babel/preset-env": "^7.26.0",
+    "@bedrockio/eslint-plugin": "^1.2.2",
     "@bedrockio/prettier-config": "^1.0.2",
-    "eslint": "^8.33.0",
-    "eslint-plugin-bedrock": "^1.0.27",
-    "eslint-plugin-import": "^2.31.0",
-    "jest": "^29.7.0",
-    "prettier-eslint": "^16.3.0",
-    "typescript": "^5.7.3"
+    "@bedrockio/yada": "^1.8.0",
+    "eslint": "^9.36.0",
+    "tsc-alias": "^1.8.16",
+    "typescript": "^5.9.3",
+    "vitest": "^3.2.4"
   },
   "volta": {
-    "node": "22.12.0",
+    "node": "22.20.0",
     "yarn": "1.22.22"
   }
 }

package/src/BaseClient.js

@@ -1,195 +1,288 @@
-import Mustache from 'mustache';
-
-import { loadTemplates, loadTemplate } from './util.js';
-
-const MESSAGES_REG = /(?:^|\n)-{3,}\s*(\w+)\s*-{3,}(.*?)(?=\n-{3,}|$)/gs;
+import { parseCode } from './utils/code.js';
+import { createMessageExtractor } from './utils/json.js';
+import { loadTemplates, renderTemplate } from './utils/templates.js';
 
 export default class BaseClient {
   constructor(options) {
-    this.options = options;
+    this.options = {
+      // @ts-ignore
+      model: this.constructor.DEFAULT_MODEL,
+      ...options,
+    };
     this.templates = null;
   }
 
+  // Public
+
   /**
-   * Interpolates vars into the provided template and
-   * runs the chat completion. The "output" option may
-   * be omitted and will default to `"text"`.
-   * {@link https://github.com/bedrockio/ai?tab=readme-ov-file#bedrockioai Documentation}
+   * Interpolates vars into the provided template as instructions and runs the
+   * prompt.
    *
-   * @param {object} options
-   * @param {string} options.model - The model to use.
-   * @param {"raw" | "text" | "json" | "messages"} [options.output] - The output to use.
-   * @param {Object.<string, any>} [options.other] - Additional props
-   *             will be interpolated in the template.
+   * @param {PromptOptions} options
    */
   async prompt(options) {
-    options = {
-      ...this.options,
-      ...options,
-    };
+    options = await this.normalizeOptions(options);
 
-    const messages = await this.getMessages(options);
-    return await this.getCompletion({
-      ...options,
-      messages,
-    });
+    const { input, output, stream, schema } = options;
+
+    const response = await this.runPrompt(options);
+
+    if (!stream) {
+      this.debug('Response:', response);
+    }
+
+    if (output === 'raw') {
+      return response;
+    }
+
+    let result;
+    if (schema) {
+      result = this.getStructuredResponse(response);
+
+      // @ts-ignore
+      if (options.hasWrappedSchema) {
+        result = result.items;
+      }
+    } else if (output === 'json') {
+      result = JSON.parse(parseCode(this.getTextResponse(response)));
+    } else {
+      result = parseCode(this.getTextResponse(response));
+    }
+
+    if (output === 'messages') {
+      return {
+        result,
+        ...this.getMessagesResponse(input, response),
+      };
+    } else {
+      return result;
+    }
   }
 
   /**
    * Streams the prompt response.
+   *
+   * @param {PromptOptions & StreamOptions} options
    * @returns {AsyncIterator}
    */
   async *stream(options) {
-    const stream = await this.getStream(options);
+    options = await this.normalizeOptions(options);
 
-    let started = false;
+    const extractor = this.getMessageExtractor(options);
 
-    // @ts-ignore
-    for await (const chunk of stream) {
-      const resolved = this.getStreamedChunk(chunk, started);
-      started = true;
+    try {
+      const stream = await this.runStream(options);
 
       // @ts-ignore
-      if (resolved) {
-        yield resolved;
+      for await (let event of stream) {
+        this.debug('Event:', event);
+
+        event = this.normalizeStreamEvent(event);
+
+        if (event) {
+          yield event;
+        }
+
+        const extractedMessages = extractor?.(event) || [];
+
+        for (let message of extractedMessages) {
+          const { key, delta, text, done } = message;
+
+          let extractEvent;
+          if (done) {
+            extractEvent = {
+              type: 'extract:done',
+              text,
+              key,
+            };
+          } else {
+            extractEvent = {
+              type: 'extract:delta',
+              delta,
+              key,
+            };
+          }
+
+          this.debug('Extract:', extractEvent);
+
+          yield extractEvent;
+        }
       }
+    } catch (error) {
+      const { message, code } = error;
+      yield {
+        type: 'error',
+        code,
+        message,
+      };
     }
   }
 
-  async getMessages(options) {
-    const { text } = options;
+  async buildTemplate(options) {
     const template = await this.resolveTemplate(options);
+    return renderTemplate(template, options);
+  }
 
-    if (template) {
-      const raw = render(template, options);
+  // Protected
 
-      const messages = [];
-      for (let match of raw.matchAll(MESSAGES_REG)) {
-        const [, role, content] = match;
-        messages.push({
-          role: role.toLowerCase(),
-          content: content.trim(),
-        });
-      }
+  runPrompt(options) {
+    void options;
+    throw new Error('Method not implemented.');
+  }
 
-      if (!messages.length) {
-        messages.push({
-          role: 'user',
-          content: raw.trim(),
-        });
+  runStream(options) {
+    void options;
+    throw new Error('Method not implemented.');
+  }
+
+  getTextResponse(response) {
+    void response;
+    throw new Error('Method not implemented.');
+  }
+
+  /**
+   * @returns {Object}
+   */
+  getStructuredResponse(response) {
+    void response;
+    throw new Error('Method not implemented.');
+  }
+
+  /**
+   * @returns {Object}
+   */
+  getMessagesResponse(input, response) {
+    void response;
+    throw new Error('Method not implemented.');
+  }
+
+  /**
+   * @returns {Object}
+   */
+  normalizeStreamEvent(event) {
+    void event;
+    throw new Error('Method not implemented.');
+  }
+
+  // Private
+
+  async normalizeOptions(options) {
+    options = {
+      input: '',
+      output: 'text',
+      ...this.options,
+      ...options,
+    };
+
+    options.input = this.normalizeInput(options);
+    options.schema = this.normalizeSchema(options);
+    options.instructions ||= await this.resolveInstructions(options);
+
+    return options;
+  }
+
+  normalizeInput(options) {
+    let { input = '', output } = options;
+
+    if (typeof input === 'string') {
+      if (output === 'json') {
+        input += '\nOutput only valid JSON.';
       }
 
-      return messages;
-    } else if (text) {
-      return [
+      input = [
         {
           role: 'user',
-          content: text,
+          content: input,
         },
       ];
-    } else {
-      throw new Error('No input provided.');
     }
-  }
 
-  async buildTemplate(options) {
-    const template = await this.resolveTemplate(options);
-    return render(template, options);
+    return input;
   }
 
-  async loadTemplates() {
-    const { templates } = this.options;
-    this.templates ||= await loadTemplates(templates);
-  }
+  normalizeSchema(options) {
+    let { schema } = options;
 
-  async resolveTemplate(options) {
-    const { template, file } = options;
-    if (template) {
-      return template;
-    } else if (file?.endsWith('.md')) {
-      return await loadTemplate(file);
-    } else if (file) {
-      await this.loadTemplates();
-      return this.templates[file];
+    if (!schema) {
+      return;
     }
-  }
 
-  async getStream(options) {
-    return await this.prompt({
-      ...options,
-      output: 'raw',
-      stream: true,
-    });
-  }
+    // Convert to JSON schema.
+    schema = schema.toJSON?.() || schema;
 
-  getCompletion(options) {
-    void options;
-    new Error('Method not implemented.');
-  }
+    if (schema?.type === 'array') {
+      schema = {
+        type: 'object',
+        properties: {
+          items: schema,
+        },
+        required: ['items'],
+        additionalProperties: false,
+      };
+      options.hasWrappedSchema = true;
+    }
 
-  getStreamedChunk(chunk, started) {
-    void chunk;
-    void started;
-    new Error('Method not implemented.');
+    return schema;
   }
-}
 
-function render(template, options) {
-  let params = {
-    ...options,
-    ...options.params,
-  };
+  getMessageExtractor(options) {
+    const { extractMessages } = options;
+    if (!extractMessages) {
+      return;
+    }
+    const messageExtractor = createMessageExtractor([extractMessages]);
+    return (event) => {
+      if (event?.type === 'delta') {
+        return messageExtractor(event.delta);
+      }
+    };
+  }
 
-  params = mapObjects(params);
-  params = wrapProxy(params);
-  return Mustache.render(template, params);
-}
+  debug(message, arg) {
+    if (this.options.debug) {
+      // TODO: replace with logger when opentelemetry is removed
+      // eslint-disable-next-line
+      console.debug(`${message}\n${JSON.stringify(arg, null, 2)}\n`);
+    }
+  }
 
-// Transform arrays and object to versions
-// that are more understandable in the context
-// of a template that may have meaningful whitespace.
-function mapObjects(params) {
-  const result = {};
-  for (let [key, value] of Object.entries(params)) {
-    if (Array.isArray(value)) {
-      value = mapArray(value);
-    } else if (typeof value === 'object') {
-      value = JSON.stringify(value, null, 2);
+  async resolveInstructions(options) {
+    if (options.template) {
+      const template = await this.resolveTemplate(options);
+      return renderTemplate(template, options);
     }
-    result[key] = value;
   }
-  return result;
-}
 
-function mapArray(arr) {
-  // Only map simple arrays of primitives.
-  if (typeof arr[0] === 'string') {
-    arr = arr
-      .map((el) => {
-        return `- ${el}`;
-      })
-      .join('\n');
+  async resolveTemplate(options) {
+    const { template } = options;
+    await this.loadTemplates();
+    return this.templates[template] || template;
   }
-  return arr;
-}
 
-// Wrap params with a proxy object that reports
-// as having all properties. If one is accessed
-// that does not exist then return the original
-// token. This way templates can be partially
-// interpolated and re-interpolated later.
-function wrapProxy(params) {
-  return new Proxy(params, {
-    has() {
-      return true;
-    },
-
-    get(target, prop) {
-      if (prop in target) {
-        return target[prop];
-      } else {
-        return `{{{${prop.toString()}}}}`;
-      }
-    },
-  });
+  async loadTemplates() {
+    const { templates } = this.options;
+    this.templates ||= await loadTemplates(templates);
+  }
 }
+
+/**
+ * @typedef {Object} PromptOptions
+ * @property {string|PromptMessage[]} input - Input to use.
+ * @property {string} [model] - The model to use.
+ * @property {boolean} stream - Stream response.
+ * @property {Object} [schema] - A JSON schema compatible object that defines the output shape.
+ * @property {"raw" | "text" | "json" | "messages"} [output] - The return value type.
+ * @property {Object} [params] - Params to be interpolated into the template.
+ *                               May also be passed as additional props to options.
+ */
+
+/**
+ * @typedef {Object} StreamOptions
+ * @property {string} [extractMessages] - Key in JSON response to extract a message stream from.
+ */
+
+/**
+ * @typedef {Object} PromptMessage
+ * @property {"system" | "user" | "assistant"} role
+ * @property {string} content
+ */