[WIP] Add experiment logging support for BrainTrust

packages/app/src/components/trivet/TestSuite.tsx

@@ -16,12 +16,15 @@ import {
 } from '@ironclad/trivet';
 import { trivetState } from '../../state/trivet';
 import Button from '@atlaskit/button';
+
 import { TryRunTests } from './api';
 import { useOpenUrl } from '../../hooks/useOpenUrl';
 import { ReactComponent as BrowserLineIcon } from 'majesticons/line/browser-line.svg';
 import { ReactComponent as AlertCircleIcon } from 'majesticons/line/alert-circle-line.svg';
 import { NoTestCasesSplash } from './NoTestCasesSplash';
 import { useTestSuite } from '../../hooks/useTestSuite';
+import { braintrustSummariesState } from '../../state/plugins/braintrust';
+import { useIsPluginEnabled } from '../../hooks/useIsPluginEnabled';
 
 const styles = css`
   min-height: 100%;
@@ -135,6 +138,10 @@ export const TestSuiteRenderer: FC<{ tryRunTests: TryRunTests }> = ({ tryRunTest
 export const TestSuite: FC<{ testSuite: TrivetTestSuite; tryRunTests: TryRunTests }> = ({ testSuite, tryRunTests }) => {
   const [{ selectedTestSuiteId, editingTestCaseId, recentTestResults, runningTests }, setState] =
     useRecoilState(trivetState);
+  const [{ brainTrustSummaries }, setBrainTrustSummariesState] = useRecoilState(braintrustSummariesState);
+
+  const brainTrustEnabled = useIsPluginEnabled('braintrust');
+
   const savedGraphs = useRecoilValue(savedGraphsState);
 
   const { addTestCase, updateTestSuite, testGraph, setEditingTestCase, deleteTestCase, duplicateTestCase } =
@@ -218,6 +225,7 @@ export const TestSuite: FC<{ testSuite: TrivetTestSuite; tryRunTests: TryRunTest
   }, [testCaseValidationResults, testSuite, updateTestSuite]);
 
   const viewDocumentation = useOpenUrl('https://rivet.ironcladapp.com/docs/trivet');
+  const brainTrustSummary = brainTrustEnabled && brainTrustSummaries?.[testSuite.id];
 
   return (
     <div css={styles}>
@@ -290,6 +298,14 @@ export const TestSuite: FC<{ testSuite: TrivetTestSuite; tryRunTests: TryRunTest
                 <Button appearance="primary" onClick={() => tryRunTests({ testSuiteIds: [testSuite.id] })}>
                   Run Test Suite
                 </Button>
+
+                {brainTrustSummary && (
+                  <div style={{ marginLeft: 10 }}>
+                    <Button appearance="primary" onClick={useOpenUrl(brainTrustSummary.experimentUrl)}>
+                      View results in BrainTrust ({brainTrustSummary.experimentName})
+                    </Button>
+                  </div>
+                )}
               </div>
               <TestCaseTable
                 testCases={testSuite.testCases}

packages/app/src/hooks/useBuiltInPlugins.ts

@@ -1,23 +1,14 @@
+import { plugins } from '@ironclad/rivet-core';
 import { orderBy } from 'lodash-es';
 
-const builtInPlugins = orderBy(
-  [
-    {
-      label: 'Anthropic',
-      value: 'anthropic',
-    },
-    {
-      label: 'AssemblyAI',
-      value: 'assemblyAi',
-    },
-    {
-      label: 'Autoevals',
-      value: 'autoevals',
-    },
-  ] as const,
+const pluginOptions = orderBy(
+  Object.entries(plugins).map(([id, plugin]) => ({
+    value: id,
+    label: plugin.name ?? plugin.id,
+  })),
   'label',
 );
 
 export function useBuiltInPlugins() {
-  return builtInPlugins;
+  return pluginOptions;
 }

packages/app/src/hooks/useIsPluginEnabled.ts

@@ -0,0 +1,8 @@
+import { BuiltInPluginID } from '@ironclad/rivet-core';
+import { useDependsOnPlugins } from './useDependsOnPlugins';
+
+export function useIsPluginEnabled(plugin: BuiltInPluginID): boolean {
+  const plugins = useDependsOnPlugins();
+
+  return plugins.some((p) => p.id === plugin);
+}

packages/app/src/hooks/useLocalExecutor.ts

@@ -24,6 +24,7 @@ import { lastRecordingState, loadedRecordingState } from '../state/execution';
 import { fillMissingSettingsFromEnvironmentVariables } from '../utils/tauri';
 import { trivetState } from '../state/trivet';
 import { runTrivet } from '@ironclad/trivet';
+import { useDependsOnPlugins } from './useDependsOnPlugins';
 
 export function useLocalExecutor() {
   const project = useRecoilValue(projectState);
@@ -134,6 +135,8 @@ export function useLocalExecutor() {
     },
   );
 
+  const plugins = useDependsOnPlugins();
+
   const tryRunTests = useStableCallback(
     async (options: { testSuiteIds?: string[]; testCaseIds?: string[]; iterationCount?: number } = {}) => {
       toast.info(
@@ -158,6 +161,7 @@ export function useLocalExecutor() {
             }))
         : testSuites;
       try {
+        const settings = await fillMissingSettingsFromEnvironmentVariables(savedSettings, plugins);
         const result = await runTrivet({
           project,
           iterationCount: options.iterationCount,
@@ -182,6 +186,21 @@ export function useLocalExecutor() {
               inputs,
             );
           },
+          // braintrustApiKey: settings.braintrustApiKey,
+          // setBrainTrustSummary: (id, summary) =>
+          //   setTrivetState((s) => {
+          //     let brainTrustSummaries = s.brainTrustSummaries;
+          //     if (summary === undefined) {
+          //       const { [id]: _, ...rest } = s.brainTrustSummaries || {};
+          //       brainTrustSummaries = rest;
+          //     } else {
+          //       brainTrustSummaries = { ...s.brainTrustSummaries, [id]: summary };
+          //     }
+          //     return {
+          //       ...s,
+          //       brainTrustSummaries,
+          //     };
+          //   }),
         });
         setTrivetState((s) => ({
           ...s,

packages/app/src/state/plugins/braintrust.ts

@@ -0,0 +1,11 @@
+import { BTExperimentSummary } from '@ironclad/trivet';
+import { atom } from 'recoil';
+
+export type BraintrustSummariesState = {
+  brainTrustSummaries?: Record<string, BTExperimentSummary>;
+};
+
+export const braintrustSummariesState = atom<BraintrustSummariesState>({
+  key: 'braintrustSummariesState',
+  default: {},
+});

packages/core/src/plugins.ts

@@ -1,11 +1,15 @@
 import anthropicPlugin from './plugins/anthropic/index.js';
 import autoevalsPlugin from './plugins/autoevals/index.js';
 import assemblyAiPlugin from './plugins/assemblyAi/index.js';
+import { braintrustPlugin } from './plugins/braintrust/plugin.js';
 
-export { anthropicPlugin, autoevalsPlugin, assemblyAiPlugin };
+export { anthropicPlugin, autoevalsPlugin, assemblyAiPlugin, braintrustPlugin };
 
 export const plugins = {
   anthropic: anthropicPlugin,
   autoevals: autoevalsPlugin,
   assemblyAi: assemblyAiPlugin,
+  braintrust: braintrustPlugin,
 };
+
+export type BuiltInPluginID = keyof typeof plugins;

packages/core/src/plugins/braintrust/plugin.ts

@@ -0,0 +1,16 @@
+import { RivetPlugin } from '../../index.js';
+
+export const braintrustPlugin: RivetPlugin = {
+  id: 'braintrust',
+  name: 'Braintrust',
+
+  configSpec: {
+    braintrustApiKey: {
+      type: 'secret',
+      label: 'Braintrust API Key',
+      description: 'API key for Braintrust',
+      pullEnvironmentVariable: 'BRAINTRUST_API_KEY',
+      helperText: 'Alternatively, set the BRAINTRUST_API_KEY environment variable',
+    },
+  },
+};

packages/trivet/package.json

@@ -40,6 +40,8 @@
   },
   "dependencies": {
     "@ironclad/rivet-core": "workspace:^",
+    "braintrust": "^0.0.41",
+    "eventemitter3": "^5.0.1",
     "lodash-es": "^4.17.21",
     "yaml": "^2.3.1"
   },