@gammatech/aijsx 0.15.3-dev.2024-10-21 → 0.16.0-dev.2024-10-28

package/dist/index.js

@@ -61,16 +61,14 @@ __export(src_exports, {
   Trace: () => Trace,
   UserMessage: () => UserMessage,
   VertexAI: () => import_vertexai2.VertexAI,
-  anthropicTokenizer: () => anthropicTokenizer,
   attachedContextSymbol: () => attachedContextSymbol,
-  computeUsage: () => computeUsage,
   createAIElement: () => createAIElement,
   createContext: () => createContext,
   createPrompt: () => createPrompt,
   createRenderContext: () => createRenderContext,
   evaluatePrompt: () => evaluatePrompt,
+  experimental__AnthropicPromptCacheDebugger: () => AnthropicPromptCacheDebugger,
   isPromptParsed: () => isPromptParsed2,
-  openaiTokenizer: () => openaiTokenizer,
   toDebugMessage: () => toDebugMessage,
   tracing: () => tracing
 });
@@ -245,7 +243,7 @@ var UserChatMessageBuilder = class {
     this.content.push(props);
     return this;
   }
-  async build() {
+  async build(opts) {
     const promises = this.content.map(
       async (part) => {
         if (typeof part === "string") {
@@ -266,7 +264,8 @@ var UserChatMessageBuilder = class {
     );
     return {
       role: "user",
-      content: await Promise.all(promises)
+      content: await Promise.all(promises),
+      cachePrompt: opts.cachePrompt
     };
   }
 };
@@ -284,34 +283,29 @@ var UserMessage = (props) => {
 var AssistantMessage = (props) => {
   return props.children;
 };
-
-// src/chat/tokenizer.ts
-var computeUsage = (messages, tokenizer2) => {
-  const promptMessages = [...messages];
-  let assistantMessage;
-  if (promptMessages[promptMessages.length - 1].role === "assistant") {
-    assistantMessage = promptMessages.pop();
-  }
-  const prompt = promptMessages.reduce((acc, m) => acc + tokenizer2(m), 0);
-  const completion = assistantMessage ? tokenizer2(assistantMessage) : 0;
-  return {
-    prompt,
-    completion,
-    total: prompt + completion
-  };
+var ChatMessageXmlAttributes = {
+  AssistantMessage: {
+    cachePrompt: true
+  },
+  UserMessage: {
+    cachePrompt: true
+  },
+  SystemMessage: {
+    cachePrompt: true
+  }
 };
 
 // src/createElement.ts
-function createAIElement(tag, props, ...children) {
+function createAIElement(tag2, props, ...children) {
   const propsToPass = {
     ...props ?? {},
     ...children.length === 0 ? {} : { children: children.length === 1 ? children[0] : children }
   };
   const result = {
-    tag,
+    tag: tag2,
     props: propsToPass,
     render: (ctx) => {
-      return tag(propsToPass, ctx);
+      return tag2(propsToPass, ctx);
     }
   };
   return result;
@@ -2044,56 +2038,14 @@ var errorToChatCompletionError = (error, requestData) => {
   );
 };
 
-// src/lib/openai/tokenizer.ts
-var import_js_tiktoken = require("js-tiktoken");
-var cl100kTokenizer = (0, import_js_tiktoken.getEncoding)("cl100k_base");
-var tokenizer = {
-  encode: (text) => cl100kTokenizer.encode(text),
-  decode: (tokens) => cl100kTokenizer.decode(tokens)
-};
-var TOKENS_PER_MESSAGE = 3;
-var textCost = (content) => {
-  return TOKENS_PER_MESSAGE + tokenizer.encode(content).length;
-};
-var COST_PER_LOW = 85;
-var COST_PER_512x512 = 170;
-var imageCost = (w, h) => {
-  const area = w * h;
-  return Math.ceil(area / (512 * 512)) * COST_PER_512x512 + COST_PER_LOW;
-};
-var openaiTokenizer = (message) => {
-  switch (message.role) {
-    case "assistant":
-    case "system":
-      return textCost(message.content || "");
-    case "user":
-      if (typeof message.content === "string") {
-        return textCost(message.content);
-      }
-      return message.content.reduce((acc, part) => {
-        if (part.type === "text") {
-          return acc + textCost(part.text);
-        } else {
-          if (part.image.detail === "low") {
-            return acc + 85;
-          }
-          if (part.image.dimensions) {
-            return acc + imageCost(
-              part.image.dimensions.width,
-              part.image.dimensions.height
-            );
-          }
-          return acc + imageCost(1024, 1024);
-        }
-      }, 0);
-  }
-};
-
 // src/chat/buildMessages.ts
 async function toXml(ctx, children) {
   const childrenXml = await ctx.render(children, {
     preserveTags: true,
-    renderedProps: IMAGE_RENDERED_PROPS
+    renderedProps: {
+      ...IMAGE_RENDERED_PROPS,
+      ...ChatMessageXmlAttributes
+    }
   });
   const topLevelTags = ["UserMessage", "AssistantMessage", "SystemMessage"];
   const chatMessageTags = [...topLevelTags, "ImagePart"];
@@ -2133,25 +2085,35 @@ async function buildChatMessages(ctx, children, opts) {
         throw new Error("Invalid User ChildNode, expecting Text or ImagePart");
       }
     }
-    return builder.build();
+    return builder.build({
+      cachePrompt: node.attributes.cachePrompt === true
+    });
   };
   return iterateChatMessageXml(nodes, async (node) => {
-    switch (node.nodeName) {
-      case "UserMessage":
-        return handleUserMessage(node);
-      case "SystemMessage":
-        return {
-          role: "system",
-          content: node.textContent
-        };
-      case "AssistantMessage":
-        return {
-          role: "assistant",
-          content: node.textContent
-        };
-      default:
-        throw new Error("Invalid top level chat message tags");
-    }
+    const cachePrompt = node.attributes.cachePrompt === true;
+    const getMessageWithoutCachePrompt = () => {
+      switch (node.nodeName) {
+        case "UserMessage":
+          return handleUserMessage(node);
+        case "SystemMessage":
+          return {
+            role: "system",
+            content: node.textContent
+          };
+        case "AssistantMessage":
+          return {
+            role: "assistant",
+            content: node.textContent
+          };
+        default:
+          throw new Error("Invalid top level chat message tags");
+      }
+    };
+    const message = await getMessageWithoutCachePrompt();
+    return cachePrompt ? {
+      ...message,
+      cachePrompt: true
+    } : message;
   });
 }
 
@@ -2251,6 +2213,9 @@ async function* OpenAIChatCompletionInner(props, ctx) {
   const chatCompletionRequest = {
     model: props.model,
     max_tokens: props.maxTokens,
+    stream_options: {
+      include_usage: true
+    },
     temperature: props.temperature,
     stop: props.stop,
     messages: openAIMessages,
@@ -2288,8 +2253,24 @@ async function* OpenAIChatCompletionInner(props, ctx) {
   }
   let finishReason = void 0;
   let content = "";
+  let tokensUsed = {
+    prompt: 0,
+    completion: 0,
+    total: 0,
+    cachedPromptTokensCreated: 0,
+    cachedPromptTokensRead: 0
+  };
   try {
     for await (const message of chatResponse) {
+      if (message.usage) {
+        tokensUsed = {
+          prompt: message.usage.prompt_tokens,
+          completion: message.usage.completion_tokens,
+          total: message.usage.total_tokens,
+          cachedPromptTokensCreated: 0,
+          cachedPromptTokensRead: message.usage.prompt_tokens_details?.cached_tokens || 0
+        };
+      }
       if (!message.choices || !message.choices[0]) {
         continue;
       }
@@ -2315,10 +2296,6 @@ async function* OpenAIChatCompletionInner(props, ctx) {
     role: "assistant",
     content
   };
-  const tokensUsed = computeUsage(
-    [...chatMessages, outputMessage],
-    openaiTokenizer
-  );
   const cost = costFn?.(props.model, tokensUsed) ?? void 0;
   const responseData = {
     ...logRequestData,
@@ -2376,32 +2353,6 @@ var import_openai3 = require("openai");
 
 // src/lib/anthropic/Anthropic.tsx
 var import_sdk = __toESM(require("@anthropic-ai/sdk"));
-
-// src/lib/anthropic/tokenizer.ts
-var import_tokenizer4 = require("@anthropic-ai/tokenizer");
-var DEFAULT_IMAGE_TOKEN_COST = 1334;
-var imageTokens = (w, h) => {
-  return Math.ceil(w * h / 750);
-};
-var anthropicTokenizer = (message) => {
-  if (message.role === "system") {
-    return (0, import_tokenizer4.countTokens)(message.content);
-  }
-  if (message.role === "assistant") {
-    return (0, import_tokenizer4.countTokens)(message.content);
-  }
-  return message.content.reduce((carry, item) => {
-    let tokens = 0;
-    if (item.type === "text") {
-      tokens = (0, import_tokenizer4.countTokens)(item.text);
-    } else {
-      tokens = item.image.dimensions ? imageTokens(item.image.dimensions.width, item.image.dimensions.height) : DEFAULT_IMAGE_TOKEN_COST;
-    }
-    return carry + tokens;
-  }, 0);
-};
-
-// src/lib/anthropic/Anthropic.tsx
 var defaultClient2 = null;
 var AnthropicClientContext = createContext(async () => {
   if (defaultClient2) {
@@ -2417,21 +2368,73 @@ var AnthropicClientContext = createContext(async () => {
   return defaultClient2;
 });
 var defaultMaxTokens = 4096;
+var populateCacheControl = (message, cachePrompt) => {
+  if (!cachePrompt) {
+    return message;
+  }
+  const content = message.content;
+  if (typeof content === "string") {
+    return {
+      role: message.role,
+      content: [
+        {
+          type: "text",
+          text: content,
+          // @ts-ignore - AnthropicSDK isn't up to date with this
+          cache_control: {
+            type: "ephemeral"
+          }
+        }
+      ]
+    };
+  }
+  if (content.length === 0) {
+    return message;
+  }
+  const lastcontent = {
+    ...content[content.length - 1],
+    cache_control: {
+      type: "ephemeral"
+    }
+  };
+  return {
+    role: message.role,
+    content: [...content.slice(0, -1), lastcontent]
+  };
+};
 var buildAnthropicMessages = (chatMesssages) => {
   let system = "";
   const messages = [];
-  chatMesssages.forEach(({ role, content }) => {
+  chatMesssages.forEach(({ role, content, cachePrompt }) => {
     if (role === "system") {
-      system = content;
+      if (cachePrompt) {
+        system = [
+          {
+            type: "text",
+            text: content,
+            // @ts-ignore - AnthropicSDK isn't up to date with this
+            cache_control: {
+              type: "ephemeral"
+            }
+          }
+        ];
+      } else {
+        system = content;
+      }
       return;
     }
     if (role === "user") {
       const userContent = content;
       if (userContent.length === 1 && userContent[0].type === "text") {
-        messages.push({
-          role,
-          content: userContent[0].text
-        });
+        messages.push(
+          populateCacheControl(
+            {
+              role,
+              content: userContent[0].text
+            },
+            cachePrompt
+          )
+        );
         return;
       }
       const c = userContent.map((part) => {
@@ -2456,17 +2459,27 @@ var buildAnthropicMessages = (chatMesssages) => {
           a.type === "image" || a.type === "text" && a.text.trim().length > 0
         )
       );
-      messages.push({
-        role,
-        content: c
-      });
+      messages.push(
+        populateCacheControl(
+          {
+            role,
+            content: c
+          },
+          cachePrompt
+        )
+      );
       return;
     }
     if (role === "assistant") {
-      messages.push({
-        role,
-        content
-      });
+      messages.push(
+        populateCacheControl(
+          {
+            role,
+            content
+          },
+          cachePrompt
+        )
+      );
       return;
     }
     throw new Error(`Invalid role: ${role}`);
@@ -2519,6 +2532,9 @@ function AnthropicChatCompletion(props, ctx) {
     }
   );
 }
+var ANTHROPIC_DEFAULT_HEADERS = {
+  "anthropic-beta": "prompt-caching-2024-07-31"
+};
 async function* AnthropicChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
@@ -2556,9 +2572,12 @@ async function* AnthropicChatCompletionInner(props, ctx) {
     stop_sequences: stopSequences,
     model: props.model
   };
-  const requestOptions = props.extraHeaders ? {
-    headers: props.extraHeaders
-  } : void 0;
+  const requestOptions = {
+    headers: {
+      ...ANTHROPIC_DEFAULT_HEADERS,
+      ...props.extraHeaders
+    }
+  };
   const chatCompletionRequestToLog = cleanChatCompletionRequest2(
     anthropicCompletionRequest
   );
@@ -2604,13 +2623,27 @@ async function* AnthropicChatCompletionInner(props, ctx) {
     throw err;
   }
   let content = "";
-  let outputUsage;
-  let inputUsage;
   let finishReason = null;
+  const tokensUsed = {
+    prompt: 0,
+    completion: 0,
+    total: 0,
+    cachedPromptTokensCreated: 0,
+    cachedPromptTokensRead: 0
+  };
   try {
     for await (const event of response) {
       if (event.type === "message_start") {
-        inputUsage = event.message.usage?.input_tokens || 0;
+        const inputUsage = event.message.usage;
+        if (inputUsage) {
+          const cacheCreated = inputUsage.cache_creation_input_tokens || 0;
+          const cacheRead = inputUsage.cache_read_input_tokens || 0;
+          const promptTokens = inputUsage.input_tokens + cacheCreated + cacheRead;
+          logger.info(`inputUsage: ${JSON.stringify(inputUsage, null, 2)}`);
+          tokensUsed.prompt = promptTokens;
+          tokensUsed.cachedPromptTokensCreated = cacheCreated;
+          tokensUsed.cachedPromptTokensRead = cacheRead;
+        }
       }
       if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
         const chunk = event.delta.text;
@@ -2619,7 +2652,8 @@ async function* AnthropicChatCompletionInner(props, ctx) {
       }
       if (event.type === "message_delta") {
         finishReason = event.delta.stop_reason;
-        outputUsage = event.usage?.output_tokens;
+        tokensUsed.completion = event.usage?.output_tokens || 0;
+        tokensUsed.total = tokensUsed.prompt + tokensUsed.completion;
         span.setAttributes({
           finishReason
         });
@@ -2637,11 +2671,6 @@ async function* AnthropicChatCompletionInner(props, ctx) {
     role: "assistant",
     content
   };
-  const tokensUsed = inputUsage !== void 0 && outputUsage !== void 0 ? {
-    prompt: inputUsage,
-    completion: outputUsage,
-    total: inputUsage + outputUsage
-  } : computeUsage([...chatMessages, outputMessage], anthropicTokenizer);
   const cost = costFn?.(props.model, tokensUsed) ?? void 0;
   const responseData = {
     ...logRequestData,
@@ -2699,6 +2728,92 @@ function cleanChatCompletionRequest2(chatCompletionRequest) {
 // src/lib/anthropic/index.ts
 var import_sdk2 = __toESM(require("@anthropic-ai/sdk"));
 
+// src/lib/anthropic/AnthropicPromptCacheDebugger.tsx
+var import_crypto = __toESM(require("crypto"));
+var renderAttrs = (attrs) => {
+  if (!attrs || Object.keys(attrs).length === 0)
+    return "";
+  return Object.entries(attrs).map(([key, value]) => ` ${key}="${value}"`).join("");
+};
+var indentOver = (str, indent) => {
+  const lines = str.split("\n");
+  return lines.map((l) => "  ".repeat(indent) + l).join("\n");
+};
+var renderTag = (tag2, indent, slice = null) => {
+  if (Array.isArray(tag2)) {
+    return tag2.map((t) => renderTag(t, indent, slice)).join("\n");
+  }
+  if (typeof tag2.children === "string") {
+    const sliced = slice == null ? tag2.children : tag2.children.slice(0, slice);
+    return indentOver(sliced, indent);
+  }
+  if (Array.isArray(tag2.children)) {
+    const children = tag2.children.map((c) => {
+      return renderTag(c, indent + 1, slice);
+    }).join("\n");
+    const res = `<${tag2.type}${renderAttrs(tag2.attrs)}>
+${children}
+</${tag2.type}>`;
+    return indentOver(res, indent);
+  }
+};
+var tag = (type, attrs, children) => {
+  let c;
+  if (typeof children === "string") {
+    c = children;
+  } else if (!Array.isArray(children)) {
+    c = [children];
+  } else {
+    c = children;
+  }
+  return {
+    type,
+    attrs,
+    children: c
+  };
+};
+var AnthropicPromptCacheDebugger = async (props, ctx) => {
+  const chatMessages = await buildChatMessages(ctx, props.children, {
+    useBase64Images: true
+  });
+  const breakpointIndices = chatMessages.map((m, i) => {
+    if (m.cachePrompt) {
+      return i;
+    }
+    return void 0;
+  }).filter((a) => a !== void 0);
+  const messagesToTags = (messages) => {
+    return messages.map((m) => {
+      const role = m.role;
+      const co = m.content;
+      let inner = "";
+      if (typeof co === "string") {
+        inner = co;
+      } else if (Array.isArray(co)) {
+        inner = co.map((c) => {
+          if (c.type === "text") {
+            return tag(c.type, {}, c.text);
+          } else if (c.type === "image") {
+            return tag(c.type, {}, c.image.data?.slice(0, 22) || "image data");
+          }
+          throw new Error("Invalid content");
+        });
+      }
+      return tag(role, {}, inner);
+    });
+  };
+  const final = breakpointIndices.map((acc, i) => {
+    const lastIndex = breakpointIndices[i - 1];
+    const thisIndex = breakpointIndices[i];
+    const toHash = messagesToTags(chatMessages.slice(0, thisIndex + 1));
+    const content = chatMessages.slice(lastIndex + 1, thisIndex + 1);
+    const hash = import_crypto.default.createHash("sha256").update(renderTag(toHash, 0)).digest("hex").slice(0, 8);
+    return tag("breakpoint", { index: i, hash }, messagesToTags(content));
+  });
+  const res = final.map((t) => renderTag(t, 0)).join("\n");
+  return res;
+};
+
 // src/lib/google/Google.tsx
 var import_vertexai = require("@google-cloud/vertexai");
 
@@ -2899,8 +3014,13 @@ async function* GoogleChatCompletionInner(props, ctx) {
     throw errorToChatCompletionError2(err, logRequestData);
   }
   let content = "";
-  let outputUsage = 0;
-  let inputUsage = 0;
+  const tokensUsed = {
+    prompt: 0,
+    completion: 0,
+    total: 0,
+    cachedPromptTokensCreated: 0,
+    cachedPromptTokensRead: 0
+  };
   let finishReason = null;
   try {
     for await (const event of response.stream) {
@@ -2913,10 +3033,10 @@ async function* GoogleChatCompletionInner(props, ctx) {
         }
         if (event.usageMetadata) {
           if (event.usageMetadata.promptTokenCount) {
-            inputUsage = event.usageMetadata.promptTokenCount;
+            tokensUsed.prompt = event.usageMetadata.promptTokenCount;
           }
           if (event.usageMetadata.candidatesTokenCount) {
-            outputUsage = event.usageMetadata.candidatesTokenCount;
+            tokensUsed.completion = event.usageMetadata.candidatesTokenCount;
           }
         }
         const chunk = event.candidates[0].content;
@@ -2954,11 +3074,7 @@ async function* GoogleChatCompletionInner(props, ctx) {
     role: "assistant",
     content
   };
-  const tokensUsed = {
-    prompt: inputUsage,
-    completion: outputUsage,
-    total: inputUsage + outputUsage
-  };
+  tokensUsed.total = tokensUsed.prompt + tokensUsed.completion;
   const cost = costFn?.(props.model, tokensUsed) ?? void 0;
   const responseData = {
     ...logRequestData,
@@ -3040,16 +3156,14 @@ var import_vertexai2 = require("@google-cloud/vertexai");
   Trace,
   UserMessage,
   VertexAI,
-  anthropicTokenizer,
   attachedContextSymbol,
-  computeUsage,
   createAIElement,
   createContext,
   createPrompt,
   createRenderContext,
   evaluatePrompt,
+  experimental__AnthropicPromptCacheDebugger,
   isPromptParsed,
-  openaiTokenizer,
   toDebugMessage,
   tracing
 });