From 7172fe912ab087328bb25af92793798ab0ee4034 Mon Sep 17 00:00:00 2001
From: Karthik Kalyanaraman <karthik@scale3labs.com>
Date: Thu, 12 Sep 2024 10:56:05 -0500
Subject: [PATCH 1/2] fix

---
 components/project/traces/trace-sheet.tsx | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/components/project/traces/trace-sheet.tsx b/components/project/traces/trace-sheet.tsx
index fe31fe5d..65118359 100644
--- a/components/project/traces/trace-sheet.tsx
+++ b/components/project/traces/trace-sheet.tsx
@@ -55,14 +55,14 @@ export function TraceSheet({
   const [span, setSpan] = useState<any | null>(null);
   const [attributes, setAttributes] = useState<any | null>(null);
   const [events, setEvents] = useState<any | null>(null);
-  const [selectedData, setSelectedData] = useState<CheckedData[]>([]);
+  const [selectedData, setSelectedData] = useState<CheckedData | null>(null);
 
   useEffect(() => {
     setSelectedTrace(trace.trace_hierarchy);
     setSelectedVendors(trace.vendors);
     if (trace.vendors.includes("langgraph")) setIncludesLanggraph(true);
     if (!open) {
-      setSelectedData([]);
+      setSelectedData(null);
       setSpansView("SPANS");
     }
   }, [trace, open]);
@@ -115,7 +115,7 @@ export function TraceSheet({
         input,
         output,
       };
-      setSelectedData((prev) => [...prev, checkedData]);
+      setSelectedData(checkedData);
     }
   }, [span]);
 
@@ -160,10 +160,11 @@ export function TraceSheet({
                   <div className="flex gap-2 items-center">
                     <AddtoDataset
                       projectId={project_id}
-                      selectedData={selectedData}
+                      selectedData={selectedData ? [selectedData] : []}
                       disabled={
-                        selectedData.length === 0 ||
-                        !selectedData.every((data) => data.input || data.output)
+                        selectedData === null ||
+                        (selectedData.input === "" &&
+                          selectedData.output === "")
                       }
                     />
                     <Button

From ce91bf04c4f982581a801c0a6728650192009db8 Mon Sep 17 00:00:00 2001
From: Karthik Kalyanaraman <105607645+karthikscale3@users.noreply.github.com>
Date: Thu, 12 Sep 2024 13:39:06 -0700
Subject: [PATCH 2/2] support for o1-preview and o1-mini (#275)

---
 lib/constants.ts              | 20 ++++++++++++++++++++
 lib/types/playground_types.ts | 12 ++++++++++++
 lib/utils.ts                  | 15 +++++++++++++--
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/lib/constants.ts b/lib/constants.ts
index 89802fb0..0298f0a4 100644
--- a/lib/constants.ts
+++ b/lib/constants.ts
@@ -34,6 +34,10 @@ export const TIKTOKEN_MODEL_MAPPING: Record<string, TiktokenEncoding> = {
   "gpt-4o-2024-05-13": "o200k_base",
   "gpt-4o-mini": "o200k_base",
   "gpt-4o-mini-2024-07-18": "o200k_base",
+  "o1-preview": "o200k_base",
+  "o1-mini": "o200k_base",
+  "o1-preview-2024-09-12": "o200k_base",
+  "o1-mini-2024-09-12": "o200k_base",
 };
 
 export type LangTraceAttributes = LLMSpanAttributes &
@@ -53,6 +57,14 @@ export interface CostTableEntry {
 
 // cost per 1000 tokens
 export const OPENAI_PRICING: Record<string, CostTableEntry> = {
+  "o1-preview": {
+    input: 0.015,
+    output: 0.06,
+  },
+  "o1-mini": {
+    input: 0.015,
+    output: 0.06,
+  },
   "gpt-4o-mini": {
     input: 0.00015,
     output: 0.0006,
@@ -244,6 +256,14 @@ export const AZURE_PRICING: Record<string, CostTableEntry> = {
     input: 0.03,
     output: 0.06,
   },
+  "o1-preview": {
+    input: 0.015,
+    output: 0.06,
+  },
+  "o1-mini": {
+    input: 0.015,
+    output: 0.06,
+  },
 };
 
 export const PAGE_SIZE = 15;
diff --git a/lib/types/playground_types.ts b/lib/types/playground_types.ts
index 1ec2da3c..f9d1da8d 100644
--- a/lib/types/playground_types.ts
+++ b/lib/types/playground_types.ts
@@ -6,6 +6,10 @@ export interface Conversation {
 }
 
 export enum OpenAIModel {
+  "o1-preview" = "o1-preview",
+  "o1-preview-2024-09-12" = "o1-preview-2024-09-12",
+  "o1-mini" = "o1-mini",
+  "o1-mini-2024-09-12" = "o1-mini-2024-09-12",
   "gpt-4o-mini" = "gpt-4o-mini",
   "gpt-4o-mini-2024-07-18" = "gpt-4o-mini-2024-07-18",
   "gpt-4o" = "gpt-4o",
@@ -34,6 +38,14 @@ export enum AnthropicModel {
 }
 
 export const openAIModels = [
+  {
+    value: "o1-preview",
+    label: "O1 Preview",
+  },
+  {
+    value: "o1-mini",
+    label: "O1 Mini",
+  },
   {
     value: "gpt-4o-mini",
     label: "GPT-4 Omni Mini",
diff --git a/lib/utils.ts b/lib/utils.ts
index e794bf55..a1ee9de5 100644
--- a/lib/utils.ts
+++ b/lib/utils.ts
@@ -528,6 +528,10 @@ export function calculatePriceFromUsage(
         correctModel = "gpt-4o";
       } else if (model.includes("gpt-4")) {
         correctModel = "gpt-4";
+      } else if (model.includes("o1-preview")) {
+        correctModel = "o1-preview";
+      } else if (model.includes("o1-mini")) {
+        correctModel = "o1-mini";
       }
     }
     costTable = OPENAI_PRICING[correctModel];
@@ -565,6 +569,10 @@ export function calculatePriceFromUsage(
         correctModel = "gpt-4o";
       } else if (model.includes("gpt-4")) {
         correctModel = "gpt-4";
+      } else if (model.includes("o1-preview")) {
+        correctModel = "o1-preview";
+      } else if (model.includes("o1-mini")) {
+        correctModel = "o1-mini";
       }
     }
     costTable = AZURE_PRICING[correctModel];
@@ -647,7 +655,7 @@ export function estimateTokens(prompt: string): number {
 export function calculateTokens(content: string, model: string): number {
   try {
     let tiktokenModel: TiktokenEncoding = "cl100k_base";
-    if (model.startsWith("gpt-4o")) {
+    if (model.startsWith("gpt-4o") || model.startsWith("o1-")) {
       tiktokenModel = "o200k_base";
     }
     return estimateTokensUsingTikToken(content, tiktokenModel);
@@ -744,7 +752,10 @@ export function getVendorFromSpan(span: Span): string {
     vendor = "gemini";
   } else if (span.name.includes("vercel") || serviceName.includes("vercel")) {
     vendor = "vercel";
-  } else if (span.name.includes("embedchain") || serviceName.includes("embedchain")) {
+  } else if (
+    span.name.includes("embedchain") ||
+    serviceName.includes("embedchain")
+  ) {
     vendor = "embedchain";
   }
   return vendor;