weflayr 0.3.2 → 0.5.1

package/index.d.ts

@@ -36,6 +36,12 @@ export interface MethodConfig {
   call: string;
   middleware?: MiddlewareFn;
   streamMiddleware?: StreamMiddlewareFn;
+  /**
+   * Property name on the response object that holds the async iterable stream.
+   * Use when the SDK returns `{ [streamPath]: AsyncIterable }` instead of a direct AsyncIterable.
+   * e.g. `'stream'` for AWS Bedrock's `ConverseStreamCommand`.
+   */
+  streamPath?: string;
 }
 
 /** Full configuration object passed to `weflayr_setup`. */

package/package.json

@@ -1,13 +1,17 @@
 {
   "name": "weflayr",
-  "version": "0.3.2",
+  "version": "0.5.1",
   "description": "Weflayr Node.js SDK — instrument any LLM client via JS Proxy",
   "main": "src/index.js",
   "types": "index.d.ts",
   "scripts": {
     "test": "node --test tests/index.test.js"
   },
-  "keywords": ["llm", "observability", "weflayr"],
+  "keywords": [
+    "llm",
+    "observability",
+    "weflayr"
+  ],
   "author": "Weflayr <contact@weflayr.com>",
   "license": "Elastic-2.0",
   "engines": {

package/src/instrument.js

@@ -87,6 +87,12 @@ function _wrapFn(fn, methodName, methodConfig) {
         return _wrapStream(result, methodName, tags, requestArgs, settings, startTime, methodConfig, eventId);
       }
 
+      if (methodConfig.streamPath && result && _isAsyncIterable(result[methodConfig.streamPath])) {
+        _sendEvent('stream_start', { event_id: eventId, method: methodName, tags });
+        const wrappedStream = _wrapStream(result[methodConfig.streamPath], methodName, tags, requestArgs, settings, startTime, methodConfig, eventId);
+        return { ...result, [methodConfig.streamPath]: wrappedStream };
+      }
+
       const middlewareData = methodConfig.middleware
         ? methodConfig.middleware(requestArgs, result) || {}
         : {};
@@ -166,14 +172,23 @@ async function* _wrapStream(stream, methodName, tags, requestArgs, settings, sta
 
 function _extractTags(args) {
   let tags = {};
-  const cleanArgs = args.map(arg => {
-    if (arg && typeof arg === 'object' && !Array.isArray(arg) && '__weflayr_tags' in arg) {
-      const { __weflayr_tags, ...rest } = arg;
-      tags = __weflayr_tags || {};
-      return rest;
+
+  function strip(obj) {
+    if (!obj || typeof obj !== 'object' || Array.isArray(obj)) return;
+    if ('__weflayr_tags' in obj) {
+      tags = obj.__weflayr_tags || {};
+      delete obj.__weflayr_tags;
     }
+    for (const key of Object.keys(obj)) {
+      strip(obj[key]);
+    }
+  }
+
+  const cleanArgs = args.map(arg => {
+    if (arg && typeof arg === 'object' && !Array.isArray(arg)) strip(arg);
     return arg;
   });
+
   return { tags, cleanArgs };
 }
 

package/src/providers/bedrock-runtime.d.ts

@@ -0,0 +1,6 @@
+import { MethodConfig, StreamAccumulator } from '../../../index';
+
+export function createStreamAccumulator(): StreamAccumulator;
+
+/** Ready-to-use method config for `client.send(new ConverseStreamCommand(...))`. */
+export const converseStream: MethodConfig;

package/src/providers/bedrock-runtime.js

@@ -0,0 +1,28 @@
+'use strict';
+
+function createStreamAccumulator() {
+  let input_tokens = null;
+  let output_tokens = null;
+
+  return {
+    onChunk(chunk) {
+      if (chunk.metadata?.usage) {
+        input_tokens = chunk.metadata.usage.inputTokens ?? null;
+        output_tokens = chunk.metadata.usage.outputTokens ?? null;
+        return true;
+      }
+      return false;
+    },
+    finalize() {
+      return { input_tokens, output_tokens };
+    },
+  };
+}
+
+const converseStream = {
+  call: 'send',
+  streamPath: 'stream',
+  streamMiddleware: createStreamAccumulator,
+};
+
+module.exports = { converseStream, createStreamAccumulator };