update pricing and differentiate output tokens in price calculation

src/sections/TokenViewer.tsx

@@ -1,4 +1,4 @@
-import { Fragment, useState } from "react";
+import { Fragment, useMemo, useState } from "react";
 import { cn } from "~/utils/cn";
 
 import BN from "bignumber.js";
@@ -29,11 +29,19 @@ const COLORS = [
 const PRICING: Record<string, BN> = {
   "gpt-4": BN("0.03").div(1000),
   "gpt-4-1106-preview": BN("0.01").div(1000),
-  "gpt-4-32k": BN("0.03").div(1000),
-  "gpt-3.5-turbo": BN("0.0010").div(1000),
+  "gpt-4-32k": BN("0.06").div(1000),
+  "gpt-3.5-turbo": BN("0.0005").div(1000),
   "gpt-3.5-instruct": BN("0.0015").div(1000),
 };
 
+const OUT_PRICING: Record<string, BN> = {
+  "gpt-4": BN("0.06").div(1000),
+  "gpt-4-1106-preview": BN("0.03").div(1000),
+  "gpt-4-32k": BN("0.12").div(1000),
+  "gpt-3.5-turbo": BN("0.0015").div(1000),
+  "gpt-3.5-instruct": BN("0.0020").div(1000),
+};
+
 function encodeWhitespace(str: string) {
   let result = str;
 
@@ -62,6 +70,42 @@ export function TokenViewer(props: {
   const tokenCount =
     props.data?.reduce((memo, i) => memo + i.tokens.length, 0) ?? 0;
   const pricing = props.model != null ? PRICING[props.model] : undefined;
+  const outPricing = props.model != null ? OUT_PRICING[props.model] : undefined;
+
+  const accuratePrice = useMemo(() => {
+    const dividedData: Array<
+      { text: string; tokens: { id: number; idx: number }[] } | undefined
+    >[] = [[]];
+    for (
+      let segIdx = 0;
+      segIdx < (props.data !== undefined ? props.data.length : 0);
+      segIdx++
+    ) {
+      if (props.data?.[segIdx]?.text === "<|im_start|>") {
+        dividedData.push([]);
+      }
+      if (props.data?.[segIdx] !== undefined) {
+        dividedData[dividedData.length - 1]?.push(props.data?.[segIdx]);
+      }
+    }
+
+    const divided = dividedData.filter(item => item.length !== 0);
+    const wentWrong = divided.some(
+      (item) =>
+        item?.[0]?.text !== "<|im_start|>"
+    );
+    if (wentWrong) {
+      return;
+    }
+    const finalPrice =
+      divided.reduce((accPrice, curDiv) => {
+        const divTokenCount =
+          curDiv?.flatMap(f => f ? [f] : []).reduce((memo, i) => memo + i.tokens.length, 0) ?? 0;
+        const finalPricing = (curDiv?.[1]?.text === "assistant") ? outPricing : pricing;
+        return accPrice.plus((finalPricing || BN("0")).multipliedBy(divTokenCount));
+      }, BN("0")) ?? 0;
+      return finalPrice;
+  }, [props.data, pricing, outPricing]);
 
   const [showWhitespace, setShowWhitespace] = useState(false);
 
@@ -77,7 +121,7 @@ export function TokenViewer(props: {
           <div className="flex-grow rounded-md border bg-slate-50 p-4 shadow-sm">
             <p className="text-sm ">Price per prompt</p>
             <p className="text-lg">
-              ${pricing?.multipliedBy(tokenCount)?.toFixed()}
+              ${accuratePrice?.toFixed() || pricing?.multipliedBy(tokenCount)?.toFixed()}
             </p>
           </div>
         )}