Update AssemblyAI LeMUR nodes

Ironclad · Jul 15, 2024 · d3786e5 · d3786e5
1 parent 5d36a49
commit d3786e5
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 18 deletions.
diff --git a/packages/core/src/plugins/assemblyAi/LemurActionItemsNode.ts b/packages/core/src/plugins/assemblyAi/LemurActionItemsNode.ts
@@ -1,6 +1,6 @@
 import { nanoid } from 'nanoid/non-secure';
 import { dedent } from 'ts-dedent';
-import { AssemblyAI, type LemurActionItemsParams } from 'assemblyai';
+import type { LemurActionItemsParams } from 'assemblyai';
 import {
   type ChartNode,
   type EditorDefinition,
@@ -20,7 +20,7 @@ import {
   getClient,
   getLemurParams,
   lemurEditorDefinitions,
-  lemurTranscriptIdsInputDefinition,
+  lemurInputDefinitions,
 } from './lemurHelpers.js';
 
 export type LemurActionItemsNode = ChartNode<'assemblyAiLemurActionItems', LemurActionItemsNodeData>;
@@ -48,7 +48,7 @@ export const LemurActionItemsNodeImpl: PluginNodeImpl<LemurActionItemsNode> = {
 
   getInputDefinitions(): NodeInputDefinition[] {
     return [
-      lemurTranscriptIdsInputDefinition,
+      ...lemurInputDefinitions,
       {
         id: 'context' as PortId,
         dataType: 'string',

diff --git a/packages/core/src/plugins/assemblyAi/LemurQaNode.ts b/packages/core/src/plugins/assemblyAi/LemurQaNode.ts
@@ -23,7 +23,7 @@ import {
   getClient,
   getLemurParams,
   lemurEditorDefinitions,
-  lemurTranscriptIdsInputDefinition,
+  lemurInputDefinitions,
 } from './lemurHelpers.js';
 import { coerceType } from '../../utils/coerceType.js';
 import { pluginNodeDefinition } from '../../model/NodeDefinition.js';
@@ -58,7 +58,7 @@ export const LemurQaNodeImpl = {
 
   getInputDefinitions(): NodeInputDefinition[] {
     return [
-      lemurTranscriptIdsInputDefinition,
+      ...lemurInputDefinitions,
       {
         id: 'questions' as PortId,
         dataType: ['string', 'string[]', 'object', 'object[]', 'any', 'any[]'],

diff --git a/packages/core/src/plugins/assemblyAi/LemurSummaryNode.ts b/packages/core/src/plugins/assemblyAi/LemurSummaryNode.ts
@@ -19,7 +19,7 @@ import {
   getClient,
   getLemurParams,
   lemurEditorDefinitions,
-  lemurTranscriptIdsInputDefinition,
+  lemurInputDefinitions,
 } from './lemurHelpers.js';
 import { pluginNodeDefinition } from '../../model/NodeDefinition.js';
 
@@ -50,7 +50,7 @@ export const LemurSummaryNodeImpl: PluginNodeImpl<LemurSummaryNode> = {
 
   getInputDefinitions(): NodeInputDefinition[] {
     return [
-      lemurTranscriptIdsInputDefinition,
+      ...lemurInputDefinitions,
       {
         id: 'context' as PortId,
         dataType: 'string',

diff --git a/packages/core/src/plugins/assemblyAi/LemurTaskNode.ts b/packages/core/src/plugins/assemblyAi/LemurTaskNode.ts
@@ -1,6 +1,6 @@
 import { nanoid } from 'nanoid/non-secure';
 import { dedent } from 'ts-dedent';
-import { AssemblyAI, type LemurTaskParams } from 'assemblyai';
+import { type LemurTaskParams } from 'assemblyai';
 import {
   type ChartNode,
   type EditorDefinition,
@@ -19,7 +19,7 @@ import {
   getClient,
   getLemurParams,
   lemurEditorDefinitions,
-  lemurTranscriptIdsInputDefinition,
+  lemurInputDefinitions,
 } from './lemurHelpers.js';
 import { pluginNodeDefinition } from '../../model/NodeDefinition.js';
 import { coerceTypeOptional } from '../../utils/coerceType.js';
@@ -51,7 +51,7 @@ export const LemurTaskNodeImpl: PluginNodeImpl<LemurTaskNode> = {
 
   getInputDefinitions(): NodeInputDefinition[] {
     return [
-      lemurTranscriptIdsInputDefinition,
+      ...lemurInputDefinitions,
       {
         id: 'prompt' as PortId,
         dataType: 'string',

diff --git a/packages/core/src/plugins/assemblyAi/lemurHelpers.ts b/packages/core/src/plugins/assemblyAi/lemurHelpers.ts
@@ -11,13 +11,13 @@ import {
 } from '../../index.js';
 import { coerceType, coerceTypeOptional } from '../../utils/coerceType.js';
 
-function getTranscriptIds(inputs: Inputs): string[] {
+function getTranscriptIds(inputs: Inputs): string[] | undefined {
   const input = inputs['transcript_ids' as PortId] as
     | StringDataValue
     | StringArrayDataValue
     | AnyDataValue
     | ArrayDataValue<AnyDataValue>;
-  if (!input) throw new Error('Transcript IDs are required.');
+  if (!input) return undefined;
 
   if (input.type === 'string' || (input.type === 'any' && typeof input.value === 'string')) {
     return [coerceType(input, 'string')];
@@ -48,47 +48,106 @@ export function getClient(context: InternalProcessContext): AssemblyAI {
 export function getLemurParams(inputs: Inputs, editorData: LemurNodeData): LemurBaseParams {
   const params: LemurBaseParams = {
     transcript_ids: getTranscriptIds(inputs),
+    input_text: coerceTypeOptional(inputs['input_text' as PortId], 'string'),
     context: coerceTypeOptional(inputs['context' as PortId], 'string') || editorData.context || undefined,
     final_model: editorData.final_model && editorData.final_model !== 'default' ? editorData.final_model : undefined,
     max_output_size: editorData.max_output_size,
+    temperature: editorData.temperature,
   };
 
   return params;
 }
 
-export const lemurTranscriptIdsInputDefinition = {
-  id: 'transcript_ids' as PortId,
-  dataType: ['string', 'string[]', 'any', 'any[]'],
-  title: 'Transcript IDs',
-} as const;
+export const lemurInputDefinitions = [
+  {
+    id: 'transcript_ids' as PortId,
+    dataType: ['string', 'string[]', 'any', 'any[]'],
+    title: 'Transcript IDs',
+  },
+  {
+    id: 'input_text' as PortId,
+    dataType: ['string'],
+    title: 'Input Text',
+  },
+] as const;
 
 export const lemurEditorDefinitions = [
   {
     type: 'dropdown',
     label: 'Final Model',
     dataKey: 'final_model',
     options: [
+      {
+        value: 'anthropic/claude-3-5-sonnet',
+        label: 'Claude 3.5 Sonnet (on Anthropic)',
+      },
+      {
+        value: 'anthropic/claude-3-opus',
+        label: 'Claude 3 Opus (on Anthropic)',
+      },
+      {
+        value: 'anthropic/claude-3-haiku',
+        label: 'Claude 3 Haiku (on Anthropic)',
+      },
+      {
+        value: 'anthropic/claude-3-sonnet',
+        label: 'Claude 3 Sonnet (on Anthropic)',
+      },
+      {
+        value: 'anthropic/claude-2-1',
+        label: 'Claude 2.1 (on Anthropic)',
+      },
+      {
+        value: 'anthropic/claude-2',
+        label: 'Claude 2.1 (on Anthropic)',
+      },
       {
         value: 'default',
         label: 'Default',
       },
+      {
+        value: 'anthropic/claude-instant-1-2',
+        label: 'Claude Instant 1.2 (on Anthropic)',
+      },
       {
         value: 'basic',
         label: 'Basic',
       },
+      {
+        value: 'assemblyai/mistral-7b',
+        label: 'Mistral 7B (hosted by AssemblyAI)',
+      },
     ],
   },
   {
     type: 'number',
     label: 'Maximum Output Size',
     dataKey: 'max_output_size',
   },
+  {
+    type: 'number',
+    label: 'Temperature',
+    dataKey: 'temperature',
+    min: 0,
+    max: 1,
+  },
 ] as const;
 
-export type FinalModel = 'default' | 'basic';
+export type FinalModel =
+  | 'anthropic/claude-3-5-sonnet'
+  | 'anthropic/claude-3-opus'
+  | 'anthropic/claude-3-haiku'
+  | 'anthropic/claude-3-sonnet'
+  | 'anthropic/claude-2-1'
+  | 'anthropic/claude-2'
+  | 'default'
+  | 'anthropic/claude-instant-1-2'
+  | 'basic'
+  | 'assemblyai/mistral-7b';
 
 export type LemurNodeData = {
   context?: string;
   final_model?: FinalModel;
   max_output_size?: number;
+  temperature?: number;
 };