From a06b6ace52988bff45012d78a29c9adbbeb32249 Mon Sep 17 00:00:00 2001
From: Tim Sweeney <tim@wandb.com>
Date: Fri, 6 Dec 2024 08:57:21 -0800
Subject: [PATCH 1/4] chore(weave): Improve large object viewer renders (#3102)

* init

* init

* Larger refactor

* Larger refactor

* Larger refactor

* Larger refactor

* Larger refactor
---
 .../Browse3/pages/CallPage/CallDetails.tsx    |  41 ++-----
 .../Browse3/pages/CallPage/ObjectViewer.tsx   |  81 ++++++++-----
 .../pages/CallPage/ObjectViewerSection.tsx    |   4 +-
 .../traceServerCachingClient.ts               | 106 ++++++++++++++++++
 .../wfReactInterface/traceServerClient.ts     |   4 +-
 .../typeViews/CustomWeaveTypeDispatcher.tsx   |   8 ++
 .../PIL.Image.Image/PILImageImage.tsx         |   5 +-
 7 files changed, 183 insertions(+), 66 deletions(-)
 create mode 100644 weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerCachingClient.ts

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx
index 8a9c011b5de..26be9f0c5d0 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx
@@ -1,4 +1,3 @@
-import {Typography} from '@mui/material';
 import Box from '@mui/material/Box';
 import _ from 'lodash';
 import React, {FC, useContext, useMemo} from 'react';
@@ -10,9 +9,7 @@ import {Button} from '../../../../../Button';
 import {useWeaveflowRouteContext, WeaveflowPeekContext} from '../../context';
 import {CustomWeaveTypeProjectContext} from '../../typeViews/CustomWeaveTypeDispatcher';
 import {CallsTable} from '../CallsPage/CallsTable';
-import {KeyValueTable} from '../common/KeyValueTable';
-import {CallLink, opNiceName} from '../common/Links';
-import {CenteredAnimatedLoader} from '../common/Loader';
+import {CallLink} from '../common/Links';
 import {useWFHooks} from '../wfReactInterface/context';
 import {CallSchema} from '../wfReactInterface/wfDataModelHooksInterface';
 import {ButtonOverlay} from './ButtonOverlay';
@@ -20,6 +17,8 @@ import {ExceptionDetails, getExceptionInfo} from './Exceptions';
 import {ObjectViewerSection} from './ObjectViewerSection';
 import {OpVersionText} from './OpVersionText';
 
+const HEADER_HEIGHT_BUFFER = 60;
+
 const Heading = styled.div`
   color: ${MOON_800};
   font-weight: 600;
@@ -104,7 +103,7 @@ export const CallDetails: FC<{
     columns
   );
 
-  const {singularChildCalls, multipleChildCallOpRefs} = useMemo(
+  const {multipleChildCallOpRefs} = useMemo(
     () => callGrouping(!childCalls.loading ? childCalls.result ?? [] : []),
     [childCalls.loading, childCalls.result]
   );
@@ -133,6 +132,7 @@ export const CallDetails: FC<{
         <Box
           sx={{
             flex: '0 0 auto',
+            maxHeight: `calc(100% - ${HEADER_HEIGHT_BUFFER}px)`,
             p: 2,
           }}>
           <CustomWeaveTypeProjectContext.Provider
@@ -143,6 +143,9 @@ export const CallDetails: FC<{
         <Box
           sx={{
             flex: '0 0 auto',
+            maxHeight: `calc(100% - ${
+              multipleChildCallOpRefs.length > 0 ? HEADER_HEIGHT_BUFFER : 0
+            }px)`,
             p: 2,
           }}>
           {'traceback' in excInfo ? (
@@ -201,7 +204,6 @@ export const CallDetails: FC<{
               sx={{
                 flex: '0 0 auto',
                 height: '500px',
-                maxHeight: '95%',
                 p: 2,
                 display: 'flex',
                 flexDirection: 'column',
@@ -240,33 +242,6 @@ export const CallDetails: FC<{
             </Box>
           );
         })}
-        {childCalls.loading && <CenteredAnimatedLoader />}
-        {/* Disabling display of singular children while we decide if we want them here. */}
-        {false && singularChildCalls.length > 0 && (
-          <Box
-            sx={{
-              flex: '0 0 auto',
-            }}>
-            {multipleChildCallOpRefs.length === 0 ? (
-              <Typography pl={1}>Child calls</Typography>
-            ) : (
-              <Typography pl={1}>Singular child calls</Typography>
-            )}
-            {singularChildCalls.map(c => (
-              <Box
-                key={c.callId}
-                sx={{
-                  flex: '0 0 auto',
-                  p: 2,
-                }}>
-                <KeyValueTable
-                  headerTitle={opNiceName(c.spanName)}
-                  data={getDisplayInputsAndOutput(c)}
-                />
-              </Box>
-            ))}
-          </Box>
-        )}
       </Box>
     </Box>
   );
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
index d741db64771..e8f39ad2880 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
@@ -23,7 +23,11 @@ import {LoadingDots} from '../../../../../LoadingDots';
 import {Browse2OpDefCode} from '../../../Browse2/Browse2OpDefCode';
 import {isWeaveRef} from '../../filters/common';
 import {StyledDataGrid} from '../../StyledDataGrid';
-import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
+import {
+  CustomWeaveTypePayload,
+  isCustomWeaveTypePayload,
+} from '../../typeViews/customWeaveType.types';
+import {getCustomWeaveTypePreferredRowHeight} from '../../typeViews/CustomWeaveTypeDispatcher';
 import {
   LIST_INDEX_EDGE_NAME,
   OBJECT_ATTR_EDGE_NAME,
@@ -48,6 +52,10 @@ import {
 } from './traverse';
 import {ValueView} from './ValueView';
 
+const DEFAULT_ROW_HEIGHT = 38;
+const CODE_ROW_HEIGHT = 350;
+const TABLE_ROW_HEIGHT = 450;
+
 type Data = Record<string, any> | any[];
 
 type ObjectViewerProps = {
@@ -440,6 +448,47 @@ export const ObjectViewer = ({
     });
   }, [apiRef, expandedIds, updateRowExpand]);
 
+  // Per https://mui.com/x/react-data-grid/row-height/#dynamic-row-height, always
+  // memoize the getRowHeight function.
+  const getRowHeight = useCallback((params: GridRowHeightParams) => {
+    const isNonRefString =
+      params.model.valueType === 'string' && !isWeaveRef(params.model.value);
+    const isArray = params.model.valueType === 'array';
+    const isTableRef =
+      isWeaveRef(params.model.value) &&
+      (parseRefMaybe(params.model.value) as any).weaveKind === 'table';
+    const {isCode} = params.model;
+    const isCustomWeaveType = isCustomWeaveTypePayload(params.model.value);
+    if (!params.model.isLeaf) {
+      // This is a group header, so we want to use the default height
+      return DEFAULT_ROW_HEIGHT;
+    } else if (isNonRefString) {
+      // This is the only special case where we will allow for dynamic height.
+      // Since strings have special renders that take up different amounts of
+      // space, we will allow for dynamic height.
+      return 'auto';
+    } else if (isCustomWeaveType) {
+      const type = (params.model.value as CustomWeaveTypePayload).weave_type
+        .type;
+      const preferredRowHeight = getCustomWeaveTypePreferredRowHeight(type);
+      if (preferredRowHeight) {
+        return preferredRowHeight;
+      }
+      return DEFAULT_ROW_HEIGHT;
+    } else if ((isArray && USE_TABLE_FOR_ARRAYS) || isTableRef) {
+      // Perfectly enough space for 1 page of data rows
+      return TABLE_ROW_HEIGHT;
+    } else if (isCode) {
+      // Probably will get negative feedback here since code that is < 20 lines
+      // will have some whitespace below the code. However, we absolutely need
+      // to have static height for all cells else the MUI data grid will jump around
+      // when cleaning up virtual rows.
+      return CODE_ROW_HEIGHT;
+    } else {
+      return DEFAULT_ROW_HEIGHT;
+    }
+  }, []);
+
   // Finally, we memoize the inner data grid component. This is important to
   // reduce the number of re-renders when the data changes.
   const inner = useMemo(() => {
@@ -473,31 +522,9 @@ export const ObjectViewer = ({
         isGroupExpandedByDefault={node => {
           return expandedIds.includes(node.id);
         }}
-        autoHeight
         columnHeaderHeight={38}
-        getRowHeight={(params: GridRowHeightParams) => {
-          const isNonRefString =
-            params.model.valueType === 'string' &&
-            !isWeaveRef(params.model.value);
-          const isArray = params.model.valueType === 'array';
-          const isTableRef =
-            isWeaveRef(params.model.value) &&
-            (parseRefMaybe(params.model.value) as any).weaveKind === 'table';
-          const {isCode} = params.model;
-          const isCustomWeaveType = isCustomWeaveTypePayload(
-            params.model.value
-          );
-          if (
-            isNonRefString ||
-            (isArray && USE_TABLE_FOR_ARRAYS) ||
-            isTableRef ||
-            isCode ||
-            isCustomWeaveType
-          ) {
-            return 'auto';
-          }
-          return 38;
-        }}
+        rowHeight={DEFAULT_ROW_HEIGHT}
+        getRowHeight={getRowHeight}
         hideFooter
         rowSelection={false}
         groupingColDef={groupingColDef}
@@ -517,10 +544,10 @@ export const ObjectViewer = ({
         }}
       />
     );
-  }, [apiRef, rows, columns, groupingColDef, expandedIds]);
+  }, [apiRef, rows, columns, getRowHeight, groupingColDef, expandedIds]);
 
   // Return the inner data grid wrapped in a div with overflow hidden.
-  return <div style={{overflow: 'hidden'}}>{inner}</div>;
+  return <div style={{overflow: 'hidden', height: '100%'}}>{inner}</div>;
 };
 
 // Helper function to build the base ref for a given path. This function is used
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewerSection.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewerSection.tsx
index d95d9a841de..d02a2e881ca 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewerSection.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewerSection.tsx
@@ -185,7 +185,7 @@ const ObjectViewerSectionNonEmpty = ({
   }, [data, isExpanded]);
 
   return (
-    <>
+    <Box sx={{height: '100%', display: 'flex', flexDirection: 'column'}}>
       <TitleRow>
         <Title>{title}</Title>
         <Button
@@ -224,7 +224,7 @@ const ObjectViewerSectionNonEmpty = ({
         )}
       </TitleRow>
       {body}
-    </>
+    </Box>
   );
 };
 
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerCachingClient.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerCachingClient.ts
new file mode 100644
index 00000000000..d1ca67c7c4c
--- /dev/null
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerCachingClient.ts
@@ -0,0 +1,106 @@
+import LRUCache from 'lru-cache';
+
+import {
+  TraceFileContentReadReq,
+  TraceFileContentReadRes,
+  TraceObjReadReq,
+  TraceObjReadRes,
+  TraceTableQueryReq,
+  TraceTableQueryRes,
+  TraceTableQueryStatsReq,
+  TraceTableQueryStatsRes,
+} from './traceServerClientTypes';
+import {DirectTraceServerClient} from './traceServerDirectClient';
+
+interface CacheConfig {
+  max: number;
+  getCacheKey: (req: any) => string;
+}
+
+export class CachingTraceServerClient extends DirectTraceServerClient {
+  private caches: Map<
+    string,
+    {cache: LRUCache<string, any>; config: CacheConfig}
+  >;
+
+  constructor(baseUrl: string) {
+    super(baseUrl);
+    this.caches = new Map();
+
+    // Configure caches for specific methods
+    this.addCache('fileContent', {
+      max: 1000,
+      getCacheKey: (req: TraceFileContentReadReq) => JSON.stringify(req),
+    });
+
+    this.addCache('objRead', {
+      max: 1000,
+      getCacheKey: (req: TraceObjReadReq) => JSON.stringify(req),
+    });
+
+    this.addCache('tableQuery', {
+      max: 1000,
+      getCacheKey: (req: TraceTableQueryReq) => JSON.stringify(req),
+    });
+
+    this.addCache('tableQueryStats', {
+      max: 1000,
+      getCacheKey: (req: TraceTableQueryStatsReq) => JSON.stringify(req),
+    });
+  }
+
+  protected addCache(methodName: string, config: CacheConfig) {
+    this.caches.set(methodName, {
+      cache: new LRUCache({max: config.max}),
+      config,
+    });
+  }
+
+  protected withCache<Req, Res>(
+    methodName: string,
+    req: Req,
+    getFresh: () => Promise<Res>
+  ): Promise<Res> {
+    const cacheInfo = this.caches.get(methodName);
+    if (!cacheInfo) {
+      return getFresh();
+    }
+
+    const {cache, config} = cacheInfo;
+    const key = config.getCacheKey(req);
+    const cached = cache.get(key) as Res | undefined;
+
+    if (cached) {
+      return Promise.resolve(cached);
+    }
+
+    return getFresh().then(result => {
+      cache.set(key, result);
+      return result;
+    });
+  }
+
+  public override fileContent(
+    req: TraceFileContentReadReq
+  ): Promise<TraceFileContentReadRes> {
+    return this.withCache('fileContent', req, () => super.fileContent(req));
+  }
+
+  public override objRead(req: TraceObjReadReq): Promise<TraceObjReadRes> {
+    return this.withCache('objRead', req, () => super.objRead(req));
+  }
+
+  public override tableQuery(
+    req: TraceTableQueryReq
+  ): Promise<TraceTableQueryRes> {
+    return this.withCache('tableQuery', req, () => super.tableQuery(req));
+  }
+
+  public override tableQueryStats(
+    req: TraceTableQueryStatsReq
+  ): Promise<TraceTableQueryStatsRes> {
+    return this.withCache('tableQueryStats', req, () =>
+      super.tableQueryStats(req)
+    );
+  }
+}
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClient.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClient.ts
index 2c38bcef551..b3a15a8a595 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClient.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClient.ts
@@ -1,5 +1,6 @@
 import _ from 'lodash';
 
+import {CachingTraceServerClient} from './traceServerCachingClient';
 import {
   CompletionsCreateReq,
   CompletionsCreateRes,
@@ -14,12 +15,11 @@ import {
   TraceRefsReadBatchReq,
   TraceRefsReadBatchRes,
 } from './traceServerClientTypes';
-import {DirectTraceServerClient} from './traceServerDirectClient';
 
 const DEFAULT_BATCH_INTERVAL = 150;
 const MAX_REFS_PER_BATCH = 1000;
 
-export class TraceServerClient extends DirectTraceServerClient {
+export class TraceServerClient extends CachingTraceServerClient {
   private readBatchCollectors: Array<{
     req: TraceRefsReadBatchReq;
     resolvePromise: (res: TraceRefsReadBatchRes) => void;
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx
index ec246ceb8cc..7edce4908da 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx
@@ -24,16 +24,24 @@ const customWeaveTypeRegistry: {
       project: string;
       data: any; // I wish this could be typed more specifically
     }>;
+    preferredRowHeight?: number;
   };
 } = {
   'PIL.Image.Image': {
     component: PILImageImage,
+    preferredRowHeight: 350,
   },
   'wave.Wave_read': {
     component: AudioPlayer,
   },
 };
 
+export const getCustomWeaveTypePreferredRowHeight = (
+  typeId: string
+): number | undefined => {
+  return customWeaveTypeRegistry[typeId]?.preferredRowHeight;
+};
+
 /**
  * This context is used to provide the entity and project to the
  * CustomWeaveTypeDispatcher. Importantly, what this does is allows the
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/PIL.Image.Image/PILImageImage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/PIL.Image.Image/PILImageImage.tsx
index fb07477497f..03944dca959 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/PIL.Image.Image/PILImageImage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/PIL.Image.Image/PILImageImage.tsx
@@ -45,8 +45,9 @@ export const PILImageImage: React.FC<{
       src={url}
       alt="Custom"
       style={{
-        maxWidth: '100%',
-        maxHeight: '100%',
+        width: '100%',
+        height: '100%',
+        objectFit: 'contain',
       }}
     />
   );

From 8d010651bb2a06d8c69cc3fa316870441e7e760d Mon Sep 17 00:00:00 2001
From: Ben Sherman <bcanfieldsherman@gmail.com>
Date: Fri, 6 Dec 2024 09:00:43 -0800
Subject: [PATCH 2/4] chore(weave): capture request id for openai completions
 (#3153)

---
 .../instructor/instructor_test.py             |  6 ++--
 tests/integrations/openai/openai_test.py      | 36 +++++++++----------
 weave/integrations/openai/openai_sdk.py       | 13 ++++---
 3 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/tests/integrations/instructor/instructor_test.py b/tests/integrations/instructor/instructor_test.py
index 761d961b9af..b9667ebb086 100644
--- a/tests/integrations/instructor/instructor_test.py
+++ b/tests/integrations/instructor/instructor_test.py
@@ -64,7 +64,7 @@ def test_instructor_openai(
     assert op_name_from_ref(call.op_name) == "openai.chat.completions.create"
     output = call.output
     output_arguments = json.loads(
-        output.choices[0].message.tool_calls[0].function.arguments
+        output["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"]
     )
     assert "person_name" in output_arguments
     assert "age" in output_arguments
@@ -112,7 +112,7 @@ async def extract_person(text: str) -> Person:
     assert op_name_from_ref(call.op_name) == "openai.chat.completions.create"
     output = call.output
     output_arguments = json.loads(
-        output.choices[0].message.tool_calls[0].function.arguments
+        output["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"]
     )
     assert "person_name" in output_arguments
     assert "age" in output_arguments
@@ -166,7 +166,7 @@ def test_instructor_iterable(
     assert call.started_at < call.ended_at
     assert op_name_from_ref(call.op_name) == "openai.chat.completions.create"
     output = call.output
-    output_arguments = json.loads(output.choices[0].message.content)
+    output_arguments = json.loads(output["choices"][0]["message"]["content"])
     assert "tasks" in output_arguments
     assert "person_name" in output_arguments["tasks"][0]
     assert "age" in output_arguments["tasks"][0]
diff --git a/tests/integrations/openai/openai_test.py b/tests/integrations/openai/openai_test.py
index 331a2580030..3022defc5a1 100644
--- a/tests/integrations/openai/openai_test.py
+++ b/tests/integrations/openai/openai_test.py
@@ -38,10 +38,10 @@ def test_openai_quickstart(client: weave.trace.weave_client.WeaveClient) -> None
     assert call.started_at < call.ended_at  # type: ignore
 
     output = call.output
-    assert output.model == "gpt-4o-2024-05-13"
-    assert output.object == "chat.completion"
+    assert output["model"] == "gpt-4o-2024-05-13"
+    assert output["object"] == "chat.completion"
 
-    usage = call.summary["usage"][output.model]  # type: ignore
+    usage = call.summary["usage"][output["model"]]  # type: ignore
     assert usage["requests"] == 1
     assert usage["completion_tokens"] == 28
     assert usage["prompt_tokens"] == 11
@@ -86,10 +86,10 @@ async def test_openai_async_quickstart(
     assert call.started_at < call.ended_at  # type: ignore
 
     output = call.output
-    assert output.model == "gpt-4o-2024-05-13"
-    assert output.object == "chat.completion"
+    assert output["model"] == "gpt-4o-2024-05-13"
+    assert output["object"] == "chat.completion"
 
-    usage = call.summary["usage"][output.model]  # type: ignore
+    usage = call.summary["usage"][output["model"]]  # type: ignore
     assert usage["requests"] == 1
     assert usage["completion_tokens"] == 28
     assert usage["prompt_tokens"] == 11
@@ -315,10 +315,10 @@ def test_openai_function_call(client: weave.trace.weave_client.WeaveClient) -> N
     assert call.started_at < call.ended_at  # type: ignore
 
     output = call.output
-    assert output.model == "gpt-4o-2024-05-13"
-    assert output.object == "chat.completion"
+    assert output["model"] == "gpt-4o-2024-05-13"
+    assert output["object"] == "chat.completion"
 
-    usage = call.summary["usage"][output.model]  # type: ignore
+    usage = call.summary["usage"][output["model"]]  # type: ignore
     assert usage["total_tokens"] == 117
     assert usage["completion_tokens"] == 18
     assert usage["prompt_tokens"] == 99
@@ -401,10 +401,10 @@ async def test_openai_function_call_async(
     assert call.started_at < call.ended_at  # type: ignore
 
     output = call.output
-    assert output.model == "gpt-4o-2024-05-13"
-    assert output.object == "chat.completion"
+    assert output["model"] == "gpt-4o-2024-05-13"
+    assert output["object"] == "chat.completion"
 
-    usage = call.summary["usage"][output.model]  # type: ignore
+    usage = call.summary["usage"][output["model"]]  # type: ignore
     assert usage["total_tokens"] == 117
     assert usage["completion_tokens"] == 18
     assert usage["prompt_tokens"] == 99
@@ -577,10 +577,10 @@ def test_openai_tool_call(client: weave.trace.weave_client.WeaveClient) -> None:
     assert call.started_at < call.ended_at  # type: ignore
 
     output = call.output
-    assert output.model == "gpt-4o-2024-05-13"
-    assert output.object == "chat.completion"
+    assert output["model"] == "gpt-4o-2024-05-13"
+    assert output["object"] == "chat.completion"
 
-    usage = call.summary["usage"][output.model]  # type: ignore
+    usage = call.summary["usage"][output["model"]]  # type: ignore
     assert usage["total_tokens"] == 117
     assert usage["completion_tokens"] == 27
     assert usage["prompt_tokens"] == 90
@@ -664,10 +664,10 @@ async def test_openai_tool_call_async(
     assert call.started_at < call.ended_at  # type: ignore
 
     output = call.output
-    assert output.model == "gpt-4o-2024-05-13"
-    assert output.object == "chat.completion"
+    assert output["model"] == "gpt-4o-2024-05-13"
+    assert output["object"] == "chat.completion"
 
-    usage = call.summary["usage"][output.model]  # type: ignore
+    usage = call.summary["usage"][output["model"]]  # type: ignore
     assert usage["total_tokens"] == 117
     assert usage["completion_tokens"] == 27
     assert usage["prompt_tokens"] == 90
diff --git a/weave/integrations/openai/openai_sdk.py b/weave/integrations/openai/openai_sdk.py
index e68661875f3..7814700d4d3 100644
--- a/weave/integrations/openai/openai_sdk.py
+++ b/weave/integrations/openai/openai_sdk.py
@@ -94,8 +94,9 @@ def _get_tool_calls(
                 )
         return _tool_calls
 
+    dump = None
     if isinstance(value, ChatCompletionChunk):
-        final_value = ChatCompletion(
+        dump = ChatCompletion(
             id=value.id,
             choices=[
                 {
@@ -116,10 +117,14 @@ def _get_tool_calls(
             object="chat.completion",
             system_fingerprint=value.system_fingerprint,
             usage=value.usage if hasattr(value, "usage") else None,
-        )
-        return final_value.model_dump(exclude_unset=True, exclude_none=True)
-    else:
+        ).model_dump(exclude_unset=True, exclude_none=True)
+    elif not hasattr(value, "model_dump"):
         return value
+    else:
+        dump = value.model_dump(exclude_unset=True, exclude_none=True)
+    if hasattr(value, "_request_id"):
+        dump["request_id"] = value._request_id
+    return dump
 
 
 def openai_accumulator(

From dd43774237789b8ecda6fcb263853a23224e1c45 Mon Sep 17 00:00:00 2001
From: Ben Sherman <bcanfieldsherman@gmail.com>
Date: Fri, 6 Dec 2024 10:11:04 -0800
Subject: [PATCH 3/4] chore(ui): update calls table column menu icons (#3162)

---
 weave-js/src/assets/icons/icon-not-visible.svg    |  5 +++++
 weave-js/src/assets/icons/icon-pin-to-right.svg   |  5 +++++
 weave-js/src/components/Icon/Icon.tsx             | 10 ++++++++++
 weave-js/src/components/Icon/index.ts             |  2 ++
 weave-js/src/components/Icon/types.ts             |  2 ++
 .../Home/Browse3/pages/CallsPage/CallsTable.tsx   | 15 ++++++++++++++-
 6 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 weave-js/src/assets/icons/icon-not-visible.svg
 create mode 100644 weave-js/src/assets/icons/icon-pin-to-right.svg

diff --git a/weave-js/src/assets/icons/icon-not-visible.svg b/weave-js/src/assets/icons/icon-not-visible.svg
new file mode 100644
index 00000000000..766810c7811
--- /dev/null
+++ b/weave-js/src/assets/icons/icon-not-visible.svg
@@ -0,0 +1,5 @@
+<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+    <path
+        d="M21.3458 10.0012C21.7133 10.1923 21.8563 10.6452 21.6652 11.0127L20.9998 10.6667C21.6652 11.0127 21.6652 11.0127 21.6652 11.0127L21.6644 11.0142L21.6634 11.0162L21.6605 11.0216L21.6519 11.0377C21.6448 11.0508 21.6351 11.0685 21.6227 11.0906C21.5979 11.1347 21.5622 11.1963 21.5155 11.2727C21.422 11.4256 21.2838 11.6385 21.0984 11.8914C20.8469 12.2344 20.5068 12.653 20.0715 13.0959C19.9136 13.2565 19.9094 13.5154 20.0687 13.6747L21.2303 14.8363C21.5232 15.1292 21.5232 15.6041 21.2303 15.897C20.9374 16.1899 20.4625 16.1899 20.1696 15.897L18.8569 14.5843C18.7183 14.4457 18.5002 14.4286 18.3376 14.538C17.7094 14.961 16.9854 15.3543 16.1606 15.6706C15.9444 15.7535 15.8335 15.9974 15.9226 16.2112L16.4922 17.5782C16.6515 17.9605 16.4707 18.3997 16.0883 18.559C15.706 18.7183 15.2669 18.5375 15.1076 18.1551L14.4192 16.5031C14.3463 16.328 14.1606 16.229 13.9735 16.2596C13.3566 16.3605 12.6991 16.4167 11.9998 16.4167C11.3006 16.4167 10.6432 16.3605 10.0263 16.2596C9.8392 16.229 9.65352 16.3281 9.58058 16.5031L8.89224 18.1551C8.73293 18.5375 8.29382 18.7183 7.91147 18.559C7.52912 18.3997 7.34831 17.9605 7.50763 17.5782L8.07719 16.2112C8.16627 15.9975 8.0554 15.7536 7.83915 15.6706C7.0144 15.3543 6.29033 14.961 5.6621 14.5381C5.49952 14.4286 5.28141 14.4458 5.14283 14.5843L3.83018 15.897C3.53729 16.1899 3.06242 16.1899 2.76952 15.897C2.47663 15.6041 2.47663 15.1292 2.76952 14.8363L3.93106 13.6748C4.09033 13.5155 4.08616 13.2566 3.92825 13.096C3.49289 12.6531 3.15279 12.2344 2.90126 11.8914C2.71579 11.6385 2.57761 11.4256 2.48417 11.2727C2.43742 11.1963 2.40178 11.1347 2.37695 11.0906C2.36452 11.0685 2.3548 11.0508 2.34772 11.0377L2.33909 11.0216L2.33626 11.0162L2.33521 11.0142C2.33521 11.0142 2.33441 11.0127 2.99977 10.6667L2.33441 11.0127C2.14331 10.6452 2.28631 10.1923 2.65381 10.0012C3.02074 9.81044 3.47274 9.95271 3.66435 10.319L3.66722 10.3243C3.67038 10.3301 3.67603 10.3405 3.68422 10.355C3.70059 10.3841 3.72706 10.43 3.76391 10.4903C3.83765 10.6109 3.9526 10.7886 4.11087 11.0044C4.42796 11.4368 4.91547 12.0175 5.58952 12.5986C6.93141 13.7554 9.01343 14.9167 11.9998 14.9167C14.9862 14.9167 17.0682 13.7554 18.4101 12.5986C19.0842 12.0175 19.5717 11.4368 19.8888 11.0044C20.047 10.7886 20.162 10.6109 20.2357 10.4903C20.2726 10.43 20.299 10.3841 20.3154 10.355C20.3236 10.3405 20.3293 10.3301 20.3324 10.3243L20.3344 10.3206L20.335 10.3196C20.5266 9.95329 20.9789 9.81044 21.3458 10.0012Z"
+        fill="currentColor" />
+</svg>
\ No newline at end of file
diff --git a/weave-js/src/assets/icons/icon-pin-to-right.svg b/weave-js/src/assets/icons/icon-pin-to-right.svg
new file mode 100644
index 00000000000..1ae05ea52ae
--- /dev/null
+++ b/weave-js/src/assets/icons/icon-pin-to-right.svg
@@ -0,0 +1,5 @@
+<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+    <path fill-rule="evenodd" clip-rule="evenodd"
+        d="M3.01719 9.2399C2.06222 10.1952 2.52073 11.7288 3.65016 12.2008C5.16656 12.8346 6.60702 13.9748 7.45814 14.976C8.31598 15.9852 9.31862 17.6021 9.90729 19.2358C10.3377 20.4302 11.8963 20.8641 12.8429 19.9172L15.8021 16.9571L20.2201 21.3765C20.5129 21.6694 20.9876 21.6694 21.2804 21.3765C21.5732 21.0836 21.5732 20.6087 21.2804 20.3158L16.8624 15.8964L19.9117 12.8461C20.8582 11.8993 20.4245 10.3402 19.2305 9.90964C17.5973 9.32079 15.9809 8.31782 14.9721 7.45971C13.9712 6.60832 12.8313 5.16741 12.1977 3.65052C11.7259 2.52073 10.1927 2.06208 9.23776 3.01736L3.01719 9.2399ZM4.22824 10.8168C4.11611 10.7699 4.03833 10.6718 4.01204 10.5595C3.98693 10.4523 4.01377 10.3643 4.07752 10.3006L10.2981 4.07802C10.3618 4.01425 10.4498 3.9874 10.557 4.01252C10.6692 4.03882 10.7673 4.11663 10.8141 4.22879C11.5256 5.93216 12.4296 7.18928 14.0007 8.60243C15.5826 10.0253 16.9977 10.6991 18.722 11.3208C18.8244 11.3577 18.8927 11.4398 18.9167 11.5384C18.9399 11.634 18.9168 11.72 18.8514 11.7855L11.7825 18.8566C11.7171 18.922 11.6311 18.9451 11.5355 18.9219C11.437 18.8979 11.3549 18.8296 11.318 18.7272C10.6965 17.0023 10.0229 15.5868 8.60049 14.0043C7.18779 12.4327 5.93107 11.5284 4.22824 10.8168Z"
+        fill="currentColor" />
+</svg>
\ No newline at end of file
diff --git a/weave-js/src/components/Icon/Icon.tsx b/weave-js/src/components/Icon/Icon.tsx
index a87f41dfe56..9c6b8f1e054 100644
--- a/weave-js/src/components/Icon/Icon.tsx
+++ b/weave-js/src/components/Icon/Icon.tsx
@@ -141,6 +141,7 @@ import {ReactComponent as ImportMolecule} from '../../assets/icons/icon-molecule
 import {ReactComponent as ImportMusicAudio} from '../../assets/icons/icon-music-audio.svg';
 import {ReactComponent as ImportNewSectionAbove} from '../../assets/icons/icon-new-section-above.svg';
 import {ReactComponent as ImportNewSectionBelow} from '../../assets/icons/icon-new-section-below.svg';
+import {ReactComponent as ImportNotVisible} from '../../assets/icons/icon-not-visible.svg';
 import {ReactComponent as ImportNumber} from '../../assets/icons/icon-number.svg';
 import {ReactComponent as ImportOpenNewTab} from '../../assets/icons/icon-open-new-tab.svg';
 import {ReactComponent as ImportOpenaiLogo} from '../../assets/icons/icon-openai-logo.svg';
@@ -160,6 +161,7 @@ import {ReactComponent as ImportPaused} from '../../assets/icons/icon-paused.svg
 import {ReactComponent as ImportPencilEdit} from '../../assets/icons/icon-pencil-edit.svg';
 import {ReactComponent as ImportPhoto} from '../../assets/icons/icon-photo.svg';
 import {ReactComponent as ImportPin} from '../../assets/icons/icon-pin.svg';
+import {ReactComponent as ImportPinToRight} from '../../assets/icons/icon-pin-to-right.svg';
 import {ReactComponent as ImportPlay} from '../../assets/icons/icon-play.svg';
 import {ReactComponent as ImportPlotly} from '../../assets/icons/icon-plotly.svg';
 import {ReactComponent as ImportPriorityCritical} from '../../assets/icons/icon-priority-critical.svg';
@@ -698,6 +700,9 @@ export const IconNewSectionAbove = (props: SVGIconProps) => (
 export const IconNewSectionBelow = (props: SVGIconProps) => (
   <ImportNewSectionBelow {...updateIconProps(props)} />
 );
+export const IconNotVisible = (props: SVGIconProps) => (
+  <ImportNotVisible {...updateIconProps(props)} />
+);
 export const IconNumber = (props: SVGIconProps) => (
   <ImportNumber {...updateIconProps(props)} />
 );
@@ -755,6 +760,9 @@ export const IconPhoto = (props: SVGIconProps) => (
 export const IconPin = (props: SVGIconProps) => (
   <ImportPin {...updateIconProps(props)} />
 );
+export const IconPinToRight = (props: SVGIconProps) => (
+  <ImportPinToRight {...updateIconProps(props)} />
+);
 export const IconPlay = (props: SVGIconProps) => (
   <ImportPlay {...updateIconProps(props)} />
 );
@@ -1201,6 +1209,7 @@ const ICON_NAME_TO_ICON: Record<IconName, ElementType> = {
   'music-audio': IconMusicAudio,
   'new-section-above': IconNewSectionAbove,
   'new-section-below': IconNewSectionBelow,
+  'not-visible': IconNotVisible,
   number: IconNumber,
   'open-new-tab': IconOpenNewTab,
   'openai-logo': IconOpenaiLogo,
@@ -1220,6 +1229,7 @@ const ICON_NAME_TO_ICON: Record<IconName, ElementType> = {
   'pencil-edit': IconPencilEdit,
   photo: IconPhoto,
   pin: IconPin,
+  'pin-to-right': IconPinToRight,
   play: IconPlay,
   plotly: IconPlotly,
   'priority-critical': IconPriorityCritical,
diff --git a/weave-js/src/components/Icon/index.ts b/weave-js/src/components/Icon/index.ts
index aa92f74b31a..d061df41656 100644
--- a/weave-js/src/components/Icon/index.ts
+++ b/weave-js/src/components/Icon/index.ts
@@ -141,6 +141,7 @@ export {
   IconMusicAudio,
   IconNewSectionAbove,
   IconNewSectionBelow,
+  IconNotVisible,
   IconNumber,
   IconOpenaiLogo,
   IconOpenNewTab,
@@ -160,6 +161,7 @@ export {
   IconPencilEdit,
   IconPhoto,
   IconPin,
+  IconPinToRight,
   IconPlay,
   IconPlotly,
   IconPriorityCritical,
diff --git a/weave-js/src/components/Icon/types.ts b/weave-js/src/components/Icon/types.ts
index 8002b47eb9c..e9ea98d1c0e 100644
--- a/weave-js/src/components/Icon/types.ts
+++ b/weave-js/src/components/Icon/types.ts
@@ -140,6 +140,7 @@ export const IconNames = {
   MusicAudio: 'music-audio',
   NewSectionAbove: 'new-section-above',
   NewSectionBelow: 'new-section-below',
+  NotVisible: 'not-visible',
   Number: 'number',
   OpenNewTab: 'open-new-tab',
   OpenaiLogo: 'openai-logo',
@@ -159,6 +160,7 @@ export const IconNames = {
   PencilEdit: 'pencil-edit',
   Photo: 'photo',
   Pin: 'pin',
+  PinToRight: 'pin-to-right',
   Play: 'play',
   Plotly: 'plotly',
   PriorityCritical: 'priority-critical',
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx
index e6e62180f64..9f845ad0cd1 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx
@@ -29,7 +29,13 @@ import {
 import {MOON_200, TEAL_300} from '@wandb/weave/common/css/color.styles';
 import {Switch} from '@wandb/weave/components';
 import {Checkbox} from '@wandb/weave/components/Checkbox/Checkbox';
-import {Icon} from '@wandb/weave/components/Icon';
+import {
+  Icon,
+  IconNotVisible,
+  IconPinToRight,
+  IconSortAscending,
+  IconSortDescending,
+} from '@wandb/weave/components/Icon';
 import React, {
   FC,
   useCallback,
@@ -1069,6 +1075,13 @@ export const CallsTable: FC<{
           },
           columnMenu: CallsCustomColumnMenu,
           pagination: PaginationButtons,
+          columnMenuSortDescendingIcon: IconSortDescending,
+          columnMenuSortAscendingIcon: IconSortAscending,
+          columnMenuHideIcon: IconNotVisible,
+          columnMenuPinLeftIcon: () => (
+            <IconPinToRight style={{transform: 'scaleX(-1)'}} />
+          ),
+          columnMenuPinRightIcon: IconPinToRight,
         }}
       />
     </FilterLayoutTemplate>

From 563fe3a0cebc10ef949dc88bb40e6461e233584d Mon Sep 17 00:00:00 2001
From: Jamie Rasmussen <112953339+jamie-rasmussen@users.noreply.github.com>
Date: Fri, 6 Dec 2024 13:14:40 -0600
Subject: [PATCH 4/4] chore(ui) add filled circle icon (#3163)

---
 weave-js/src/assets/icons/icon-filled-circle.svg | 3 +++
 weave-js/src/components/Icon/Icon.tsx            | 5 +++++
 weave-js/src/components/Icon/index.ts            | 1 +
 weave-js/src/components/Icon/types.ts            | 1 +
 4 files changed, 10 insertions(+)
 create mode 100644 weave-js/src/assets/icons/icon-filled-circle.svg

diff --git a/weave-js/src/assets/icons/icon-filled-circle.svg b/weave-js/src/assets/icons/icon-filled-circle.svg
new file mode 100644
index 00000000000..2b57c31558b
--- /dev/null
+++ b/weave-js/src/assets/icons/icon-filled-circle.svg
@@ -0,0 +1,3 @@
+<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M19 12C19 15.866 15.866 19 12 19C8.13401 19 5 15.866 5 12C5 8.13401 8.13401 5 12 5C15.866 5 19 8.13401 19 12Z" fill="currentColor"/>
+</svg>
diff --git a/weave-js/src/components/Icon/Icon.tsx b/weave-js/src/components/Icon/Icon.tsx
index 9c6b8f1e054..948c8a456b4 100644
--- a/weave-js/src/components/Icon/Icon.tsx
+++ b/weave-js/src/components/Icon/Icon.tsx
@@ -74,6 +74,7 @@ import {ReactComponent as ImportExpandUncollapse} from '../../assets/icons/icon-
 import {ReactComponent as ImportExportShareUpload} from '../../assets/icons/icon-export-share-upload.svg';
 import {ReactComponent as ImportFacebookSocial} from '../../assets/icons/icon-facebook-social.svg';
 import {ReactComponent as ImportFailed} from '../../assets/icons/icon-failed.svg';
+import {ReactComponent as ImportFilledCircle} from '../../assets/icons/icon-filled-circle.svg';
 import {ReactComponent as ImportFilterAlt} from '../../assets/icons/icon-filter-alt.svg';
 import {ReactComponent as ImportFlashBolt} from '../../assets/icons/icon-flash-bolt.svg';
 import {ReactComponent as ImportFolderAlt} from '../../assets/icons/icon-folder-alt.svg';
@@ -499,6 +500,9 @@ export const IconFacebookSocial = (props: SVGIconProps) => (
 export const IconFailed = (props: SVGIconProps) => (
   <ImportFailed {...updateIconProps(props)} />
 );
+export const IconFilledCircle = (props: SVGIconProps) => (
+  <ImportFilledCircle {...updateIconProps(props)} />
+);
 export const IconFilterAlt = (props: SVGIconProps) => (
   <ImportFilterAlt {...updateIconProps(props)} />
 );
@@ -1142,6 +1146,7 @@ const ICON_NAME_TO_ICON: Record<IconName, ElementType> = {
   'export-share-upload': IconExportShareUpload,
   'facebook-social': IconFacebookSocial,
   failed: IconFailed,
+  'filled-circle': IconFilledCircle,
   'filter-alt': IconFilterAlt,
   'flash-bolt': IconFlashBolt,
   'folder-alt': IconFolderAlt,
diff --git a/weave-js/src/components/Icon/index.ts b/weave-js/src/components/Icon/index.ts
index d061df41656..81839a71019 100644
--- a/weave-js/src/components/Icon/index.ts
+++ b/weave-js/src/components/Icon/index.ts
@@ -74,6 +74,7 @@ export {
   IconExportShareUpload,
   IconFacebookSocial,
   IconFailed,
+  IconFilledCircle,
   IconFilterAlt,
   IconFlashBolt,
   IconFolderAlt,
diff --git a/weave-js/src/components/Icon/types.ts b/weave-js/src/components/Icon/types.ts
index e9ea98d1c0e..55f46c52833 100644
--- a/weave-js/src/components/Icon/types.ts
+++ b/weave-js/src/components/Icon/types.ts
@@ -73,6 +73,7 @@ export const IconNames = {
   ExportShareUpload: 'export-share-upload',
   FacebookSocial: 'facebook-social',
   Failed: 'failed',
+  FilledCircle: 'filled-circle',
   FilterAlt: 'filter-alt',
   FlashBolt: 'flash-bolt',
   FolderAlt: 'folder-alt',