@gammatech/aijsx 0.11.4-dev.2024-07-08 → 0.12.0-dev.2024-07-08

package/dist/index.d.mts

@@ -1,5 +1,5 @@
-import { L as LogChatCompletionRequest, A as AINode, C as ChatMessage, R as RenderContext, a as LogImplementation, S as SpanProcessor, b as ContextValues, c as Context, T as Tracer, d as ReadableSpan, e as SpanExporter, f as AIComponent, g as SpanAttributes, P as Prompt, h as ChatCompletionRequestPayloads, D as DebugMessage, E as EvaluatorResult, i as EvaluatorFn, j as PromptParsed, J as JSX } from './jsx-dev-runtime-vXV-PExh.mjs';
-export { y as AIElement, l as AIFragment, B as BoundLogger, n as ChatRole, u as CombinedLogger, s as ConsoleLogger, I as ImagePart, o as ImagePartProps, v as Literal, p as LogChatCompletionResponse, q as LogLevel, r as Logger, N as NoopLogImplementation, V as OutputParser, F as PropsOfAIComponent, w as RenderResult, z as Renderable, K as Span, G as SpanContext, M as SpanEvent, H as SpanStatus, Q as TracingContext, O as TracingContextKey, U as TracingContextManager, x as attachedContextSymbol, k as createAIElement, m as createContext, t as toDebugMessage } from './jsx-dev-runtime-vXV-PExh.mjs';
+import { L as LogChatCompletionRequest, A as AINode, C as ChatMessage, R as RenderContext, a as LogImplementation, S as SpanProcessor, b as ContextValues, c as Context, T as Tracer, d as ReadableSpan, e as SpanAttributes, f as SpanExporter, g as AIComponent, P as Prompt, h as ChatCompletionRequestPayloads, D as DebugMessage, E as EvaluatorResult, i as EvaluatorFn, j as PromptParsed, J as JSX } from './jsx-dev-runtime-S--1HQKU.mjs';
+export { y as AIElement, l as AIFragment, B as BoundLogger, n as ChatRole, u as CombinedLogger, s as ConsoleLogger, I as ImagePart, o as ImagePartProps, v as Literal, p as LogChatCompletionResponse, q as LogLevel, r as Logger, N as NoopLogImplementation, V as OutputParser, F as PropsOfAIComponent, w as RenderResult, z as Renderable, K as Span, G as SpanContext, M as SpanEvent, H as SpanStatus, Q as TracingContext, O as TracingContextKey, U as TracingContextManager, x as attachedContextSymbol, k as createAIElement, m as createContext, t as toDebugMessage } from './jsx-dev-runtime-S--1HQKU.mjs';
 import { ZodObject, ZodRawShape, ZodTypeAny, ZodString, z } from 'zod';
 import { OpenAI } from 'openai';
 export { OpenAI as OpenAIClient } from 'openai';
@@ -100,19 +100,27 @@ declare class AITraceAPI {
 }
 declare const tracing: AITraceAPI;
 
+type SpanProcessingStatus = 'running' | 'processing' | 'done';
 declare class SpanTree {
     roots: SpanTreeNode[];
-    addSpan(span: ReadableSpan): boolean;
+    addSpan(span: ReadableSpan, status: SpanProcessingStatus): boolean;
     findNode(id: string): SpanTreeNode | null;
 }
 declare class SpanTreeNode {
-    readonly id: string;
-    readonly name: string;
+    span: ReadableSpan;
+    status: SpanProcessingStatus;
     readonly parent: SpanTreeNode | null;
     readonly children: SpanTreeNode[];
-    readonly attrs: Record<string, any>;
-    constructor(id: string, name: string, parent: SpanTreeNode | null, children: SpanTreeNode[], attrs: Record<string, any>);
+    childrenDonePromise: Promise<void>;
+    resolveChildrenDonePromise: () => void;
+    constructor(span: ReadableSpan, status: SpanProcessingStatus, parent?: SpanTreeNode | null, children?: SpanTreeNode[]);
+    get name(): string;
+    get id(): string;
+    get attrs(): SpanAttributes;
     addChild(child: SpanTreeNode): void;
+    descendents(fn: (node: SpanTreeNode) => void): void;
+    everyDescendent(fn: (node: SpanTreeNode) => boolean): boolean;
+    getRootByName(name: string): SpanTreeNode | null;
     getNearestAttribute<K>(key: string): K | null;
     getRootAttribute<K>(key: string): K | null;
 }
@@ -133,34 +141,16 @@ type CostFn = (tokensUsed: {
  * 3. sum totalCost on the root span
  * 4. sum totalUsage on the root span
  */
-declare abstract class EnrichingSpanProcessor implements SpanProcessor {
+declare class AISpanProcessor implements SpanProcessor {
     private readonly exporter;
-    protected finalSpans: {
-        status: 'running' | 'processing' | 'done';
-        span: ReadableSpan;
-    }[];
-    protected onAllDonePromise: Promise<void> | null;
-    protected onAllDonePromiseResolve: (() => void) | null;
     protected spanTree: SpanTree;
     constructor(exporter: SpanExporter);
     onStart(span: ReadableSpan): Promise<void>;
     onEnd(span: ReadableSpan): Promise<void>;
-    abstract onTraceComplete(rootSpan: ReadableSpan): Promise<void>;
-    abstract enrichSpan(span: ReadableSpan): Promise<ReadableSpan>;
     shutdown(): Promise<void> | undefined;
-}
-declare class AISpanProcessor extends EnrichingSpanProcessor {
-    protected totalCost: number;
-    protected totalUsage: {
-        prompt: number;
-        completion: number;
-        total: number;
-    };
-    onTraceComplete(rootSpan: ReadableSpan): Promise<void>;
+    onRootPromptDone(rootSpan: ReadableSpan): Promise<void>;
     enrichSpan(span: ReadableSpan): Promise<ReadableSpan>;
     protected addPromptKeys(span: ReadableSpan): Promise<ReadableSpan>;
-    protected updateCost(span: ReadableSpan): Promise<ReadableSpan>;
-    protected updateUsage(span: ReadableSpan): Promise<ReadableSpan>;
     protected evaluatePrompt(span: ReadableSpan): Promise<ReadableSpan>;
 }
 
@@ -321,4 +311,4 @@ type GoogleChatCompletionProps = {
 };
 declare function GoogleChatCompletion(props: GoogleChatCompletionProps, ctx: RenderContext): JSX.Element;
 
-export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EnrichingSpanProcessor, EvaluatorFn, EvaluatorResult, Fallback, type GetChatCompletionClientAndProvider, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, RetryLastErrorContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };
+export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EvaluatorFn, EvaluatorResult, Fallback, type GetChatCompletionClientAndProvider, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, RetryLastErrorContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };

package/dist/index.d.ts

@@ -1,5 +1,5 @@
-import { L as LogChatCompletionRequest, A as AINode, C as ChatMessage, R as RenderContext, a as LogImplementation, S as SpanProcessor, b as ContextValues, c as Context, T as Tracer, d as ReadableSpan, e as SpanExporter, f as AIComponent, g as SpanAttributes, P as Prompt, h as ChatCompletionRequestPayloads, D as DebugMessage, E as EvaluatorResult, i as EvaluatorFn, j as PromptParsed, J as JSX } from './jsx-dev-runtime-vXV-PExh.js';
-export { y as AIElement, l as AIFragment, B as BoundLogger, n as ChatRole, u as CombinedLogger, s as ConsoleLogger, I as ImagePart, o as ImagePartProps, v as Literal, p as LogChatCompletionResponse, q as LogLevel, r as Logger, N as NoopLogImplementation, V as OutputParser, F as PropsOfAIComponent, w as RenderResult, z as Renderable, K as Span, G as SpanContext, M as SpanEvent, H as SpanStatus, Q as TracingContext, O as TracingContextKey, U as TracingContextManager, x as attachedContextSymbol, k as createAIElement, m as createContext, t as toDebugMessage } from './jsx-dev-runtime-vXV-PExh.js';
+import { L as LogChatCompletionRequest, A as AINode, C as ChatMessage, R as RenderContext, a as LogImplementation, S as SpanProcessor, b as ContextValues, c as Context, T as Tracer, d as ReadableSpan, e as SpanAttributes, f as SpanExporter, g as AIComponent, P as Prompt, h as ChatCompletionRequestPayloads, D as DebugMessage, E as EvaluatorResult, i as EvaluatorFn, j as PromptParsed, J as JSX } from './jsx-dev-runtime-S--1HQKU.js';
+export { y as AIElement, l as AIFragment, B as BoundLogger, n as ChatRole, u as CombinedLogger, s as ConsoleLogger, I as ImagePart, o as ImagePartProps, v as Literal, p as LogChatCompletionResponse, q as LogLevel, r as Logger, N as NoopLogImplementation, V as OutputParser, F as PropsOfAIComponent, w as RenderResult, z as Renderable, K as Span, G as SpanContext, M as SpanEvent, H as SpanStatus, Q as TracingContext, O as TracingContextKey, U as TracingContextManager, x as attachedContextSymbol, k as createAIElement, m as createContext, t as toDebugMessage } from './jsx-dev-runtime-S--1HQKU.js';
 import { ZodObject, ZodRawShape, ZodTypeAny, ZodString, z } from 'zod';
 import { OpenAI } from 'openai';
 export { OpenAI as OpenAIClient } from 'openai';
@@ -100,19 +100,27 @@ declare class AITraceAPI {
 }
 declare const tracing: AITraceAPI;
 
+type SpanProcessingStatus = 'running' | 'processing' | 'done';
 declare class SpanTree {
     roots: SpanTreeNode[];
-    addSpan(span: ReadableSpan): boolean;
+    addSpan(span: ReadableSpan, status: SpanProcessingStatus): boolean;
     findNode(id: string): SpanTreeNode | null;
 }
 declare class SpanTreeNode {
-    readonly id: string;
-    readonly name: string;
+    span: ReadableSpan;
+    status: SpanProcessingStatus;
     readonly parent: SpanTreeNode | null;
     readonly children: SpanTreeNode[];
-    readonly attrs: Record<string, any>;
-    constructor(id: string, name: string, parent: SpanTreeNode | null, children: SpanTreeNode[], attrs: Record<string, any>);
+    childrenDonePromise: Promise<void>;
+    resolveChildrenDonePromise: () => void;
+    constructor(span: ReadableSpan, status: SpanProcessingStatus, parent?: SpanTreeNode | null, children?: SpanTreeNode[]);
+    get name(): string;
+    get id(): string;
+    get attrs(): SpanAttributes;
     addChild(child: SpanTreeNode): void;
+    descendents(fn: (node: SpanTreeNode) => void): void;
+    everyDescendent(fn: (node: SpanTreeNode) => boolean): boolean;
+    getRootByName(name: string): SpanTreeNode | null;
     getNearestAttribute<K>(key: string): K | null;
     getRootAttribute<K>(key: string): K | null;
 }
@@ -133,34 +141,16 @@ type CostFn = (tokensUsed: {
  * 3. sum totalCost on the root span
  * 4. sum totalUsage on the root span
  */
-declare abstract class EnrichingSpanProcessor implements SpanProcessor {
+declare class AISpanProcessor implements SpanProcessor {
     private readonly exporter;
-    protected finalSpans: {
-        status: 'running' | 'processing' | 'done';
-        span: ReadableSpan;
-    }[];
-    protected onAllDonePromise: Promise<void> | null;
-    protected onAllDonePromiseResolve: (() => void) | null;
     protected spanTree: SpanTree;
     constructor(exporter: SpanExporter);
     onStart(span: ReadableSpan): Promise<void>;
     onEnd(span: ReadableSpan): Promise<void>;
-    abstract onTraceComplete(rootSpan: ReadableSpan): Promise<void>;
-    abstract enrichSpan(span: ReadableSpan): Promise<ReadableSpan>;
     shutdown(): Promise<void> | undefined;
-}
-declare class AISpanProcessor extends EnrichingSpanProcessor {
-    protected totalCost: number;
-    protected totalUsage: {
-        prompt: number;
-        completion: number;
-        total: number;
-    };
-    onTraceComplete(rootSpan: ReadableSpan): Promise<void>;
+    onRootPromptDone(rootSpan: ReadableSpan): Promise<void>;
     enrichSpan(span: ReadableSpan): Promise<ReadableSpan>;
     protected addPromptKeys(span: ReadableSpan): Promise<ReadableSpan>;
-    protected updateCost(span: ReadableSpan): Promise<ReadableSpan>;
-    protected updateUsage(span: ReadableSpan): Promise<ReadableSpan>;
     protected evaluatePrompt(span: ReadableSpan): Promise<ReadableSpan>;
 }
 
@@ -321,4 +311,4 @@ type GoogleChatCompletionProps = {
 };
 declare function GoogleChatCompletion(props: GoogleChatCompletionProps, ctx: RenderContext): JSX.Element;
 
-export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EnrichingSpanProcessor, EvaluatorFn, EvaluatorResult, Fallback, type GetChatCompletionClientAndProvider, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, RetryLastErrorContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };
+export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EvaluatorFn, EvaluatorResult, Fallback, type GetChatCompletionClientAndProvider, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, RetryLastErrorContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };

package/dist/index.js

@@ -40,7 +40,6 @@ __export(src_exports, {
   CombinedLogger: () => CombinedLogger,
   ConsoleLogger: () => ConsoleLogger,
   DefaultMaxRetriesContext: () => DefaultMaxRetriesContext,
-  EnrichingSpanProcessor: () => EnrichingSpanProcessor,
   Fallback: () => Fallback,
   GoogleChatCompletion: () => GoogleChatCompletion,
   GoogleClientContext: () => GoogleClientContext,
@@ -1631,33 +1630,17 @@ async function* Fallback({ shouldFallback = () => true, fallback, children }, ct
 // src/tracing/SpanTree.ts
 var SpanTree = class {
   roots = [];
-  addSpan(span) {
+  addSpan(span, status) {
     const parentId = span.parentSpanId;
     if (!parentId) {
-      this.roots.push(
-        new SpanTreeNode(
-          span.spanContext.spanId,
-          span.name,
-          null,
-          [],
-          span.attributes
-        )
-      );
+      this.roots.push(new SpanTreeNode(span, status));
       return true;
     }
     const parent = this.findNode(parentId);
     if (!parent) {
       return false;
     }
-    parent.addChild(
-      new SpanTreeNode(
-        span.spanContext.spanId,
-        span.name,
-        parent,
-        [],
-        span.attributes
-      )
-    );
+    parent.addChild(new SpanTreeNode(span, status, parent));
     return true;
   }
   findNode(id) {
@@ -1683,16 +1666,59 @@ function findSpanNode(root, id) {
   return null;
 }
 var SpanTreeNode = class {
-  constructor(id, name, parent, children, attrs) {
-    this.id = id;
-    this.name = name;
+  constructor(span, status, parent = null, children = []) {
+    this.span = span;
+    this.status = status;
     this.parent = parent;
     this.children = children;
-    this.attrs = attrs;
+    this.childrenDonePromise = new Promise((resolve) => {
+      this.resolveChildrenDonePromise = resolve;
+    });
+  }
+  childrenDonePromise;
+  resolveChildrenDonePromise = () => {
+  };
+  get name() {
+    return this.span.name;
+  }
+  get id() {
+    return this.span.spanContext.spanId;
+  }
+  get attrs() {
+    return this.span.attributes;
   }
   addChild(child) {
     this.children.push(child);
   }
+  descendents(fn) {
+    for (const child of this.children) {
+      fn(child);
+      child.descendents(fn);
+    }
+  }
+  everyDescendent(fn) {
+    let result = true;
+    this.descendents((node) => {
+      if (!result) {
+        return;
+      }
+      if (!fn(node)) {
+        result = false;
+      }
+    });
+    return result;
+  }
+  getRootByName(name) {
+    let node = this;
+    let found = null;
+    while (node) {
+      if (node.name === name) {
+        found = node;
+      }
+      node = node.parent;
+    }
+    return found;
+  }
   getNearestAttribute(key) {
     return getNearestAttribute(this.parent, key);
   }
@@ -1734,65 +1760,69 @@ async function evaluatePrompt(prompt, variables, result) {
 }
 
 // src/tracing/AISpanProcessor.ts
-var EnrichingSpanProcessor = class {
+var AISpanProcessor = class {
   constructor(exporter) {
     this.exporter = exporter;
-    this.onAllDonePromise = new Promise((resolve) => {
-      this.onAllDonePromiseResolve = resolve;
-    });
   }
-  finalSpans = [];
-  onAllDonePromise = null;
-  onAllDonePromiseResolve = null;
   spanTree = new SpanTree();
   async onStart(span) {
-    this.spanTree.addSpan(span);
-    this.finalSpans.push({
-      status: "running",
-      span
-    });
+    this.spanTree.addSpan(span, "running");
   }
   async onEnd(span) {
-    const entry = this.finalSpans.find(
-      (a) => a.span.spanContext.spanId === span.spanContext.spanId
-    );
+    const entry = this.spanTree.findNode(span.spanContext.spanId);
     if (!entry) {
       throw new Error("onEnd called without onStart");
     }
-    const enrichedSpan = await this.enrichSpan(span);
     entry.status = "processing";
-    entry.span = enrichedSpan;
+    entry.span = await this.enrichSpan(span);
     entry.status = "done";
-    if (this.finalSpans.every((a) => a.status === "done")) {
-      this.onAllDonePromiseResolve?.();
+    const rootPrompt = entry.getRootByName("ai.prompt");
+    if (!rootPrompt) {
+      this.exporter.export([entry.span]);
+      return;
     }
-    if (span.parentSpanId == null) {
-      await this.onAllDonePromise;
-      await this.onTraceComplete(enrichedSpan);
+    if (rootPrompt.everyDescendent(({ status }) => status === "done")) {
+      rootPrompt.resolveChildrenDonePromise();
     }
-    this.exporter.export([enrichedSpan]);
+    if (span.spanContext.spanId === rootPrompt.id) {
+      await rootPrompt.childrenDonePromise;
+      await this.onRootPromptDone(entry.span);
+    }
+    this.exporter.export([entry.span]);
   }
   shutdown() {
     return this.exporter.shutdown?.();
   }
-};
-var AISpanProcessor = class extends EnrichingSpanProcessor {
-  totalCost = 0;
-  totalUsage = {
-    prompt: 0,
-    completion: 0,
-    total: 0
-  };
-  async onTraceComplete(rootSpan) {
-    rootSpan.attributes.totalCost = this.totalCost;
-    rootSpan.attributes.totalUsage = this.totalUsage;
+  async onRootPromptDone(rootSpan) {
+    let totalCost = 0;
+    const totalUsage = {
+      prompt: 0,
+      completion: 0,
+      total: 0
+    };
+    const root = this.spanTree.findNode(rootSpan.spanContext.spanId);
+    root?.descendents(({ span }) => {
+      if (!isChatCompletionSpan(span)) {
+        return;
+      }
+      const { cost, tokensUsed } = span.attributes;
+      if (typeof cost === "number") {
+        totalCost += cost;
+      }
+      if (!tokensUsed || tokensUsed.total == null || tokensUsed.completion == null || tokensUsed.prompt == null) {
+        return;
+      }
+      totalUsage.prompt += tokensUsed.prompt;
+      totalUsage.completion += tokensUsed.completion;
+      totalUsage.total += tokensUsed.total;
+    });
+    rootSpan.attributes.totalCost = totalCost;
+    rootSpan.attributes.totalUsage = totalUsage;
   }
   async enrichSpan(span) {
     const result1 = await this.evaluatePrompt(span);
-    const result2 = await this.updateCost(result1);
-    const result3 = await this.updateUsage(result2);
-    const result4 = await this.addPromptKeys(result3);
-    return result4;
+    const result2 = await this.addPromptKeys(result1);
+    return result2;
   }
   async addPromptKeys(span) {
     if (!isChatCompletionSpan(span) && !isPromptSpan(span)) {
@@ -1815,29 +1845,6 @@ var AISpanProcessor = class extends EnrichingSpanProcessor {
     };
     return final;
   }
-  async updateCost(span) {
-    if (!isChatCompletionSpan(span)) {
-      return span;
-    }
-    const { cost } = span.attributes;
-    if (typeof cost === "number") {
-      this.totalCost += cost;
-    }
-    return span;
-  }
-  async updateUsage(span) {
-    if (!isChatCompletionSpan(span)) {
-      return span;
-    }
-    const tokensUsed = span.attributes.tokensUsed;
-    if (!tokensUsed || tokensUsed.total == null || tokensUsed.completion == null || tokensUsed.prompt == null) {
-      return span;
-    }
-    this.totalUsage.prompt += tokensUsed.prompt;
-    this.totalUsage.completion += tokensUsed.completion;
-    this.totalUsage.total += tokensUsed.total;
-    return span;
-  }
   async evaluatePrompt(span) {
     if (!isPrompt2(span.attributes.prompt)) {
       return span;
@@ -2942,7 +2949,6 @@ var import_vertexai2 = require("@google-cloud/vertexai");
   CombinedLogger,
   ConsoleLogger,
   DefaultMaxRetriesContext,
-  EnrichingSpanProcessor,
   Fallback,
   GoogleChatCompletion,
   GoogleClientContext,

package/dist/index.mjs

@@ -1528,33 +1528,17 @@ async function* Fallback({ shouldFallback = () => true, fallback, children }, ct
 // src/tracing/SpanTree.ts
 var SpanTree = class {
   roots = [];
-  addSpan(span) {
+  addSpan(span, status) {
     const parentId = span.parentSpanId;
     if (!parentId) {
-      this.roots.push(
-        new SpanTreeNode(
-          span.spanContext.spanId,
-          span.name,
-          null,
-          [],
-          span.attributes
-        )
-      );
+      this.roots.push(new SpanTreeNode(span, status));
       return true;
     }
     const parent = this.findNode(parentId);
     if (!parent) {
       return false;
     }
-    parent.addChild(
-      new SpanTreeNode(
-        span.spanContext.spanId,
-        span.name,
-        parent,
-        [],
-        span.attributes
-      )
-    );
+    parent.addChild(new SpanTreeNode(span, status, parent));
     return true;
   }
   findNode(id) {
@@ -1580,16 +1564,59 @@ function findSpanNode(root, id) {
   return null;
 }
 var SpanTreeNode = class {
-  constructor(id, name, parent, children, attrs) {
-    this.id = id;
-    this.name = name;
+  constructor(span, status, parent = null, children = []) {
+    this.span = span;
+    this.status = status;
     this.parent = parent;
     this.children = children;
-    this.attrs = attrs;
+    this.childrenDonePromise = new Promise((resolve) => {
+      this.resolveChildrenDonePromise = resolve;
+    });
+  }
+  childrenDonePromise;
+  resolveChildrenDonePromise = () => {
+  };
+  get name() {
+    return this.span.name;
+  }
+  get id() {
+    return this.span.spanContext.spanId;
+  }
+  get attrs() {
+    return this.span.attributes;
   }
   addChild(child) {
     this.children.push(child);
   }
+  descendents(fn) {
+    for (const child of this.children) {
+      fn(child);
+      child.descendents(fn);
+    }
+  }
+  everyDescendent(fn) {
+    let result = true;
+    this.descendents((node) => {
+      if (!result) {
+        return;
+      }
+      if (!fn(node)) {
+        result = false;
+      }
+    });
+    return result;
+  }
+  getRootByName(name) {
+    let node = this;
+    let found = null;
+    while (node) {
+      if (node.name === name) {
+        found = node;
+      }
+      node = node.parent;
+    }
+    return found;
+  }
   getNearestAttribute(key) {
     return getNearestAttribute(this.parent, key);
   }
@@ -1631,65 +1658,69 @@ async function evaluatePrompt(prompt, variables, result) {
 }
 
 // src/tracing/AISpanProcessor.ts
-var EnrichingSpanProcessor = class {
+var AISpanProcessor = class {
   constructor(exporter) {
     this.exporter = exporter;
-    this.onAllDonePromise = new Promise((resolve) => {
-      this.onAllDonePromiseResolve = resolve;
-    });
   }
-  finalSpans = [];
-  onAllDonePromise = null;
-  onAllDonePromiseResolve = null;
   spanTree = new SpanTree();
   async onStart(span) {
-    this.spanTree.addSpan(span);
-    this.finalSpans.push({
-      status: "running",
-      span
-    });
+    this.spanTree.addSpan(span, "running");
   }
   async onEnd(span) {
-    const entry = this.finalSpans.find(
-      (a) => a.span.spanContext.spanId === span.spanContext.spanId
-    );
+    const entry = this.spanTree.findNode(span.spanContext.spanId);
     if (!entry) {
       throw new Error("onEnd called without onStart");
     }
-    const enrichedSpan = await this.enrichSpan(span);
     entry.status = "processing";
-    entry.span = enrichedSpan;
+    entry.span = await this.enrichSpan(span);
     entry.status = "done";
-    if (this.finalSpans.every((a) => a.status === "done")) {
-      this.onAllDonePromiseResolve?.();
+    const rootPrompt = entry.getRootByName("ai.prompt");
+    if (!rootPrompt) {
+      this.exporter.export([entry.span]);
+      return;
     }
-    if (span.parentSpanId == null) {
-      await this.onAllDonePromise;
-      await this.onTraceComplete(enrichedSpan);
+    if (rootPrompt.everyDescendent(({ status }) => status === "done")) {
+      rootPrompt.resolveChildrenDonePromise();
     }
-    this.exporter.export([enrichedSpan]);
+    if (span.spanContext.spanId === rootPrompt.id) {
+      await rootPrompt.childrenDonePromise;
+      await this.onRootPromptDone(entry.span);
+    }
+    this.exporter.export([entry.span]);
   }
   shutdown() {
     return this.exporter.shutdown?.();
   }
-};
-var AISpanProcessor = class extends EnrichingSpanProcessor {
-  totalCost = 0;
-  totalUsage = {
-    prompt: 0,
-    completion: 0,
-    total: 0
-  };
-  async onTraceComplete(rootSpan) {
-    rootSpan.attributes.totalCost = this.totalCost;
-    rootSpan.attributes.totalUsage = this.totalUsage;
+  async onRootPromptDone(rootSpan) {
+    let totalCost = 0;
+    const totalUsage = {
+      prompt: 0,
+      completion: 0,
+      total: 0
+    };
+    const root = this.spanTree.findNode(rootSpan.spanContext.spanId);
+    root?.descendents(({ span }) => {
+      if (!isChatCompletionSpan(span)) {
+        return;
+      }
+      const { cost, tokensUsed } = span.attributes;
+      if (typeof cost === "number") {
+        totalCost += cost;
+      }
+      if (!tokensUsed || tokensUsed.total == null || tokensUsed.completion == null || tokensUsed.prompt == null) {
+        return;
+      }
+      totalUsage.prompt += tokensUsed.prompt;
+      totalUsage.completion += tokensUsed.completion;
+      totalUsage.total += tokensUsed.total;
+    });
+    rootSpan.attributes.totalCost = totalCost;
+    rootSpan.attributes.totalUsage = totalUsage;
   }
   async enrichSpan(span) {
     const result1 = await this.evaluatePrompt(span);
-    const result2 = await this.updateCost(result1);
-    const result3 = await this.updateUsage(result2);
-    const result4 = await this.addPromptKeys(result3);
-    return result4;
+    const result2 = await this.addPromptKeys(result1);
+    return result2;
   }
   async addPromptKeys(span) {
     if (!isChatCompletionSpan(span) && !isPromptSpan(span)) {
@@ -1712,29 +1743,6 @@ var AISpanProcessor = class extends EnrichingSpanProcessor {
     };
     return final;
   }
-  async updateCost(span) {
-    if (!isChatCompletionSpan(span)) {
-      return span;
-    }
-    const { cost } = span.attributes;
-    if (typeof cost === "number") {
-      this.totalCost += cost;
-    }
-    return span;
-  }
-  async updateUsage(span) {
-    if (!isChatCompletionSpan(span)) {
-      return span;
-    }
-    const tokensUsed = span.attributes.tokensUsed;
-    if (!tokensUsed || tokensUsed.total == null || tokensUsed.completion == null || tokensUsed.prompt == null) {
-      return span;
-    }
-    this.totalUsage.prompt += tokensUsed.prompt;
-    this.totalUsage.completion += tokensUsed.completion;
-    this.totalUsage.total += tokensUsed.total;
-    return span;
-  }
   async evaluatePrompt(span) {
     if (!isPrompt2(span.attributes.prompt)) {
       return span;
@@ -2846,7 +2854,6 @@ export {
   CombinedLogger,
   ConsoleLogger,
   DefaultMaxRetriesContext,
-  EnrichingSpanProcessor,
   Fallback,
   GoogleChatCompletion,
   GoogleClientContext,

package/dist/{jsx-dev-runtime-vXV-PExh.d.mts → jsx-dev-runtime-S--1HQKU.d.mts}

@@ -354,4 +354,4 @@ declare const jsxs: typeof jsx;
 /** @hidden */
 declare const Fragment: typeof AIFragment;
 
-export { type AINode as A, BoundLogger as B, type ChatMessage as C, type DebugMessage as D, type EvaluatorResult as E, type PropsOfAIComponent as F, type SpanContext as G, type SpanStatus as H, ImagePart as I, JSX as J, type Span as K, type LogChatCompletionRequest as L, type SpanEvent as M, NoopLogImplementation as N, type TracingContextKey as O, type Prompt as P, type TracingContext as Q, type RenderContext as R, type SpanProcessor as S, type Tracer as T, type TracingContextManager as U, type OutputParser as V, jsx as W, jsxDEV as X, jsxs as Y, Fragment as Z, LogImplementation as a, type ContextValues as b, type Context as c, type ReadableSpan as d, type SpanExporter as e, type AIComponent as f, type SpanAttributes as g, type ChatCompletionRequestPayloads as h, type EvaluatorFn as i, type PromptParsed as j, createAIElement as k, AIFragment as l, createContext as m, type ChatRole as n, type ImagePartProps as o, type LogChatCompletionResponse as p, type LogLevel as q, type Logger as r, ConsoleLogger as s, toDebugMessage as t, CombinedLogger as u, type Literal as v, type RenderResult as w, attachedContextSymbol as x, type AIElement as y, type Renderable as z };
+export { type AINode as A, BoundLogger as B, type ChatMessage as C, type DebugMessage as D, type EvaluatorResult as E, type PropsOfAIComponent as F, type SpanContext as G, type SpanStatus as H, ImagePart as I, JSX as J, type Span as K, type LogChatCompletionRequest as L, type SpanEvent as M, NoopLogImplementation as N, type TracingContextKey as O, type Prompt as P, type TracingContext as Q, type RenderContext as R, type SpanProcessor as S, type Tracer as T, type TracingContextManager as U, type OutputParser as V, jsx as W, jsxDEV as X, jsxs as Y, Fragment as Z, LogImplementation as a, type ContextValues as b, type Context as c, type ReadableSpan as d, type SpanAttributes as e, type SpanExporter as f, type AIComponent as g, type ChatCompletionRequestPayloads as h, type EvaluatorFn as i, type PromptParsed as j, createAIElement as k, AIFragment as l, createContext as m, type ChatRole as n, type ImagePartProps as o, type LogChatCompletionResponse as p, type LogLevel as q, type Logger as r, ConsoleLogger as s, toDebugMessage as t, CombinedLogger as u, type Literal as v, type RenderResult as w, attachedContextSymbol as x, type AIElement as y, type Renderable as z };

package/dist/{jsx-dev-runtime-vXV-PExh.d.ts → jsx-dev-runtime-S--1HQKU.d.ts}

@@ -354,4 +354,4 @@ declare const jsxs: typeof jsx;
 /** @hidden */
 declare const Fragment: typeof AIFragment;
 
-export { type AINode as A, BoundLogger as B, type ChatMessage as C, type DebugMessage as D, type EvaluatorResult as E, type PropsOfAIComponent as F, type SpanContext as G, type SpanStatus as H, ImagePart as I, JSX as J, type Span as K, type LogChatCompletionRequest as L, type SpanEvent as M, NoopLogImplementation as N, type TracingContextKey as O, type Prompt as P, type TracingContext as Q, type RenderContext as R, type SpanProcessor as S, type Tracer as T, type TracingContextManager as U, type OutputParser as V, jsx as W, jsxDEV as X, jsxs as Y, Fragment as Z, LogImplementation as a, type ContextValues as b, type Context as c, type ReadableSpan as d, type SpanExporter as e, type AIComponent as f, type SpanAttributes as g, type ChatCompletionRequestPayloads as h, type EvaluatorFn as i, type PromptParsed as j, createAIElement as k, AIFragment as l, createContext as m, type ChatRole as n, type ImagePartProps as o, type LogChatCompletionResponse as p, type LogLevel as q, type Logger as r, ConsoleLogger as s, toDebugMessage as t, CombinedLogger as u, type Literal as v, type RenderResult as w, attachedContextSymbol as x, type AIElement as y, type Renderable as z };
+export { type AINode as A, BoundLogger as B, type ChatMessage as C, type DebugMessage as D, type EvaluatorResult as E, type PropsOfAIComponent as F, type SpanContext as G, type SpanStatus as H, ImagePart as I, JSX as J, type Span as K, type LogChatCompletionRequest as L, type SpanEvent as M, NoopLogImplementation as N, type TracingContextKey as O, type Prompt as P, type TracingContext as Q, type RenderContext as R, type SpanProcessor as S, type Tracer as T, type TracingContextManager as U, type OutputParser as V, jsx as W, jsxDEV as X, jsxs as Y, Fragment as Z, LogImplementation as a, type ContextValues as b, type Context as c, type ReadableSpan as d, type SpanAttributes as e, type SpanExporter as f, type AIComponent as g, type ChatCompletionRequestPayloads as h, type EvaluatorFn as i, type PromptParsed as j, createAIElement as k, AIFragment as l, createContext as m, type ChatRole as n, type ImagePartProps as o, type LogChatCompletionResponse as p, type LogLevel as q, type Logger as r, ConsoleLogger as s, toDebugMessage as t, CombinedLogger as u, type Literal as v, type RenderResult as w, attachedContextSymbol as x, type AIElement as y, type Renderable as z };

package/dist/jsx-dev-runtime.d.mts

@@ -1,2 +1,2 @@
-export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-vXV-PExh.mjs';
+export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-S--1HQKU.mjs';
 import 'zod';

package/dist/jsx-dev-runtime.d.ts

@@ -1,2 +1,2 @@
-export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-vXV-PExh.js';
+export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-S--1HQKU.js';
 import 'zod';

package/dist/jsx-runtime.d.mts

@@ -1,2 +1,2 @@
-export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-vXV-PExh.mjs';
+export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-S--1HQKU.mjs';
 import 'zod';

package/dist/jsx-runtime.d.ts

@@ -1,2 +1,2 @@
-export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-vXV-PExh.js';
+export { Z as Fragment, J as JSX, W as jsx, X as jsxDEV, Y as jsxs } from './jsx-dev-runtime-S--1HQKU.js';
 import 'zod';

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@gammatech/aijsx",
-  "version": "0.11.4-dev.2024-07-08",
+  "version": "0.12.0-dev.2024-07-08",
   "description": "Rewrite of aijsx",
   "author": "Jordan Garcia",
   "license": "MIT",