chore(weave): add obj create and table update bindings to wf react in…

…terface
wandb · Jan 7, 2025 · 063e8c1 · 063e8c1
1 parent fb3e3a7
commit 063e8c1
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 3 deletions.
diff --git a/...c/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClient.ts b/...c/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClient.ts
@@ -8,6 +8,8 @@ import {
   FeedbackCreateRes,
   FeedbackPurgeReq,
   FeedbackPurgeRes,
+  TableUpdateReq,
+  TableUpdateRes,
   TraceCallsDeleteReq,
   TraceCallUpdateReq,
   TraceObjCreateReq,
@@ -115,6 +117,10 @@ export class TraceServerClient extends CachingTraceServerClient {
     return res;
   }
 
+  public tableUpdate(req: TableUpdateReq): Promise<TableUpdateRes> {
+    return super.tableUpdate(req);
+  }
+
   public feedbackCreate(req: FeedbackCreateReq): Promise<FeedbackCreateRes> {
     const res = super.feedbackCreate(req).then(createRes => {
       const listeners = this.onFeedbackListeners[req.weave_ref] ?? [];

diff --git a/...ponents/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientTypes.ts b/...ponents/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientTypes.ts
@@ -347,3 +347,35 @@ export type ActionsExecuteBatchReq = {
 };
 
 export type ActionsExecuteBatchRes = {};
+
+export type TableUpdateSpec = TableAppendSpec | TablePopSpec | TableInsertSpec;
+
+export interface TableAppendSpec {
+  append: {
+    row: Record<string, any>;
+  };
+}
+
+export interface TablePopSpec {
+  pop: {
+    index: number;
+  };
+}
+
+export interface TableInsertSpec {
+  insert: {
+    index: number;
+    row: Record<string, any>;
+  };
+}
+
+export type TableUpdateReq = {
+  project_id: string;
+  base_digest: string;
+  updates: TableUpdateSpec[];
+};
+
+export type TableUpdateRes = {
+  digest: string;
+  updated_row_digests: string[];
+};
diff --git a/...onents/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerDirectClient.ts b/...onents/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerDirectClient.ts
@@ -27,6 +27,8 @@ import {
   FeedbackPurgeRes,
   FeedbackQueryReq,
   FeedbackQueryRes,
+  TableUpdateReq,
+  TableUpdateRes,
   TraceCallReadReq,
   TraceCallReadRes,
   TraceCallSchema,
@@ -253,6 +255,13 @@ export class DirectTraceServerClient {
     );
   }
 
+  public tableUpdate(req: TableUpdateReq): Promise<TableUpdateRes> {
+    return this.makeRequest<TableUpdateReq, TableUpdateRes>(
+      '/table/update',
+      req
+    );
+  }
+
   public tableQuery(req: TraceTableQueryReq): Promise<TraceTableQueryRes> {
     return this.makeRequest<TraceTableQueryReq, TraceTableQueryRes>(
       '/table/query',

diff --git a/...rc/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/tsDataModelHooks.ts b/...rc/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/tsDataModelHooks.ts
@@ -1713,21 +1713,75 @@ export const traceCallToUICallSchema = (
   };
 };
 
+export const useObjCreate = (): ((
+  projectId: string,
+  objectId: string,
+  val: any,
+  baseObjectClass?: string
+) => Promise<string>) => {
+  const getTsClient = useGetTraceServerClientContext();
+
+  return useCallback(
+    (
+      projectId: string,
+      objectId: string,
+      val: any,
+      baseObjectClass?: string
+    ) => {
+      return getTsClient()
+        .objCreate({
+          obj: {
+            project_id: projectId,
+            object_id: objectId,
+            val,
+            builtin_object_class: baseObjectClass,
+          },
+        })
+        .then(res => {
+          return res.digest;
+        });
+    },
+    [getTsClient]
+  );
+};
+
+export const useTableUpdate = (): ((
+  projectId: string,
+  digest: string,
+  updates: traceServerTypes.TableUpdateSpec[]
+) => Promise<traceServerTypes.TableUpdateRes>) => {
+  const getTsClient = useGetTraceServerClientContext();
+
+  return useCallback(
+    (
+      projectId: string,
+      baseDigest: string,
+      updates: traceServerTypes.TableUpdateSpec[]
+    ) => {
+      return getTsClient().tableUpdate({
+        project_id: projectId,
+        base_digest: baseDigest,
+        updates,
+      });
+    },
+    [getTsClient]
+  );
+};
+
 /// Utility Functions ///
 
 export const convertISOToDate = (iso: string): Date => {
   return new Date(iso);
 };
 
-// Export //
-
 export const tsWFDataModelHooks: WFDataModelHooksInterface = {
   useCall,
   useCalls,
   useCallsStats,
   useCallsDeleteFunc,
   useCallUpdateFunc,
   useCallsExport,
+  useObjCreate,
   useOpVersion,
   useOpVersions,
   useObjectVersion,
@@ -1738,6 +1792,7 @@ export const tsWFDataModelHooks: WFDataModelHooksInterface = {
   useFileContent,
   useTableRowsQuery,
   useTableQueryStats,
+  useTableUpdate,
   derived: {
     useChildCallsForCompare,
     useGetRefsType,

diff --git a/...ents/PagePanelComponents/Home/Browse3/pages/wfReactInterface/wfDataModelHooksInterface.ts b/...ents/PagePanelComponents/Home/Browse3/pages/wfReactInterface/wfDataModelHooksInterface.ts
@@ -14,7 +14,7 @@ import {WeaveKind} from '../../../../../../react';
 import {KNOWN_BASE_OBJECT_CLASSES, OP_CATEGORIES} from './constants';
 import {Query} from './traceServerClientInterface/query'; // TODO: This import is not ideal, should delete this whole interface
 import * as traceServerClientTypes from './traceServerClientTypes'; // TODO: This import is not ideal, should delete this whole interface
-import {ContentType} from './traceServerClientTypes';
+import {ContentType, TableUpdateSpec} from './traceServerClientTypes';
 
 export type OpCategory = (typeof OP_CATEGORIES)[number];
 export type KnownBaseObjectClassType =
@@ -214,6 +214,12 @@ export type WFDataModelHooksInterface = {
     expandedRefCols?: string[],
     includeFeedback?: boolean
   ) => Promise<Blob>;
+  useObjCreate: () => (
+    projectId: string,
+    objectId: string,
+    val: any,
+    baseObjectClass?: string
+  ) => Promise<string>;
   useOpVersion: (key: OpVersionKey | null) => Loadable<OpVersionSchema | null>;
   useOpVersions: (
     entity: string,
@@ -270,6 +276,11 @@ export type WFDataModelHooksInterface = {
     key: FeedbackKey | null,
     sortBy?: traceServerClientTypes.SortBy[]
   ) => LoadableWithError<any[] | null> & Refetchable;
+  useTableUpdate: () => (
+    projectId: string,
+    baseDigest: string,
+    updates: traceServerClientTypes.TableUpdateSpec[]
+  ) => Promise<traceServerClientTypes.TableUpdateRes>;
   derived: {
     useChildCallsForCompare: (
       entity: string,