diff --git a/docs/docs/tutorial-tracing_2.md b/docs/docs/tutorial-tracing_2.md index 108571cf650..da1980f1155 100644 --- a/docs/docs/tutorial-tracing_2.md +++ b/docs/docs/tutorial-tracing_2.md @@ -5,7 +5,6 @@ In the [Track LLM inputs & outputs](/quickstart) tutorial, the basics of trackin In this tutorial you will learn how to: - **Track data** as it flows though your application - **Track metadata** at call time -- **Export data** that was logged to Weave ## Tracking nested function calls diff --git a/weave-js/src/common/components/elements/LegacyWBIcon.tsx b/weave-js/src/common/components/elements/LegacyWBIcon.tsx index b1fce5a4895..fa440a9ba03 100644 --- a/weave-js/src/common/components/elements/LegacyWBIcon.tsx +++ b/weave-js/src/common/components/elements/LegacyWBIcon.tsx @@ -26,6 +26,10 @@ export interface LegacyWBIconProps { style?: any; 'data-test'?: any; + + role?: string; + ariaHidden?: string; + ariaLabel?: string; } const LegacyWBIconComp = React.forwardRef( @@ -42,6 +46,10 @@ const LegacyWBIconComp = React.forwardRef( onMouseLeave, style, 'data-test': dataTest, + role, + title, + ariaHidden, + ariaLabel, }, ref ) => { @@ -59,6 +67,10 @@ const LegacyWBIconComp = React.forwardRef( onMouseLeave, style, 'data-test': dataTest, + role, + title, + 'aria-hidden': ariaHidden, + 'aria-label': ariaLabel, }; if (ref == null) { return ; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse2/CellValue.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse2/CellValue.tsx index e37595dbd0e..55c52832283 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse2/CellValue.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse2/CellValue.tsx @@ -6,6 +6,8 @@ import {parseRef} from '../../../../react'; import {ValueViewNumber} from '../Browse3/pages/CallPage/ValueViewNumber'; import {ValueViewPrimitive} from '../Browse3/pages/CallPage/ValueViewPrimitive'; import {isRef} from '../Browse3/pages/common/util'; +import {isCustomWeaveTypePayload} from '../Browse3/typeViews/customWeaveType.types'; +import {CustomWeaveTypeDispatcher} from '../Browse3/typeViews/CustomWeaveTypeDispatcher'; import {CellValueBoolean} from './CellValueBoolean'; import {CellValueImage} from './CellValueImage'; import {CellValueString} from './CellValueString'; @@ -64,5 +66,8 @@ export const CellValue = ({value, isExpanded = false}: CellValueProps) => { ); } + if (isCustomWeaveTypePayload(value)) { + return ; + } return ; }; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse2/browse2Util.ts b/weave-js/src/components/PagePanelComponents/Home/Browse2/browse2Util.ts index f8aff50ee95..9bfb3e3e1d4 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse2/browse2Util.ts +++ b/weave-js/src/components/PagePanelComponents/Home/Browse2/browse2Util.ts @@ -1,9 +1,31 @@ -export const flattenObject = ( +/** + * Flatten an object, but preserve any object that has a `_type` field. + * This is critical for handling "Weave Types" - payloads that should be + * treated as holistic objects, rather than flattened. + */ +export const flattenObjectPreservingWeaveTypes = (obj: { + [key: string]: any; +}) => { + return flattenObject(obj, '', {}, (key, value) => { + return ( + typeof value !== 'object' || + value == null || + value._type !== 'CustomWeaveType' + ); + }); +}; + +const flattenObject = ( obj: {[key: string]: any}, parentKey: string = '', - result: {[key: string]: any} = {} + result: {[key: string]: any} = {}, + shouldFlatten: (key: string, value: any) => boolean = () => true ) => { - if (typeof obj !== 'object' || obj === null) { + if ( + typeof obj !== 'object' || + obj === null || + !shouldFlatten(parentKey, obj) + ) { return obj; } const keys = Object.keys(obj); @@ -14,31 +36,14 @@ export const flattenObject = ( const newKey = parentKey ? `${parentKey}.${key}` : key; if (Array.isArray(obj[key])) { result[newKey] = obj[key]; - } else if (typeof obj[key] === 'object') { - flattenObject(obj[key], newKey, result); + } else if ( + typeof obj[key] === 'object' && + shouldFlatten(newKey, obj[key]) + ) { + flattenObject(obj[key], newKey, result, shouldFlatten); } else { result[newKey] = obj[key]; } }); return result; }; -export const unflattenObject = (obj: {[key: string]: any}) => { - const result: {[key: string]: any} = {}; - for (const key in obj) { - if (!obj.hasOwnProperty(key)) { - continue; - } - const keys = key.split('.'); - let current = result; - for (let i = 0; i < keys.length; i++) { - const k = keys[i]; - if (i === keys.length - 1) { - current[k] = obj[key]; - } else { - current[k] = current[k] || {}; - } - current = current[k]; - } - } - return result; -}; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/NotFoundPanel.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/NotFoundPanel.tsx new file mode 100644 index 00000000000..54aeb4477aa --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/NotFoundPanel.tsx @@ -0,0 +1,20 @@ +import {ErrorPanel} from '@wandb/weave/components/ErrorPanel'; +import React, {FC, useContext} from 'react'; + +import {Button} from '../../../Button'; +import {useClosePeek, WeaveflowPeekContext} from './context'; + +export const NotFoundPanel: FC<{title: string}> = ({title}) => { + const close = useClosePeek(); + const {isPeeking} = useContext(WeaveflowPeekContext); + return ( +
+
+ {isPeeking &&
+
+ +
+
+ ); +}; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/CellFilterWrapper.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/CellFilterWrapper.tsx index 2547dd19f93..afbb66e8b39 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/CellFilterWrapper.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/CellFilterWrapper.tsx @@ -11,6 +11,7 @@ type CellFilterWrapperProps = { field: string; operation: string | null; value: any; + style?: React.CSSProperties; }; export const CellFilterWrapper = ({ @@ -19,6 +20,7 @@ export const CellFilterWrapper = ({ field, operation, value, + style, }: CellFilterWrapperProps) => { const onClickCapture = onAddFilter ? (e: React.MouseEvent) => { @@ -31,5 +33,9 @@ export const CellFilterWrapper = ({ } : undefined; - return
{children}
; + return ( +
+ {children} +
+ ); }; 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 323e6ad4c71..15e9ca8e64b 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 @@ -8,6 +8,7 @@ import styled from 'styled-components'; import {MOON_800} from '../../../../../../common/css/color.styles'; 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'; @@ -117,7 +118,10 @@ export const CallDetails: FC<{ flex: '0 0 auto', p: 2, }}> - + + + ) : ( - + + + )} {multipleChildCallOpRefs.map(opVersionRef => { @@ -251,13 +258,15 @@ const getDisplayInputsAndOutput = (call: CallSchema) => { const span = call.rawSpan; const inputKeys = span.inputs._keys ?? - Object.keys(span.inputs).filter(k => !k.startsWith('_')); + Object.keys(span.inputs).filter(k => !k.startsWith('_') || k === '_type'); const inputs = _.fromPairs(inputKeys.map(k => [k, span.inputs[k]])); const callOutput = span.output ?? {}; const outputKeys = callOutput._keys ?? - Object.keys(callOutput).filter(k => k === '_result' || !k.startsWith('_')); + Object.keys(callOutput).filter( + k => k === '_result' || !k.startsWith('_') || k === '_type' + ); const output = _.fromPairs(outputKeys.map(k => [k, callOutput[k]])); return {inputs, output}; }; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx index 517bb1b4e24..78af2e03159 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx @@ -1,5 +1,4 @@ import Box from '@mui/material/Box'; -import {ErrorPanel} from '@wandb/weave/components/ErrorPanel'; import {Loading} from '@wandb/weave/components/Loading'; import {useViewTraceEvent} from '@wandb/weave/integrations/analytics/useViewEvents'; import React, {FC, useCallback} from 'react'; @@ -9,12 +8,9 @@ import {makeRefCall} from '../../../../../../util/refs'; import {Button} from '../../../../../Button'; import {Tailwind} from '../../../../../Tailwind'; import {Browse2OpDefCode} from '../../../Browse2/Browse2OpDefCode'; -import { - TRACETREE_PARAM, - useClosePeek, - useWeaveflowCurrentRouteContext, -} from '../../context'; +import {TRACETREE_PARAM, useWeaveflowCurrentRouteContext} from '../../context'; import {FeedbackGrid} from '../../feedback/FeedbackGrid'; +import {NotFoundPanel} from '../../NotFoundPanel'; import {isEvaluateOp} from '../common/heuristics'; import {CenteredAnimatedLoader} from '../common/Loader'; import {SimplePageLayoutWithHeader} from '../common/SimplePageLayout'; @@ -26,7 +22,6 @@ import {CallDetails} from './CallDetails'; import {CallOverview} from './CallOverview'; import {CallSummary} from './CallSummary'; import {CallTraceView, useCallFlattenedTraceTree} from './CallTraceView'; - export const CallPage: FC<{ entity: string; project: string; @@ -34,7 +29,6 @@ export const CallPage: FC<{ path?: string; }> = props => { const {useCall} = useWFHooks(); - const close = useClosePeek(); const call = useCall({ entity: props.entity, @@ -45,16 +39,7 @@ export const CallPage: FC<{ if (call.loading) { return ; } else if (call.result === null) { - return ( -
-
-
-
- -
-
- ); + return ; } return ; }; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallSummary.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallSummary.tsx index 793191fbd80..daa866cb4fe 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallSummary.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallSummary.tsx @@ -32,7 +32,7 @@ export const CallSummary: React.FC<{ ); return ( -
+
parseRef(props.tableRefUri) as WeaveObjectRef, + [props.tableRefUri] + ); + // Determines if the table itself is truncated const isTruncated = useMemo(() => { return (fetchQuery.result ?? []).length > MAX_ROWS; @@ -96,16 +106,19 @@ export const WeaveCHTable: FC<{ ); return ( - + + + ); }; @@ -133,7 +146,7 @@ export const DataTableView: FC<{ if (val == null) { return {}; } else if (typeof val === 'object' && !Array.isArray(val)) { - return flattenObject(val); + return flattenObjectPreservingWeaveTypes(val); } return {'': val}; }); 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 8868f8a1d7d..24e10a3dd09 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 @@ -21,6 +21,7 @@ import {LoadingDots} from '../../../../../LoadingDots'; import {Browse2OpDefCode} from '../../../Browse2/Browse2OpDefCode'; import {parseRefMaybe} from '../../../Browse2/SmallRef'; import {StyledDataGrid} from '../../StyledDataGrid'; +import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types'; import {isRef} from '../common/util'; import { LIST_INDEX_EDGE_NAME, @@ -163,6 +164,37 @@ export const ObjectViewer = ({ } > = []; traverse(resolvedData, context => { + // Ops should be migrated to the generic CustomWeaveType pattern, but for + // now they are custom handled. + const isOpPayload = context.value?.weave_type?.type === 'Op'; + + if (isCustomWeaveTypePayload(context.value) && !isOpPayload) { + /** + * This block adds an "empty" key that is used to render the custom + * weave type. In the event that a custom type has both properties AND + * custom views, then we might need to extend / modify this part. + */ + const refBackingData = context.value?._ref; + let depth = context.depth; + let path = context.path; + if (refBackingData) { + contexts.push({ + ...context, + isExpandableRef: true, + }); + depth += 1; + path = context.path.plus(''); + } + contexts.push({ + depth, + isLeaf: true, + path, + value: context.value, + valueType: context.valueType, + }); + return 'skip'; + } + if (context.depth !== 0) { const contextTail = context.path.tail(); const isNullDescription = @@ -207,7 +239,8 @@ export const ObjectViewer = ({ if (USE_TABLE_FOR_ARRAYS && context.valueType === 'array') { return 'skip'; } - if (context.value?._ref && context.value?.weave_type?.type === 'Op') { + if (context.value?._ref && isOpPayload) { + // This should be moved to the CustomWeaveType pattern. contexts.push({ depth: context.depth + 1, isLeaf: true, @@ -377,11 +410,15 @@ export const ObjectViewer = ({ isRef(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 + isCode || + isCustomWeaveType ) { return 'auto'; } 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 d5176339465..6a72ee6e6d1 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 @@ -14,6 +14,8 @@ import {isWeaveObjectRef, parseRef} from '../../../../../../react'; import {Alert} from '../../../../../Alert'; import {Button} from '../../../../../Button'; import {CodeEditor} from '../../../../../CodeEditor'; +import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types'; +import {CustomWeaveTypeDispatcher} from '../../typeViews/CustomWeaveTypeDispatcher'; import {isRef} from '../common/util'; import {OBJECT_ATTR_EDGE_NAME} from '../wfReactInterface/constants'; import {WeaveCHTable, WeaveCHTableSourceRefContext} from './DataTableView'; @@ -119,7 +121,7 @@ const ObjectViewerSectionNonEmpty = ({ ); } return null; - }, [apiRef, mode, data, expandedIds]); + }, [mode, apiRef, data, expandedIds]); const setTreeExpanded = useCallback( (setIsExpanded: boolean) => { @@ -215,9 +217,20 @@ export const ObjectViewerSection = ({ noHide, isExpanded, }: ObjectViewerSectionProps) => { - const numKeys = Object.keys(data).length; const currentRef = useContext(WeaveCHTableSourceRefContext); + if (isCustomWeaveTypePayload(data)) { + return ( + <> + + {title} + + + + ); + } + + const numKeys = Object.keys(data).length; if (numKeys === 0) { return ( <> diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ValueView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ValueView.tsx index 1536ff6e9a1..581aba7a729 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ValueView.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ValueView.tsx @@ -1,7 +1,9 @@ import React, {useMemo} from 'react'; -import {parseRef} from '../../../../../../react'; +import {isWeaveObjectRef, parseRef} from '../../../../../../react'; import {parseRefMaybe, SmallRef} from '../../../Browse2/SmallRef'; +import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types'; +import {CustomWeaveTypeDispatcher} from '../../typeViews/CustomWeaveTypeDispatcher'; import {isRef} from '../common/util'; import { DataTableView, @@ -77,5 +79,36 @@ export const ValueView = ({data, isExpanded}: ValueViewProps) => { return
{JSON.stringify(data.value)}
; } + if (data.valueType === 'object') { + if (isCustomWeaveTypePayload(data.value)) { + // This is a little ugly, but essentially if the data is coming from an + // expanded ref, then we want to use that ref to get the entity and project. + // Else we just use the current entity and project. + let entityForWeaveType: string | undefined; + let projectForWeaveType: string | undefined; + + if (valueIsExpandedRef(data)) { + const parsedRef = parseRef((data.value as any)._ref); + if (isWeaveObjectRef(parsedRef)) { + entityForWeaveType = parsedRef.entityName; + projectForWeaveType = parsedRef.projectName; + } + } + + // If we have have a custom view for this weave type, use it. + return ( + + ); + } + } + return
{data.value.toString()}
; }; + +const valueIsExpandedRef = (data: ValueData) => { + return data.value != null && (data.value as any)._ref != null; +}; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx index 7d8ba65cb86..c0007df46dd 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx @@ -284,6 +284,10 @@ function buildCallsTableColumns( const {cols: newCols, groupingModel} = buildDynamicColumns( filteredDynamicColumnNames, + row => { + const [rowEntity, rowProject] = row.project_id.split('/'); + return {entity: rowEntity, project: rowProject}; + }, (row, key) => (row as any)[key], key => expandedRefCols.has(key), key => columnsWithRefs.has(key), diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts index 2ab5664005e..59570f4c14e 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts @@ -1,7 +1,7 @@ import _ from 'lodash'; import {useMemo} from 'react'; -import {flattenObject} from '../../../../../Browse2/browse2Util'; +import {flattenObjectPreservingWeaveTypes} from '../../../../../Browse2/browse2Util'; import { buildCompositeMetricsMap, CompositeScoreMetrics, @@ -138,8 +138,10 @@ export const useFilteredAggregateRows = (state: EvaluationComparisonState) => { evaluationCallId: predictAndScoreRes.evaluationCallId, inputDigest: datasetRow.digest, inputRef: predictAndScoreRes.exampleRef, - input: flattenObject({input: datasetRow.val}), - output: flattenObject({output}), + input: flattenObjectPreservingWeaveTypes({ + input: datasetRow.val, + }), + output: flattenObjectPreservingWeaveTypes({output}), scores: Object.fromEntries( [...Object.entries(state.data.scoreMetrics)].map( ([scoreKey, scoreVal]) => { diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx index 92e51f0e57d..f8c85adeae7 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx @@ -4,6 +4,8 @@ import React, {useMemo} from 'react'; import {maybePluralizeWord} from '../../../../../core/util/string'; import {LoadingDots} from '../../../../LoadingDots'; +import {NotFoundPanel} from '../NotFoundPanel'; +import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher'; import {WeaveCHTableSourceRefContext} from './CallPage/DataTableView'; import {ObjectViewerSection} from './CallPage/ObjectViewerSection'; import {WFHighLevelCallFilter} from './CallsPage/callsTableFilter'; @@ -58,7 +60,7 @@ export const ObjectVersionPage: React.FC<{ if (objectVersion.loading) { return ; } else if (objectVersion.result == null) { - return
Object not found
; + return ; } return ( @@ -207,7 +209,7 @@ const ObjectVersionPageInner: React.FC<{ { label: 'Values', content: ( - + ) : ( - + + + )} diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionsPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionsPage.tsx index aa6a9161a25..93af4fcacc5 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionsPage.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionsPage.tsx @@ -239,9 +239,16 @@ const ObjectVersionsTable: React.FC<{ }); }); - const {cols: newCols, groupingModel} = - buildDynamicColumns(dynamicFields, (row, key) => { - const obj: ObjectVersionSchema = (row as any).obj; + const {cols: newCols, groupingModel} = buildDynamicColumns<{ + obj: ObjectVersionSchema; + }>( + dynamicFields, + row => ({ + entity: row.obj.entity, + project: row.obj.project, + }), + (row, key) => { + const obj: ObjectVersionSchema = row.obj; const res = obj.val?.[key]; if (isTableRef(res)) { // This whole block is a hack to make the table ref clickable. This @@ -258,7 +265,8 @@ const ObjectVersionsTable: React.FC<{ return makeRefExpandedPayload(targetRefUri, res); } return res; - }); + } + ); cols.push(...newCols); groups = groupingModel; } diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/OpVersionPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/OpVersionPage.tsx index 030b8980675..8b76964845a 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/OpVersionPage.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/OpVersionPage.tsx @@ -1,6 +1,7 @@ import React, {useMemo} from 'react'; import {LoadingDots} from '../../../../LoadingDots'; +import {NotFoundPanel} from '../NotFoundPanel'; import {OpCodeViewer} from '../OpCodeViewer'; import { CallsLink, @@ -35,7 +36,7 @@ export const OpVersionPage: React.FC<{ if (opVersion.loading) { return ; } else if (opVersion.result == null) { - return
Op version not found
; + return ; } return ; }; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/common/tabularListViews/columnBuilder.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/common/tabularListViews/columnBuilder.tsx index 9ff6ae4a191..1e81bbd9643 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/common/tabularListViews/columnBuilder.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/common/tabularListViews/columnBuilder.tsx @@ -9,13 +9,15 @@ import React from 'react'; import {isWeaveObjectRef, parseRef} from '../../../../../../../react'; import {ErrorBoundary} from '../../../../../../ErrorBoundary'; -import {flattenObject} from '../../../../Browse2/browse2Util'; +import {flattenObjectPreservingWeaveTypes} from '../../../../Browse2/browse2Util'; import {CellValue} from '../../../../Browse2/CellValue'; import {CollapseHeader} from '../../../../Browse2/CollapseGroupHeader'; import {ExpandHeader} from '../../../../Browse2/ExpandHeader'; import {NotApplicable} from '../../../../Browse2/NotApplicable'; import {SmallRef} from '../../../../Browse2/SmallRef'; import {CellFilterWrapper} from '../../../filters/CellFilterWrapper'; +import {isCustomWeaveTypePayload} from '../../../typeViews/customWeaveType.types'; +import {CustomWeaveTypeProjectContext} from '../../../typeViews/CustomWeaveTypeDispatcher'; import { OBJECT_ATTR_EDGE_NAME, WEAVE_PRIVATE_PREFIX, @@ -60,7 +62,15 @@ export function prepareFlattenedDataForTable( ): Array { return data.map(r => { // First, flatten the inner object - let flattened = flattenObject(r ?? {}); + let flattened = flattenObjectPreservingWeaveTypes(r ?? {}); + + // In the rare case that we have custom objects in the root (this only occurs if you directly) + // publish a custom object. Then we want to instead nest it under an empty key! + if (isCustomWeaveTypePayload(flattened)) { + flattened = { + ' ': flattened, + }; + } flattened = replaceTableRefsInFlattenedData(flattened); @@ -182,6 +192,7 @@ const isExpandedRefWithValueAsTableRef = ( export const buildDynamicColumns = ( filteredDynamicColumnNames: string[], + entityProjectFromRow: (row: T) => {entity: string; project: string}, valueForKey: (row: T, key: string) => any, columnIsExpanded?: (col: string) => boolean, columnCanBeExpanded?: (col: string) => boolean, @@ -269,6 +280,7 @@ export const buildDynamicColumns = ( return val; }, renderCell: cellParams => { + const {entity, project} = entityProjectFromRow(cellParams.row); const val = valueForKey(cellParams.row, key); if (val === undefined) { return ( @@ -287,7 +299,12 @@ export const buildDynamicColumns = ( onAddFilter={onAddFilter} field={key} operation={null} - value={val}> + value={val} + style={{ + width: '100%', + height: '100%', + alignContent: 'center', + }}> {/* In the future, we may want to move this isExpandedRefWithValueAsTableRef condition into `CellValue`. However, at the moment, `ExpandedRefWithValueAsTableRef` is a Table-specific data structure and we might not want to leak that into the @@ -295,7 +312,10 @@ export const buildDynamicColumns = ( {isExpandedRefWithValueAsTableRef(val) ? ( ) : ( - + + + )} diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/tsDataModelHooks.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/tsDataModelHooks.ts index 93f239b2d23..b659a35b845 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/tsDataModelHooks.ts +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/tsDataModelHooks.ts @@ -669,6 +669,13 @@ const useOpVersion = ( }; } + if (opVersionRes.obj == null) { + return { + loading: false, + result: null, + }; + } + const returnedResult = convertTraceServerObjectVersionToOpSchema( opVersionRes.obj ); @@ -812,6 +819,13 @@ const useObjectVersion = ( }; } + if (objectVersionRes.obj == null) { + return { + loading: false, + result: null, + }; + } + const returnedResult: ObjectVersionSchema = convertTraceServerObjectVersionToSchema(objectVersionRes.obj); diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx new file mode 100644 index 00000000000..cc8389fecb1 --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/CustomWeaveTypeDispatcher.tsx @@ -0,0 +1,80 @@ +import React from 'react'; + +import {CustomWeaveTypePayload} from './customWeaveType.types'; +import {PILImageImage} from './PIL.Image.Image/PILImageImage'; + +type CustomWeaveTypeDispatcherProps = { + data: CustomWeaveTypePayload; + // Entity and Project can be optionally provided as props, but if they are not + // provided, they must be provided in context. Failure to provide them will + // result in a console warning and a fallback to a default component. + // + // This pattern is used because in many cases we are rendering data from + // hierarchical data structures, and we want to avoid passing entity and project + // down through the tree. + entity?: string; + project?: string; +}; + +const customWeaveTypeRegistry: { + [typeId: string]: { + component: React.FC<{ + entity: string; + project: string; + data: any; // I wish this could be typed more specifically + }>; + }; +} = { + 'PIL.Image.Image': { + component: PILImageImage, + }, +}; + +/** + * This context is used to provide the entity and project to the + * CustomWeaveTypeDispatcher. Importantly, what this does is allows the + * developer to inject an entity/project context around some component tree, and + * then any CustomWeaveTypeDispatchers within that tree will be assumed to be + * within that entity/project context. This is far cleaner than passing + * entity/project down through the tree. We just have to remember in the future + * case when we support multiple entities/projects in the same tree, we will + * need to update this context if you end up traversing into a different + * entity/project. This should already be accounted for in all the current + * use-cases. + */ +export const CustomWeaveTypeProjectContext = React.createContext<{ + entity: string; + project: string; +} | null>(null); + +/** + * This is the primary entry-point for dispatching custom weave types. Currently + * we just have 1, but as we add more, we might want to add a more robust + * "registry" + */ +export const CustomWeaveTypeDispatcher: React.FC< + CustomWeaveTypeDispatcherProps +> = ({data, entity, project}) => { + const projectContext = React.useContext(CustomWeaveTypeProjectContext); + const typeId = data.weave_type.type; + const comp = customWeaveTypeRegistry[typeId]?.component; + const defaultReturn = Custom Weave Type: {data.weave_type.type}; + + if (comp) { + const applicableEntity = entity || projectContext?.entity; + const applicableProject = project || projectContext?.project; + if (applicableEntity == null || applicableProject == null) { + console.warn( + 'CustomWeaveTypeDispatch: entity and project must be provided in context or as props' + ); + return defaultReturn; + } + return React.createElement(comp, { + entity: applicableEntity, + project: applicableProject, + data, + }); + } + + return defaultReturn; +}; 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 new file mode 100644 index 00000000000..fb07477497f --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/PIL.Image.Image/PILImageImage.tsx @@ -0,0 +1,53 @@ +import React from 'react'; + +import {LoadingDots} from '../../../../../LoadingDots'; +import {useWFHooks} from '../../pages/wfReactInterface/context'; +import {CustomWeaveTypePayload} from '../customWeaveType.types'; + +type PILImageImageTypePayload = CustomWeaveTypePayload< + 'PIL.Image.Image', + {'image.png': string} +>; + +export const isPILImageImageType = ( + data: CustomWeaveTypePayload +): data is PILImageImageTypePayload => { + return data.weave_type.type === 'PIL.Image.Image'; +}; + +export const PILImageImage: React.FC<{ + entity: string; + project: string; + data: PILImageImageTypePayload; +}> = props => { + const {useFileContent} = useWFHooks(); + const imageBinary = useFileContent( + props.entity, + props.project, + props.data.files['image.png'] + ); + + if (imageBinary.loading) { + return ; + } else if (imageBinary.result == null) { + return ; + } + + const arrayBuffer = imageBinary.result as any as ArrayBuffer; + const blob = new Blob([arrayBuffer], {type: 'image/png'}); + const url = URL.createObjectURL(blob); + + // TODO: It would be nice to have a more general image render - similar to the + // ValueViewImage that does things like light box, general scaling, + // downloading, etc.. + return ( + Custom + ); +}; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/customWeaveType.types.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/customWeaveType.types.ts new file mode 100644 index 00000000000..4677e58c407 --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/typeViews/customWeaveType.types.ts @@ -0,0 +1,48 @@ +export type CustomWeaveTypePayload< + T extends string = string, + FP extends {[filename: string]: string} = {[filename: string]: string} +> = { + _type: 'CustomWeaveType'; + weave_type: { + type: T; + }; + files: FP; + load_op?: string | CustomWeaveTypePayload<'Op', {'obj.py': string}>; +} & {[extra: string]: any}; + +export const isCustomWeaveTypePayload = ( + data: any +): data is CustomWeaveTypePayload => { + if (typeof data !== 'object' || data == null) { + return false; + } + if (data._type !== 'CustomWeaveType') { + return false; + } + if ( + typeof data.weave_type !== 'object' || + data.weave_type == null || + typeof data.weave_type.type !== 'string' + ) { + return false; + } + if (typeof data.files !== 'object' || data.files == null) { + return false; + } + if (data.weave_type.type === 'Op') { + if (data.load_op != null) { + return false; + } + } else { + if (data.load_op == null) { + return false; + } + if ( + typeof data.load_op !== 'string' && + !isCustomWeaveTypePayload(data.load_op) + ) { + return false; + } + } + return true; +}; diff --git a/weave/frontend/index.html b/weave/frontend/index.html index 74808c202a6..c96a45a9ae8 100644 --- a/weave/frontend/index.html +++ b/weave/frontend/index.html @@ -91,7 +91,7 @@ - + diff --git a/weave/frontend/sha1.txt b/weave/frontend/sha1.txt index 01a425480b3..7a517474529 100644 --- a/weave/frontend/sha1.txt +++ b/weave/frontend/sha1.txt @@ -1 +1 @@ -cd3b2e94bf9dc8702f53efb3844074379cd3b951 +18eebed493dc14f0fbaa3ab62505c6bdfd42ad6f diff --git a/weave/init_message.py b/weave/init_message.py index ffb6da21cf6..34f107139dd 100644 --- a/weave/init_message.py +++ b/weave/init_message.py @@ -44,10 +44,17 @@ def _print_version_check() -> None: if use_message: print(use_message) - orig_module = wandb._wandb_module - wandb._wandb_module = "weave" - weave_messages = wandb.sdk.internal.update.check_available(weave.__version__) - wandb._wandb_module = orig_module + weave_messages = None + if hasattr(weave, "_wandb_module"): + try: + orig_module = wandb._wandb_module # type: ignore + wandb._wandb_module = "weave" # type: ignore + weave_messages = wandb.sdk.internal.update.check_available( + weave.__version__ + ) + wandb._wandb_module = orig_module # type: ignore + except Exception: + weave_messages = None if weave_messages: use_message = ( diff --git a/weave/tests/test_op.py b/weave/tests/test_op.py index a0c2d3240a2..bfa2271e8e5 100644 --- a/weave/tests/test_op.py +++ b/weave/tests/test_op.py @@ -249,3 +249,28 @@ def my_op(self, a: int) -> str: # type: ignore[empty-body] "a": types.Int(), } assert SomeWeaveObj.my_op.concrete_output_type == types.String() + + +def test_op_internal_tracing_enabled(client): + # This test verifies the behavior of `_tracing_enabled` which + # is not a user-facing API and is used internally to toggle + # tracing on and off. + @weave.op + def my_op(): + return "hello" + + my_op() # <-- this call will be traced + + assert len(list(my_op.calls())) == 1 + + my_op._tracing_enabled = False + + my_op() # <-- this call will not be traced + + assert len(list(my_op.calls())) == 1 + + my_op._tracing_enabled = True + + my_op() # <-- this call will be traced + + assert len(list(my_op.calls())) == 2 diff --git a/weave/trace/custom_objs.py b/weave/trace/custom_objs.py index 7f8b7215c32..6a8c0022d3e 100644 --- a/weave/trace/custom_objs.py +++ b/weave/trace/custom_objs.py @@ -156,6 +156,9 @@ def decode_custom_obj( raise ValueError(f"No serializer found for {weave_type}") load_instance_op = serializer.load + # Disables tracing so that calls to loading data itself don't get traced + load_instance_op._tracing_enabled = False # type: ignore + art = MemTraceFilesArtifact( encoded_path_contents, metadata={}, diff --git a/weave/trace/op.py b/weave/trace/op.py index 2247c3df1eb..8c81869f5a7 100644 --- a/weave/trace/op.py +++ b/weave/trace/op.py @@ -116,6 +116,15 @@ class Op(Protocol): __call__: Callable[..., Any] __self__: Any + # `_tracing_enabled` is a runtime-only flag that can be used to disable + # call tracing for an op. It is not persisted as a property of the op, but is + # respected by the current execution context. It is an underscore property + # because it is not intended to be used by users directly, but rather assists + # with internal Weave behavior. If we find a need to expose this to users, we + # should consider a more user-friendly API (perhaps a setter/getter) & whether + # it disables child ops as well. + _tracing_enabled: bool + def _set_on_output_handler(func: Op, on_output: OnOutputHandlerType) -> None: if func._on_output_handler is not None: @@ -328,6 +337,8 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any: return await func(*args, **kwargs) if weave_client_context.get_weave_client() is None: return await func(*args, **kwargs) + if not wrapper._tracing_enabled: # type: ignore + return await func(*args, **kwargs) call = _create_call(wrapper, *args, **kwargs) # type: ignore res, _ = await _execute_call(wrapper, call, *args, **kwargs) # type: ignore return res @@ -339,6 +350,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return func(*args, **kwargs) if weave_client_context.get_weave_client() is None: return func(*args, **kwargs) + if not wrapper._tracing_enabled: # type: ignore + return func(*args, **kwargs) call = _create_call(wrapper, *args, **kwargs) # type: ignore res, _ = _execute_call(wrapper, call, *args, **kwargs) # type: ignore return res @@ -366,6 +379,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: wrapper._set_on_output_handler = partial(_set_on_output_handler, wrapper) # type: ignore wrapper._on_output_handler = None # type: ignore + wrapper._tracing_enabled = True # type: ignore + return cast(Op, wrapper) return create_wrapper(func) diff --git a/weave/trace/serializer.py b/weave/trace/serializer.py index 7d5bf90fc93..334811111d8 100644 --- a/weave/trace/serializer.py +++ b/weave/trace/serializer.py @@ -52,7 +52,7 @@ def id(self) -> str: # "Op" in the database. if ser_id.endswith(".Op"): return "Op" - return self.target_class.__name__ + return ser_id SERIALIZERS = [] diff --git a/weave/trace/vals.py b/weave/trace/vals.py index 5deb9c4bf4e..8933dcc201b 100644 --- a/weave/trace/vals.py +++ b/weave/trace/vals.py @@ -292,7 +292,13 @@ def _remote_iter(self) -> Generator[dict, None, None]: for item in response.rows: new_ref = self.ref.with_item(item.digest) if self.ref else None - yield make_trace_obj(item.val, new_ref, self.server, self.root) + res = from_json( + item.val, + self.table_ref.entity + "/" + self.table_ref.project, + self.server, + ) + res = make_trace_obj(res, new_ref, self.server, self.root) + yield res if len(response.rows) < page_size: break diff --git a/weave/type_serializers/Image/__init__.py b/weave/type_serializers/Image/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/weave/type_serializers/Image/image.py b/weave/type_serializers/Image/image.py new file mode 100644 index 00000000000..c9b9402a4de --- /dev/null +++ b/weave/type_serializers/Image/image.py @@ -0,0 +1,44 @@ +"""Defines the custom Image weave type.""" + +from weave.trace import serializer +from weave.trace.custom_objs import MemTraceFilesArtifact + +dependencies_met = False + +try: + from PIL import Image + + dependencies_met = True +except ImportError: + pass + + +def save(obj: "Image.Image", artifact: MemTraceFilesArtifact, name: str) -> None: + # Note: I am purposely ignoring the `name` here and hard-coding the filename to "image.png". + # There is an extensive internal discussion here: + # https://weightsandbiases.slack.com/archives/C03BSTEBD7F/p1723670081582949 + # + # In summary, there is an outstanding design decision to be made about how to handle the + # `name` parameter. One school of thought is that using the `name` parameter allows multiple + # object to use the same artifact more cleanly. However, another school of thought is that + # the payload should not be dependent on an external name - resulting in different payloads + # for the same logical object. + # + # Using `image.png` is fine for now since we don't have any cases of multiple objects + # using the same artifact. Moreover, since we package the deserialization logic with the + # object payload, we can always change the serialization logic later without breaking + # existing payloads. + with artifact.new_file("image.png", binary=True) as f: + obj.save(f, format="png") # type: ignore + + +def load(artifact: MemTraceFilesArtifact, name: str) -> "Image.Image": + # Note: I am purposely ignoring the `name` here and hard-coding the filename. See comment + # on save. + path = artifact.path("image.png") + return Image.open(path) + + +def register() -> None: + if dependencies_met: + serializer.register_serializer(Image.Image, save, load) diff --git a/weave/type_serializers/Image/image_test.py b/weave/type_serializers/Image/image_test.py new file mode 100644 index 00000000000..cf42c07c4d6 --- /dev/null +++ b/weave/type_serializers/Image/image_test.py @@ -0,0 +1,102 @@ +from PIL import Image + +import weave +from weave.weave_client import WeaveClient, get_ref + +"""When testing types, it is important to test: +Objects: +1. Publishing Directly +2. Publishing as a property +3. Using as a cell in a table + +Calls: +4. Using as inputs, output, and output component (raw) +5. Using as inputs, output, and output component (refs) + +""" + + +def test_image_publish(client: WeaveClient) -> None: + img = Image.new("RGB", (512, 512), "purple") + weave.publish(img) + + ref = get_ref(img) + + assert ref is not None + gotten_img = weave.ref(ref.uri()).get() + assert img.tobytes() == gotten_img.tobytes() + + +class ImageWrapper(weave.Object): + img: Image.Image + + +def test_image_as_property(client: WeaveClient) -> None: + img = Image.new("RGB", (512, 512), "purple") + img_wrapper = ImageWrapper(img=img) + assert img_wrapper.img == img + + weave.publish(img_wrapper) + + ref = get_ref(img_wrapper) + assert ref is not None + + gotten_img_wrapper = weave.ref(ref.uri()).get() + assert gotten_img_wrapper.img.tobytes() == img.tobytes() + + +def test_image_as_dataset_cell(client: WeaveClient) -> None: + img = Image.new("RGB", (512, 512), "purple") + dataset = weave.Dataset(rows=[{"img": img}]) + assert dataset.rows[0]["img"] == img + + weave.publish(dataset) + + ref = get_ref(dataset) + assert ref is not None + + gotten_dataset = weave.ref(ref.uri()).get() + assert gotten_dataset.rows[0]["img"].tobytes() == img.tobytes() + + +@weave.op +def image_as_solo_output(publish_first: bool) -> Image.Image: + img = Image.new("RGB", (512, 512), "purple") + if publish_first: + weave.publish(img) + return img + + +@weave.op +def image_as_input_and_output_part(in_img: Image.Image) -> dict: + return {"out_img": in_img} + + +def test_image_as_call_io(client: WeaveClient) -> None: + non_published_img = image_as_solo_output(publish_first=False) + img_dict = image_as_input_and_output_part(non_published_img) + + exp_bytes = non_published_img.tobytes() + assert img_dict["out_img"].tobytes() == exp_bytes + + image_as_solo_output_call = image_as_solo_output.calls()[0] + image_as_input_and_output_part_call = image_as_input_and_output_part.calls()[0] + + assert image_as_solo_output_call.output.tobytes() == exp_bytes + assert image_as_input_and_output_part_call.inputs["in_img"].tobytes() == exp_bytes + assert image_as_input_and_output_part_call.output["out_img"].tobytes() == exp_bytes + + +def test_image_as_call_io_refs(client: WeaveClient) -> None: + non_published_img = image_as_solo_output(publish_first=True) + img_dict = image_as_input_and_output_part(non_published_img) + + exp_bytes = non_published_img.tobytes() + assert img_dict["out_img"].tobytes() == exp_bytes + + image_as_solo_output_call = image_as_solo_output.calls()[0] + image_as_input_and_output_part_call = image_as_input_and_output_part.calls()[0] + + assert image_as_solo_output_call.output.tobytes() == exp_bytes + assert image_as_input_and_output_part_call.inputs["in_img"].tobytes() == exp_bytes + assert image_as_input_and_output_part_call.output["out_img"].tobytes() == exp_bytes diff --git a/weave/type_serializers/__init__.py b/weave/type_serializers/__init__.py new file mode 100644 index 00000000000..396af8f791e --- /dev/null +++ b/weave/type_serializers/__init__.py @@ -0,0 +1,3 @@ +from .Image import image + +image.register() diff --git a/weave/weave_client.py b/weave/weave_client.py index ec7ab401adc..47344959a85 100644 --- a/weave/weave_client.py +++ b/weave/weave_client.py @@ -23,6 +23,7 @@ from weave.trace.op import op as op_deco from weave.trace.refs import CallRef, ObjectRef, OpRef, Ref, TableRef from weave.trace.serialize import from_json, isinstance_namedtuple, to_json +from weave.trace.serializer import get_serializer_for_obj from weave.trace.vals import WeaveObject, WeaveTable, make_trace_obj from weave.trace_server.ids import generate_id from weave.trace_server.trace_server_interface import ( @@ -288,7 +289,7 @@ def make_client_call( parent_id=server_call.parent_id, id=server_call.id, inputs=from_json(server_call.inputs, server_call.project_id, server), - output=output, + output=from_json(output, server_call.project_id, server), summary=dict(server_call.summary) if server_call.summary is not None else None, display_name=server_call.display_name, attributes=server_call.attributes, @@ -729,10 +730,18 @@ def _project_id(self) -> str: @trace_sentry.global_trace_sentry.watch() def _save_object(self, val: Any, name: str, branch: str = "latest") -> ObjectRef: self._save_nested_objects(val, name=name) + + # typically, this condition would belong inside of the + # `_save_nested_objects` switch. However, we don't want to recursively + # publish all custom objects. Instead we only want to do this at the + # top-most level if requested + if get_serializer_for_obj(val) is not None: + self._save_and_attach_ref(val) + return self._save_object_basic(val, name, branch) def _save_object_basic( - self, val: Any, name: str, branch: str = "latest" + self, val: Any, name: Optional[str] = None, branch: str = "latest" ) -> ObjectRef: # The WeaveTable case is special because object saving happens inside # _save_object_nested and it has a special table_ref -- skip it here. @@ -745,6 +754,14 @@ def _save_object_basic( return val json_val = to_json(val, self._project_id(), self.server) + if name is None: + if json_val.get("_type") == "CustomWeaveType": + custom_name = json_val.get("weave_type", {}).get("type") + name = custom_name + + if name is None: + raise ValueError("Name must be provided for object saving") + response = self.server.obj_create( ObjCreateReq( obj=ObjSchemaForInsert( @@ -810,11 +827,10 @@ def _save_nested_objects(self, obj: Any, name: Optional[str] = None) -> Any: @trace_sentry.global_trace_sentry.watch() def _save_table(self, table: Table) -> TableRef: + rows = to_json(table.rows, self._project_id(), self.server) response = self.server.table_create( TableCreateReq( - table=TableSchemaForInsert( - project_id=self._project_id(), rows=table.rows - ) + table=TableSchemaForInsert(project_id=self._project_id(), rows=rows) ) ) return TableRef( @@ -848,8 +864,16 @@ def _objects(self, filter: Optional[ObjectVersionFilter] = None) -> list[ObjSche def _save_op(self, op: Op, name: Optional[str] = None) -> Ref: if op.ref is not None: return op.ref + if name is None: name = op.name + + return self._save_and_attach_ref(op, name) + + def _save_and_attach_ref(self, op: Any, name: Optional[str] = None) -> Ref: + if (ref := getattr(op, "ref", None)) is not None: + return ref + op_def_ref = self._save_object_basic(op, name) # setattr(op, "ref", op_def_ref) fails here