Merge branch 'master' into improve-chatnvidia-langchain-integration

.github/workflows/test.yaml

@@ -81,11 +81,16 @@ jobs:
         env:
           CI: 1
           WANDB_ENABLE_TEST_CONTAINER: true
+          LOGGING_ENABLED: true
         ports:
           - '8080:8080'
           - '8083:8083'
           - '9015:9015'
-        options: --health-cmd "curl --fail http://localhost:8080/healthz || exit 1" --health-interval=5s --health-timeout=3s
+        options: >-
+          --health-cmd "wget -q -O /dev/null http://localhost:8080/healthz || exit 1"
+          --health-interval=5s
+          --health-timeout=3s
+          --health-start-period=10s
     outputs:
       tests_should_run: ${{ steps.test_check.outputs.tests_should_run }}
     steps:
@@ -254,11 +259,16 @@ jobs:
         env:
           CI: 1
           WANDB_ENABLE_TEST_CONTAINER: true
+          LOGGING_ENABLED: true
         ports:
           - '8080:8080'
           - '8083:8083'
           - '9015:9015'
-        options: --health-cmd "curl --fail http://localhost:8080/healthz || exit 1" --health-interval=5s --health-timeout=3s
+        options: >-
+          --health-cmd "wget -q -O /dev/null http://localhost:8080/healthz || exit 1"
+          --health-interval=5s
+          --health-timeout=3s
+          --health-start-period=10s
       weave_clickhouse:
         image: clickhouse/clickhouse-server
         ports:
@@ -267,6 +277,8 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v3
+      - name: Enable debug logging
+        run: echo "ACTIONS_STEP_DEBUG=true" >> $GITHUB_ENV
       - name: Set up Python ${{ matrix.python-version-major }}.${{ matrix.python-version-minor }}
         uses: actions/setup-python@v5
         with:

sdks/node/package.json

@@ -1,6 +1,6 @@
 {
   "name": "weave",
-  "version": "0.7.3",
+  "version": "0.7.4",
   "description": "AI development toolkit",
   "types": "dist/index.d.ts",
   "main": "dist/index.js",

sdks/node/src/__tests__/weaveClient.test.ts

@@ -108,7 +108,7 @@ describe('WeaveClient', () => {
     beforeEach(() => {
       mockTraceServerApi = {
         call: {
-          callStartBatchCallUpsertBatchPost: jest.fn(),
+          callStartBatchCallUpsertBatchPost: jest.fn().mockResolvedValue({}),
         },
       } as any;
       mockWandbServerApi = {} as any;
@@ -121,6 +121,26 @@ describe('WeaveClient', () => {
       (client as any).BATCH_INTERVAL = 10;
     });
 
+    it('should handle oversized batch items', async () => {
+      const bigPayloadSize = 11 * 1024 * 1024;
+      const smallData = {mode: 'start', data: {id: '2', payload: 'small'}};
+      const bigData = {
+        mode: 'start',
+        data: {id: '1', payload: 'x'.repeat(bigPayloadSize)},
+      };
+      (client as any).callQueue.push(smallData, bigData);
+
+      await (client as any).processBatch();
+
+      expect(
+        mockTraceServerApi.call.callStartBatchCallUpsertBatchPost
+      ).toHaveBeenCalledWith({
+        batch: [{mode: 'start', req: smallData.data}],
+      });
+
+      expect((client as any).callQueue).toContain(bigData);
+    });
+
     it('should batch multiple calls together', async () => {
       // Add test calls to queue
       (client as any).callQueue.push(

sdks/node/src/weaveClient.ts

@@ -1,4 +1,5 @@
 import {AsyncLocalStorage} from 'async_hooks';
+import * as fs from 'fs';
 import {uuidv7} from 'uuidv7';
 
 import {Dataset} from './dataset';
@@ -32,6 +33,8 @@ import {packageVersion} from './utils/userAgent';
 import {WandbServerApi} from './wandb/wandbServerApi';
 import {ObjectRef, WeaveObject, getClassChain} from './weaveObject';
 
+const WEAVE_ERRORS_LOG_FNAME = 'weaveErrors.log';
+
 export type CallStackEntry = {
   callId: string;
   traceId: string;
@@ -71,13 +74,18 @@ class CallStack {
 type CallStartParams = StartedCallSchemaForInsert;
 type CallEndParams = EndedCallSchemaForInsert;
 
+// We count characters item by item, and try to limit batches to about this size.
+const MAX_BATCH_SIZE_CHARS = 10 * 1024 * 1024;
+
 export class WeaveClient {
   private stackContext = new AsyncLocalStorage<CallStack>();
   private callQueue: Array<{mode: 'start' | 'end'; data: any}> = [];
   private batchProcessTimeout: NodeJS.Timeout | null = null;
   private isBatchProcessing: boolean = false;
   private batchProcessingPromises: Set<Promise<void>> = new Set();
   private readonly BATCH_INTERVAL: number = 200;
+  private errorCount = 0;
+  private readonly MAX_ERRORS = 10;
 
   constructor(
     public traceServerApi: TraceServerApi<any>,
@@ -114,25 +122,43 @@ export class WeaveClient {
 
     this.isBatchProcessing = true;
 
-    // We count characters item by item, and try to limit batches to about
-    // this size.
-    const maxBatchSizeChars = 5 * 1024 * 1024;
-
     let batchToProcess = [];
     let currentBatchSize = 0;
 
-    while (this.callQueue.length > 0 && currentBatchSize < maxBatchSizeChars) {
-      const item = this.callQueue[0];
+    while (
+      this.callQueue.length > 0 &&
+      currentBatchSize < MAX_BATCH_SIZE_CHARS
+    ) {
+      const item = this.callQueue.shift();
+      if (item === undefined) {
+        throw new Error('Call queue is empty');
+      }
+
       const itemSize = JSON.stringify(item).length;
+      if (itemSize > MAX_BATCH_SIZE_CHARS) {
+        fs.appendFileSync(
+          WEAVE_ERRORS_LOG_FNAME,
+          `Item size ${itemSize} exceeds max batch size ${MAX_BATCH_SIZE_CHARS}.  Item: ${JSON.stringify(item)}\n`
+        );
+      }
 
-      if (currentBatchSize + itemSize <= maxBatchSizeChars) {
-        batchToProcess.push(this.callQueue.shift()!);
+      if (currentBatchSize + itemSize <= MAX_BATCH_SIZE_CHARS) {
+        batchToProcess.push(item);
         currentBatchSize += itemSize;
       } else {
+        // doesn't fit, put it back
+        this.callQueue.unshift(item);
         break;
       }
     }
 
+    if (batchToProcess.length === 0) {
+      this.batchProcessTimeout = null;
+      return;
+    }
+
+    this.isBatchProcessing = true;
+
     const batchReq = {
       batch: batchToProcess.map(item => ({
         mode: item.mode,
@@ -146,8 +172,20 @@ export class WeaveClient {
       );
     } catch (error) {
       console.error('Error processing batch:', error);
+      this.errorCount++;
+      fs.appendFileSync(
+        WEAVE_ERRORS_LOG_FNAME,
+        `Error processing batch: ${error}\n`
+      );
+
       // Put failed items back at the front of the queue
       this.callQueue.unshift(...batchToProcess);
+
+      // Exit if we have too many errors
+      if (this.errorCount > this.MAX_ERRORS) {
+        console.error(`Exceeded max errors: ${this.MAX_ERRORS}; exiting`);
+        process.exit(1);
+      }
     } finally {
       this.isBatchProcessing = false;
       this.batchProcessTimeout = null;
@@ -734,7 +772,9 @@ function mergeSummaries(left: Summary, right: Summary): Summary {
       if (typeof leftValue === 'number' && typeof result[key] === 'number') {
         result[key] = leftValue + result[key];
       } else if (
+        leftValue != null &&
         typeof leftValue === 'object' &&
+        result[key] != null &&
         typeof result[key] === 'object'
       ) {
         result[key] = mergeSummaries(leftValue, result[key]);

wb_schema.gql

@@ -104,10 +104,24 @@ type UpdateUserPayload {
   clientMutationId: String
 }
 
+input InsertSecretInput {
+  entityName: String!
+  secretName: String!
+    @constraints(max: 255, pattern: "^[A-Za-z_][A-Za-z0-9_]*$")
+  secretValue: String!
+  clientMutationId: String
+}
+
+type InsertSecretPayload {
+  success: Boolean!
+  clientMutationId: String
+}
+
 type Mutation {
   updateUser(input: UpdateUserInput!): UpdateUserPayload @audit
   deleteView(input: DeleteViewInput!): DeleteViewPayload
   upsertView(input: UpsertViewInput!): UpsertViewPayload @audit
+  insertSecret(input: InsertSecretInput!): InsertSecretPayload
   updateArtifactSequence(
     input: UpdateArtifactSequenceInput!
   ): UpdateArtifactCollectionPayload
@@ -275,6 +289,12 @@ type RowType {
   row: JSON!
 }
 
+type Secret {
+  entityId: Int!
+  name: String!
+  createdAt: DateTime!
+}
+
 type Entity implements Node {
   id: ID!
   name: String!
@@ -296,6 +316,7 @@ type Entity implements Node {
     filters: JSONString
     collectionTypes: [ArtifactCollectionType!]
   ): ArtifactCollectionConnection
+  secrets: [Secret!]!
 }
 
 type EntityConnection {

weave-js/src/common/components/elements/ModifiedDropdown.test.tsx

@@ -52,6 +52,18 @@ describe('testing simple search', () => {
       value: '_step',
       key: '_step',
     },
+    {
+      icon: 'wbic-ic-up-arrow',
+      text: 'Stepperoni',
+      value: '_stepperoni',
+      key: '_stepperoni',
+    },
+    {
+      icon: 'wbic-ic-up-arrow',
+      text: '99',
+      value: 99,
+      key: '_99',
+    },
     {
       icon: 'calendar',
       text: 'Relative Time (Wall)',
@@ -137,6 +149,40 @@ describe('testing simple search', () => {
     expect(results.every(r => (r.value as string).includes('loss'))).toBe(true);
   });
 
+  it('simpleSearch matches case-insensitive regex strings', () => {
+    const results = simpleSearch(options, 'LOSS', {
+      allowRegexSearch: true,
+    });
+    expect(results.every(r => (r.value as string).includes('loss'))).toBe(true);
+  });
+
+  it('simpleSearch matches case-insensitive regex strings', () => {
+    const results = simpleSearch(options, 'tep$', {
+      allowRegexSearch: true,
+    });
+    expect(results.length).toBe(1);
+    expect(results.every(r => (r.value as string).includes('_step'))).toBe(
+      true
+    );
+  });
+
+  it('simpleSearch matches options with number values', () => {
+    const results = simpleSearch(options, '99$', {
+      allowRegexSearch: true,
+    });
+    expect(results.length).toBe(1);
+    results.forEach(r => {
+      expect(r.value).toEqual(99);
+    });
+  });
+
+  it('simpleSearch matches all results on * regex string', () => {
+    const results = simpleSearch(options, '*', {
+      allowRegexSearch: true,
+    });
+    expect(results.length).toBe(options.length);
+  });
+
   it('simpleSearch can disallow matching regex patterns', () => {
     const results = simpleSearch(options, '.*s.*s.*');
     expect(results.length).toBe(0);

weave-js/src/common/components/elements/ModifiedDropdown.tsx

@@ -45,9 +45,23 @@ type LabelCoord = {
 
 const ITEM_LIMIT_VALUE = '__item_limit';
 
+/**
+ * The functionality here is similar to `searchRegexFromQuery` in `panelbank.ts`
+ */
 export function getAsValidRegex(s: string): RegExp | null {
+  let cleanS = s.trim();
+
+  // if the query is a single '*', match everything (even though * isn't technically a valid regex)
+  if (cleanS === '*') {
+    cleanS = '.*';
+  }
+
+  if (cleanS.length === 0) {
+    return null;
+  }
+
   try {
-    return new RegExp(s);
+    return new RegExp(cleanS, 'i');
   } catch (e) {
     return null;
   }
@@ -64,14 +78,18 @@ export const simpleSearch = (
 
   return _.chain(options)
     .filter(o => {
-      const text = JSON.stringify(o.text).toLowerCase();
-      return regex ? regex.test(text) : _.includes(text, query.toLowerCase());
+      const t = typeof o.text === 'string' ? o.text : JSON.stringify(o.text);
+
+      return regex
+        ? regex.test(t)
+        : _.includes(t.toLowerCase(), query.toLowerCase());
     })
     .sortBy(o => {
-      const valJSON = typeof o.text === 'string' ? `"${query}"` : query;
-      return JSON.stringify(o.text).toLowerCase() === valJSON.toLowerCase()
-        ? 0
-        : 1;
+      const oString =
+        typeof o.text === 'string' ? o.text : JSON.stringify(o.text);
+      const qString = typeof query === 'string' ? query : JSON.stringify(query);
+
+      return oString.toLowerCase() === qString.toLowerCase() ? 0 : 1;
     })
     .value();
 };
@@ -148,10 +166,10 @@ const ModifiedDropdown: FC<ModifiedDropdownProps> = React.memo(
               _.concat(currentOptions, search(propsOptions, query) as Option[])
             );
           } else {
-            const updatedOptions = currentOptions.concat(
-              simpleSearch(propsOptions, query, {allowRegexSearch}) as Option[]
-            );
-            setOptions(updatedOptions);
+            const matchedOptions = simpleSearch(propsOptions, query, {
+              allowRegexSearch,
+            }) as Option[];
+            setOptions([...currentOptions, ...matchedOptions]);
           }
         }, debounceTime || 400),
       [allowRegexSearch, debounceTime, multiple, propsOptions, search, value]