@gammatech/aijsx 0.1.3 → 0.2.0-beta.1

package/dist/index.js

@@ -47,14 +47,14 @@ __export(src_exports, {
   SystemMessage: () => SystemMessage,
   UserMessage: () => UserMessage,
   attachedContextSymbol: () => attachedContextSymbol,
-  childrenToConversationMessage: () => childrenToConversationMessage,
   computeUsage: () => computeUsage,
   countAnthropicTokens: () => import_tokenizer3.countTokens,
   createAIElement: () => createAIElement,
   createContext: () => createContext,
   createRenderContext: () => createRenderContext,
   defaultMaxTokens: () => defaultMaxTokens,
-  tokenCountForConversationMessage: () => tokenCountForConversationMessage,
+  tokenCountForOpenAIMessage: () => tokenCountForOpenAIMessage,
+  tokenCountForOpenAIVisionMessage: () => tokenCountForOpenAIVisionMessage,
   tokenLimitForChatModel: () => tokenLimitForChatModel,
   tokenizer: () => tokenizer
 });
@@ -70,32 +70,9 @@ var UserMessage = (props) => {
 var AssistantMessage = (props) => {
   return props.children;
 };
-var childrenToConversationMessage = (c) => {
-  const children = Array.isArray(c) ? c : [c];
-  return children.map((child) => {
-    if (child.tag.name === "UserMessage") {
-      return {
-        type: "user",
-        element: child
-      };
-    } else if (child.tag.name === "SystemMessage") {
-      return {
-        type: "system",
-        element: child
-      };
-    } else if (child.tag.name === "AssistantMessage") {
-      return {
-        type: "assistant",
-        element: child
-      };
-    } else {
-      throw new Error("OpenAI: unknown message type");
-    }
-  });
-};
 var computeUsage = (messages) => {
-  const prompt = messages.filter((m) => m.type === "user" || m.type === "system").reduce((acc, m) => acc + m.tokens, 0);
-  const completion = messages.filter((m) => m.type === "assistant").reduce((acc, m) => acc + m.tokens, 0);
+  const prompt = messages.filter((m) => m.role === "user" || m.role === "system").reduce((acc, m) => acc + m.tokens, 0);
+  const completion = messages.filter((m) => m.role === "assistant").reduce((acc, m) => acc + m.tokens, 0);
   return {
     prompt,
     completion,
@@ -125,7 +102,7 @@ function createAIElement(tag, props, ...children) {
   return result;
 }
 function isAIElement(value) {
-  return value !== null && typeof value === "object" && "tag" in value;
+  return value !== null && typeof value === "object" && "tag" in value && "render" in value;
 }
 function isLiteral(value) {
   return typeof value === "string" || typeof value === "number" || typeof value === "undefined" || typeof value === "boolean" || // capture null + undefined
@@ -314,6 +291,108 @@ function createRenderContext({
 // src/types.ts
 var attachedContextSymbol = Symbol("AI.attachedContext");
 
+// src/xml.ts
+var import_fast_xml_parser = require("fast-xml-parser");
+var XmlNode = class {
+  constructor(parent, nodeName, attributes, value, childNodes) {
+    this.parent = parent;
+    this.nodeName = nodeName;
+    this.attributes = attributes;
+    this.value = value;
+    this.childNodes = childNodes;
+    this.childNodes = childNodes;
+    if (childNodes) {
+      childNodes.forEach((n) => n.parent = this);
+    }
+  }
+  toObject() {
+    if (this.value) {
+      return {
+        nodeName: this.nodeName,
+        value: this.value
+      };
+    }
+    if (this.childNodes) {
+      return {
+        nodeName: this.nodeName,
+        attributes: this.attributes,
+        childNodes: this.childNodes.map((n) => n.toObject())
+      };
+    }
+  }
+  get textContent() {
+    return this.value || this.childNodes?.map((n) => n.textContent).join("");
+  }
+  collapse(nodes) {
+    if (this.value !== void 0) {
+      return this;
+    }
+    if (this.childNodes) {
+      this.childNodes.forEach((n) => n.collapse(nodes));
+    }
+    if (
+      // is the root node
+      !this.parent || // is a preserved node tag
+      nodes.includes(this.nodeName)
+    ) {
+      return this;
+    }
+    const parent = this.parent;
+    const index = parent.childNodes.indexOf(this);
+    parent.childNodes.splice(index, 1, ...this.childNodes);
+    return this;
+  }
+};
+var XmlFragment = class extends XmlNode {
+  // initialize with value here because ts isn't smart enough
+  // to understand that the 4th argument to super initializes
+  // childNodes
+  childNodes = [];
+  constructor(childNodes) {
+    super(null, "#fragment", {}, void 0, childNodes);
+  }
+};
+function parseXml(input) {
+  const parser = new import_fast_xml_parser.XMLParser({
+    isArray: (_tag, path) => {
+      if (path.indexOf(".") === -1) {
+        return false;
+      }
+      return true;
+    },
+    ignoreAttributes: false,
+    attributeNamePrefix: "",
+    preserveOrder: true,
+    trimValues: false
+  });
+  const constructNode = (parent, nodeObject) => {
+    if (nodeObject.hasOwnProperty("#text")) {
+      return new XmlNode(parent, "#text", {}, nodeObject["#text"]);
+    }
+    const nodeName = Object.keys(nodeObject)[0];
+    const attributeName = Object.keys(nodeObject)[1];
+    const childObjects = nodeObject[nodeName];
+    const attributes = Object.entries(nodeObject[attributeName] || {}).reduce((acc, [key, value]) => {
+      try {
+        acc[key] = JSON.parse(value);
+      } catch (e) {
+        console.error(
+          `Error parsing attribute value (attr=${key} value=${value}) for element ${nodeName}`
+        );
+      }
+      return acc;
+    }, {});
+    const node = parent === null ? new XmlFragment([]) : new XmlNode(parent, nodeName, attributes, void 0, []);
+    node.childNodes = childObjects.map((child) => constructNode(node, child));
+    return node;
+  };
+  const parsed = parser.parse(`<#fragment>${input}</#fragment>`);
+  return constructNode(null, parsed[0]);
+}
+function escape(html) {
+  return html.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;").replace(/"/g, "&quot;").replace(/'/g, "&#039;");
+}
+
 // src/context.ts
 var LoggerContext = createContext(
   new NoopLogImplementation()
@@ -469,8 +548,11 @@ var StreamRenderContext = class _StreamRenderContext {
     this.contextValues = contextValues;
     const logImpl = this.getContext(LoggerContext);
     this.logger = new BoundLogger(logImpl, this);
-    this.render = (renderable) => {
-      const generator = this.renderStream(renderable);
+    this.render = (renderable, opts) => {
+      const generator = this.renderStream(
+        renderable,
+        opts || { preserveTags: false }
+      );
       const result = {
         then: (onFulfilled, onRejected) => accumResults(generator).then(onFulfilled, onRejected),
         [Symbol.asyncIterator]: () => generator
@@ -478,9 +560,14 @@ var StreamRenderContext = class _StreamRenderContext {
       return result;
     };
     const self = this;
-    this.renderStream = async function* (renderable) {
+    this.renderStream = async function* (renderable, opts) {
+      const preserveTags = opts.preserveTags ?? false;
+      const renderedProps = opts.renderedProps || {};
+      const renderEscaped = (val) => {
+        return preserveTags ? escape(val) : val;
+      };
       if (isLiteral(renderable)) {
-        yield renderLiteral(renderable);
+        yield renderEscaped(renderLiteral(renderable));
         return;
       }
       if (isAIElement(renderable)) {
@@ -488,8 +575,29 @@ var StreamRenderContext = class _StreamRenderContext {
         const childRenderId = uuidv4();
         const newCtx = self.enter(renderable, childRenderId, ctxValues);
         const logger = newCtx.logger;
+        const isFragment = renderable.tag.name === "AIFragment";
         try {
-          return yield* newCtx.render(renderable.render(newCtx));
+          if (preserveTags && !isFragment) {
+            const propsToRender = Object.entries(
+              renderable.props
+            ).reduce(
+              // create an array of escaped + json serialized key="{value}"s for attributes
+              (acc, [name, value]) => {
+                if (renderedProps[renderable.tag.name]?.[name]) {
+                  acc.push(`${name}="${escape(JSON.stringify(value))}"`);
+                }
+                return acc;
+              },
+              []
+            ).join(" ");
+            const attrs = propsToRender ? ` ${propsToRender}` : "";
+            yield `<${renderable.tag.name}${attrs}>`;
+          }
+          yield* newCtx.renderStream(renderable.render(newCtx), opts);
+          if (preserveTags && !isFragment) {
+            yield `</${renderable.tag.name}>`;
+          }
+          return;
         } catch (ex) {
           logger.logException(ex);
           throw ex;
@@ -497,10 +605,10 @@ var StreamRenderContext = class _StreamRenderContext {
       }
       if (Array.isArray(renderable)) {
         if (renderable.every((r) => isLiteral(r))) {
-          yield renderable.map((r) => renderLiteral(r)).join("");
+          yield renderable.map((r) => renderEscaped(renderLiteral(r))).join("");
           return;
         }
-        const streams = renderable.filter((a) => !!a).map((r) => self.renderStream(r));
+        const streams = renderable.filter((a) => !!a).map((r) => self.renderStream(r, opts));
         const result = coalesceParallelStreams(streams);
         while (true) {
           const { value, done } = await result.next();
@@ -521,7 +629,7 @@ var StreamRenderContext = class _StreamRenderContext {
       const next = await renderable.then(
         (r) => r
       );
-      return yield* self.render(next);
+      return yield* self.renderStream(next, opts);
     };
   }
   render;
@@ -623,23 +731,44 @@ function tokenLimitForChatModel(model) {
     }
   }
 }
-function tokenCountForConversationMessage(message) {
+function tokenCountForOpenAIMessage(message) {
   const TOKENS_PER_MESSAGE = 3;
-  switch (message.type) {
+  switch (message.role) {
     case "assistant":
     case "system":
     case "user":
-      return TOKENS_PER_MESSAGE + tokenizer.encode(message.content).length;
+      return (
+        // TODO this isn't working for vision
+        TOKENS_PER_MESSAGE + tokenizer.encode(message.content).length
+      );
   }
 }
-
-// src/jsx-runtime.ts
-function jsx(type, config, maybeKey) {
-  const configWithKey = maybeKey !== void 0 ? { ...config, key: maybeKey } : config;
-  const children = config && Array.isArray(config.children) ? config.children : [];
-  return createAIElement(type, configWithKey, ...children);
+function tokenCountForOpenAIVisionMessage(message) {
+  const TOKENS_PER_MESSAGE = 3;
+  const textCost = (content) => {
+    return TOKENS_PER_MESSAGE + tokenizer.encode(content).length;
+  };
+  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_url.detail || part.image_url.detail === "low") {
+            return acc + 85;
+          } else {
+            return acc + (170 * 4 + 85);
+          }
+        }
+      }, 0);
+  }
 }
-var jsxs = jsx;
 
 // src/lib/openai/OpenAI.tsx
 var defaultClient = null;
@@ -651,35 +780,60 @@ var OpenAIClientContext = createContext(() => {
   defaultClient = new import_openai.OpenAI({ apiKey });
   return defaultClient;
 });
+function buildOpenAIMessages(childrenXml) {
+  const messages = [];
+  const chatMessageTags2 = ["UserMessage", "AssistantMessage", "SystemMessage"];
+  const parsed = parseXml(childrenXml).collapse(chatMessageTags2);
+  const topLevelValid = parsed.childNodes.every(
+    (node) => chatMessageTags2.includes(node.nodeName)
+  );
+  if (!topLevelValid) {
+    throw new Error("Invalid top level chat message tags");
+  }
+  for (const node of parsed.childNodes) {
+    if (node.nodeName === "UserMessage") {
+      messages.push({
+        content: node.textContent,
+        role: "user"
+      });
+    } else if (node.nodeName === "AssistantMessage") {
+      messages.push({
+        content: node.textContent,
+        role: "assistant"
+      });
+    } else if (node.nodeName === "SystemMessage") {
+      messages.push({
+        content: node.textContent,
+        role: "system"
+      });
+    }
+  }
+  return messages;
+}
 async function* OpenAIChatCompletion(props, { logger, render, getContext }) {
   const startTime = performance.now();
   const client = getContext(OpenAIClientContext)();
   if (!client) {
     throw new Error("[OpenAI] must supply OpenAI model via context");
   }
-  const renderedMessages = await Promise.all(
-    childrenToConversationMessage(props.children).map(async (message) => {
-      const partiallyRendered = {
-        ...message,
-        content: await render(message.element)
-      };
-      return {
-        ...partiallyRendered,
-        tokens: tokenCountForConversationMessage(partiallyRendered)
-      };
+  const openAIMessages = buildOpenAIMessages(
+    await render(props.children, {
+      preserveTags: true
     })
   );
-  const chatMessages = renderedMessages.map((m) => {
+  const renderedMessages = openAIMessages.map((message) => {
     return {
-      content: m.content,
-      role: m.type
+      role: message.role,
+      // TODO support gpt4 vision
+      content: message.content,
+      tokens: tokenCountForOpenAIMessage(message)
     };
   });
   const chatCompletionRequest = {
     model: props.model,
     max_tokens: props.maxTokens,
     temperature: props.temperature,
-    messages: chatMessages,
+    messages: openAIMessages,
     stream: true
   };
   const logRequestData = {
@@ -721,11 +875,10 @@ async function* OpenAIChatCompletion(props, { logger, render, getContext }) {
     }
   }
   const outputMessage = {
-    type: "assistant",
-    element: /* @__PURE__ */ jsx(AssistantMessage, { children: content }),
+    role: "assistant",
     content,
-    tokens: tokenCountForConversationMessage({
-      type: "assistant",
+    tokens: tokenCountForOpenAIMessage({
+      role: "assistant",
       content
     })
   };
@@ -758,6 +911,48 @@ var AnthropicClientContext = createContext(
   }
 );
 var defaultMaxTokens = 4096;
+var chatMessageTags = ["UserMessage", "AssistantMessage", "SystemMessage"];
+function buildChatMessages(childrenXml) {
+  const messages = [];
+  const parsed = parseXml(childrenXml).collapse(chatMessageTags);
+  const topLevelValid = parsed.childNodes.every(
+    (node) => chatMessageTags.includes(node.nodeName)
+  );
+  if (!topLevelValid) {
+    throw new Error("Invalid top level chat message tags");
+  }
+  for (const node of parsed.childNodes) {
+    if (node.nodeName === "UserMessage") {
+      const content = `${import_sdk.default.HUMAN_PROMPT} ${node.textContent}`;
+      messages.push({
+        role: "user",
+        content,
+        tokens: (0, import_tokenizer2.countTokens)(content)
+      });
+    } else if (node.nodeName === "AssistantMessage") {
+      const content = `${import_sdk.default.AI_PROMPT} ${node.textContent}`;
+      messages.push({
+        role: "assistant",
+        content,
+        tokens: (0, import_tokenizer2.countTokens)(content)
+      });
+    } else if (node.nodeName === "SystemMessage") {
+      const userContent = `${import_sdk.default.HUMAN_PROMPT} For subsequent replies you will adhere to the following instructions: ${node.textContent}`;
+      messages.push({
+        role: "user",
+        content: userContent,
+        tokens: (0, import_tokenizer2.countTokens)(userContent)
+      });
+      const assistantContent = `${import_sdk.default.AI_PROMPT} Okay, I will do that.`;
+      messages.push({
+        role: "assistant",
+        content: assistantContent,
+        tokens: (0, import_tokenizer2.countTokens)(assistantContent)
+      });
+    }
+  }
+  return messages;
+}
 async function* AnthropicChatCompletion(props, { render, logger, getContext }) {
   const startTime = performance.now();
   const client = getContext(AnthropicClientContext)();
@@ -766,41 +961,17 @@ async function* AnthropicChatCompletion(props, { render, logger, getContext }) {
       "[AnthropicChatCompletion] must supply AnthropicClient via context"
     );
   }
-  const renderedMessages = await Promise.all(
-    childrenToConversationMessage(props.children).flatMap((message) => {
-      if (message.type === "system") {
-        return [
-          {
-            type: "user",
-            element: /* @__PURE__ */ jsxs(UserMessage, { children: [
-              "For subsequent replies you will adhere to the following instructions: ",
-              message.element
-            ] })
-          },
-          {
-            type: "assistant",
-            element: /* @__PURE__ */ jsx(AssistantMessage, { children: "Okay, I will do that." })
-          }
-        ];
-      }
-      return [message];
-    }).map(async (message) => {
-      const prefix = message.type === "user" ? import_sdk.default.HUMAN_PROMPT : import_sdk.default.AI_PROMPT;
-      const rendered = await render(message.element);
-      const content2 = `${prefix} ${rendered.trim()}`;
-      return {
-        ...message,
-        content: content2,
-        tokens: (0, import_tokenizer2.countTokens)(content2)
-      };
+  const inputMessages = buildChatMessages(
+    await render(props.children, {
+      preserveTags: true
     })
   );
-  const chatMessages = renderedMessages.map((m) => {
-    return m.content;
-  });
-  chatMessages.push(import_sdk.default.AI_PROMPT);
+  const prompt = [
+    ...inputMessages.map((message) => message.content),
+    import_sdk.default.AI_PROMPT
+  ].join("");
   const anthropicCompletionRequest = {
-    prompt: chatMessages.join("\n\n"),
+    prompt,
     max_tokens_to_sample: props.maxTokens ?? defaultMaxTokens,
     temperature: props.temperature,
     model: props.model,
@@ -811,7 +982,7 @@ async function* AnthropicChatCompletion(props, { render, logger, getContext }) {
     model: props.model,
     provider: props.provider,
     providerRegion: props.providerRegion,
-    inputMessages: renderedMessages,
+    inputMessages,
     request: anthropicCompletionRequest
   };
   logger.chatCompletionRequest("anthropic", logRequestData);
@@ -843,8 +1014,7 @@ async function* AnthropicChatCompletion(props, { render, logger, getContext }) {
     yield text;
   }
   const outputMessage = {
-    type: "assistant",
-    element: /* @__PURE__ */ jsx(AssistantMessage, { children: content }),
+    role: "assistant",
     content,
     tokens: (0, import_tokenizer2.countTokens)(content)
   };
@@ -853,7 +1023,7 @@ async function* AnthropicChatCompletion(props, { render, logger, getContext }) {
     finishReason: "stop",
     latency: performance.now() - startTime,
     outputMessage,
-    tokensUsed: computeUsage([...renderedMessages, outputMessage])
+    tokensUsed: computeUsage([...inputMessages, outputMessage])
   };
   logger.chatCompletionResponse("anthropic", responseData);
 }
@@ -881,14 +1051,14 @@ var import_tokenizer3 = require("@anthropic-ai/tokenizer");
   SystemMessage,
   UserMessage,
   attachedContextSymbol,
-  childrenToConversationMessage,
   computeUsage,
   countAnthropicTokens,
   createAIElement,
   createContext,
   createRenderContext,
   defaultMaxTokens,
-  tokenCountForConversationMessage,
+  tokenCountForOpenAIMessage,
+  tokenCountForOpenAIVisionMessage,
   tokenLimitForChatModel,
   tokenizer
 });