diff --git a/tests/integrations/instructor/instructor_test.py b/tests/integrations/instructor/instructor_test.py
index 761d961b9af9..b9667ebb0866 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 331a25800305..3022defc5a1f 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-js/src/assets/icons/icon-filled-circle.svg b/weave-js/src/assets/icons/icon-filled-circle.svg
new file mode 100644
index 000000000000..2b57c31558b1
--- /dev/null
+++ b/weave-js/src/assets/icons/icon-filled-circle.svg
@@ -0,0 +1,3 @@
+
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 000000000000..766810c7811b
--- /dev/null
+++ b/weave-js/src/assets/icons/icon-not-visible.svg
@@ -0,0 +1,5 @@
+
\ 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 000000000000..1ae05ea52ae8
--- /dev/null
+++ b/weave-js/src/assets/icons/icon-pin-to-right.svg
@@ -0,0 +1,5 @@
+
\ 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 a87f41dfe56e..948c8a456b42 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';
@@ -141,6 +142,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 +162,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';
@@ -497,6 +500,9 @@ export const IconFacebookSocial = (props: SVGIconProps) => (
export const IconFailed = (props: SVGIconProps) => (
);
+export const IconFilledCircle = (props: SVGIconProps) => (
+
+);
export const IconFilterAlt = (props: SVGIconProps) => (
);
@@ -698,6 +704,9 @@ export const IconNewSectionAbove = (props: SVGIconProps) => (
export const IconNewSectionBelow = (props: SVGIconProps) => (
);
+export const IconNotVisible = (props: SVGIconProps) => (
+
+);
export const IconNumber = (props: SVGIconProps) => (
);
@@ -755,6 +764,9 @@ export const IconPhoto = (props: SVGIconProps) => (
export const IconPin = (props: SVGIconProps) => (
);
+export const IconPinToRight = (props: SVGIconProps) => (
+
+);
export const IconPlay = (props: SVGIconProps) => (
);
@@ -1134,6 +1146,7 @@ const ICON_NAME_TO_ICON: Record = {
'export-share-upload': IconExportShareUpload,
'facebook-social': IconFacebookSocial,
failed: IconFailed,
+ 'filled-circle': IconFilledCircle,
'filter-alt': IconFilterAlt,
'flash-bolt': IconFlashBolt,
'folder-alt': IconFolderAlt,
@@ -1201,6 +1214,7 @@ const ICON_NAME_TO_ICON: Record = {
'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 +1234,7 @@ const ICON_NAME_TO_ICON: Record = {
'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 aa92f74b31a0..81839a710199 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,
@@ -141,6 +142,7 @@ export {
IconMusicAudio,
IconNewSectionAbove,
IconNewSectionBelow,
+ IconNotVisible,
IconNumber,
IconOpenaiLogo,
IconOpenNewTab,
@@ -160,6 +162,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 8002b47eb9c4..55f46c528333 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',
@@ -140,6 +141,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 +161,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/CallPage/CallDetails.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx
index 8a9c011b5de2..26be9f0c5d0e 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<{
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<{
);
})}
- {childCalls.loading && }
- {/* Disabling display of singular children while we decide if we want them here. */}
- {false && singularChildCalls.length > 0 && (
-
- {multipleChildCallOpRefs.length === 0 ? (
- Child calls
- ) : (
- Singular child calls
- )}
- {singularChildCalls.map(c => (
-
-
-
- ))}
-
- )}
);
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 d741db647714..e8f39ad28802 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 | 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 {inner}
;
+ return {inner}
;
};
// 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 d95d9a841de3..d02a2e881ca7 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 (
- <>
+
{title}
);
};
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 e6e62180f64c..9f845ad0cd1e 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: () => (
+
+ ),
+ columnMenuPinRightIcon: IconPinToRight,
}}
/>
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 000000000000..d1ca67c7c4c9
--- /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; 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(
+ methodName: string,
+ req: Req,
+ getFresh: () => Promise
+ ): Promise {
+ 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 {
+ return this.withCache('fileContent', req, () => super.fileContent(req));
+ }
+
+ public override objRead(req: TraceObjReadReq): Promise {
+ return this.withCache('objRead', req, () => super.objRead(req));
+ }
+
+ public override tableQuery(
+ req: TraceTableQueryReq
+ ): Promise {
+ return this.withCache('tableQuery', req, () => super.tableQuery(req));
+ }
+
+ public override tableQueryStats(
+ req: TraceTableQueryStatsReq
+ ): Promise {
+ 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 2c38bcef5519..b3a15a8a5953 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 ec246ceb8cca..7edce4908daa 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 fb07477497ff..03944dca959d 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',
}}
/>
);
diff --git a/weave/integrations/openai/openai_sdk.py b/weave/integrations/openai/openai_sdk.py
index e68661875f39..7814700d4d3a 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(