Add codegen support for ChatMessageTrigger

ee/codegen/src/__test__/__snapshots__/generate-code.test.ts.snap

@@ -38,6 +38,41 @@ class UseApiWithSecret(FinalOutputNode[BaseState, Any]):
 "
 `;
 
+exports[`generateCode > should generate code for %1 chat-message-trigger.ts > triggers/chat_message.py 1`] = `
+"from vellum.workflows.references import LazyReference
+from vellum.workflows.triggers import ChatMessageTrigger
+
+from ..nodes.some_node import SomeNode
+
+
+class ChatMessage(ChatMessageTrigger):
+    message: str
+
+    class Config(ChatMessageTrigger.Config):
+        output = LazyReference(lambda: SomeNode.Outputs.result)
+
+    class Display(ChatMessageTrigger.Display):
+        label = "Chat Message"
+        x = 100
+        y = 200
+        z_index = 1
+        icon = "vellum:icon:message"
+        color = "blue"
+"
+`;
+
+exports[`generateCode > should generate code for %1 chat-message-trigger.ts > workflow.py 1`] = `
+"from vellum.workflows import BaseWorkflow
+
+from .nodes.some_node import SomeNode
+from .triggers.chat_message import ChatMessage
+
+
+class Workflow(BaseWorkflow):
+    graph = ChatMessage >> SomeNode
+"
+`;
+
 exports[`generateCode > should generate code for %1 code-execution-node-with-await-all.ts > nodes/code_execution_with_await_all/__init__.py 1`] = `
 "from vellum.workflows.nodes.displayable import CodeExecutionNode
 from vellum.workflows.state import BaseState

ee/codegen/src/__test__/generate-code-fixtures/chat-message-trigger.ts

@@ -0,0 +1,83 @@
+export default {
+  workflow_raw_data: {
+    nodes: [
+      {
+        id: "some-node",
+        type: "GENERIC",
+        label: "Some Node",
+        display_data: null,
+        base: {
+          name: "BaseNode",
+          module: ["vellum", "workflows", "nodes", "bases", "base"],
+        },
+        definition: {
+          name: "SomeNode",
+          module: ["testing", "nodes", "some_node"],
+        },
+        trigger: {
+          id: "some-target",
+          merge_behavior: "AWAIT_ATTRIBUTES",
+        },
+        ports: [
+          {
+            id: "some-default-port-id",
+            name: "default",
+            type: "DEFAULT",
+          },
+        ],
+        outputs: [
+          {
+            id: "some-output-id",
+            name: "result",
+            type: "STRING",
+          },
+        ],
+        attributes: [],
+      },
+    ],
+    edges: [
+      {
+        id: "edge-2",
+        source_node_id: "chat-message-trigger",
+        source_handle_id: "chat-message-trigger",
+        target_node_id: "some-node",
+        target_handle_id: "some-target",
+        type: "DEFAULT",
+      },
+    ],
+    output_values: [],
+  },
+  input_variables: [],
+  output_variables: [],
+  triggers: [
+    {
+      id: "chat-message-trigger",
+      type: "CHAT_MESSAGE",
+      attributes: [
+        {
+          id: "message-attribute-id",
+          key: "message",
+          type: "JSON",
+        },
+      ],
+      exec_config: {
+        output: {
+          type: "NODE_OUTPUT",
+          node_id: "some-node",
+          node_output_id: "some-output-id",
+        },
+      },
+      display_data: {
+        label: "Chat Message",
+        position: {
+          x: 100,
+          y: 200,
+        },
+        z_index: 1,
+        icon: "vellum:icon:message",
+        color: "blue",
+      },
+    },
+  ],
+  assertions: ["workflow.py", "triggers/chat_message.py"],
+};

ee/codegen/src/__test__/utils/triggers.test.ts

@@ -69,4 +69,19 @@ describe("getTriggerClassInfo", () => {
       modulePath: ["tests", "fixtures", "triggers", "slack_new_message"],
     });
   });
+
+  it("should return correct info for CHAT_MESSAGE trigger", () => {
+    const trigger: WorkflowTrigger = {
+      id: "chat-message-trigger-id",
+      type: WorkflowTriggerType.CHAT_MESSAGE,
+      attributes: [{ id: "attr-1", type: "JSON", key: "message" }],
+    };
+
+    const result = getTriggerClassInfo(trigger, workflowContextFactory());
+
+    expect(result).toEqual({
+      className: "ChatMessageTrigger",
+      modulePath: ["code", "triggers", "chat_message"],
+    });
+  });
 });

ee/codegen/src/context/trigger-context/chat-message-trigger.ts

@@ -0,0 +1,37 @@
+import { GENERATED_TRIGGERS_MODULE_NAME } from "src/constants";
+import { BaseTriggerContext } from "src/context/trigger-context/base";
+import { ChatMessageTrigger } from "src/types/vellum";
+import { createPythonClassName, toPythonSafeSnakeCase } from "src/utils/casing";
+
+export class ChatMessageTriggerContext extends BaseTriggerContext<ChatMessageTrigger> {
+  protected getTriggerModuleInfo(): {
+    moduleName: string;
+    className: string;
+    modulePath: string[];
+  } {
+    const label = this.triggerData.displayData?.label || "chat_message";
+    const rawModuleName = toPythonSafeSnakeCase(label);
+    let moduleName = rawModuleName;
+    let numRenameAttempts = 0;
+    while (this.workflowContext.isTriggerModuleNameUsed(moduleName)) {
+      moduleName = `${rawModuleName}_${numRenameAttempts + 1}`;
+      numRenameAttempts += 1;
+    }
+    const className = createPythonClassName(
+      this.triggerData.displayData?.label || "ChatMessageTrigger",
+      { force: true }
+    );
+
+    const modulePath = [
+      ...this.workflowContext.modulePath.slice(0, -1),
+      GENERATED_TRIGGERS_MODULE_NAME,
+      moduleName,
+    ];
+
+    return {
+      moduleName,
+      className,
+      modulePath,
+    };
+  }
+}

ee/codegen/src/context/trigger-context/create-trigger-context.ts

@@ -1,7 +1,9 @@
 import { WorkflowContext } from "src/context";
+import { ChatMessageTriggerContext } from "src/context/trigger-context/chat-message-trigger";
 import { IntegrationTriggerContext } from "src/context/trigger-context/integration-trigger";
 import { ScheduledTriggerContext } from "src/context/trigger-context/scheduled-trigger";
 import {
+  ChatMessageTrigger,
   IntegrationTrigger,
   ScheduledTrigger,
   WorkflowTrigger,
@@ -37,6 +39,14 @@ export function createTriggerContext({
       workflowContext.addTriggerContext(triggerContext);
       break;
     }
+    case "CHAT_MESSAGE": {
+      const triggerContext = new ChatMessageTriggerContext({
+        workflowContext,
+        triggerData: triggerData as ChatMessageTrigger,
+      });
+      workflowContext.addTriggerContext(triggerContext);
+      break;
+    }
     case "MANUAL":
       // For now, we don't create contexts for MANUAL triggers
       // as they don't have associated files

ee/codegen/src/context/trigger-context/index.ts

@@ -1,4 +1,5 @@
 export { BaseTriggerContext } from "./base";
+export { ChatMessageTriggerContext } from "./chat-message-trigger";
 export { createTriggerContext } from "./create-trigger-context";
 export { IntegrationTriggerContext } from "./integration-trigger";
 export { ScheduledTriggerContext } from "./scheduled-trigger";

ee/codegen/src/generators/triggers/chat-message-trigger.ts

@@ -0,0 +1,114 @@
+import { python } from "@fern-api/python-ast";
+
+import { VELLUM_WORKFLOW_TRIGGERS_MODULE_PATH } from "src/constants";
+import { Class } from "src/generators/extensions/class";
+import { ClassInstantiation } from "src/generators/extensions/class-instantiation";
+import { Field } from "src/generators/extensions/field";
+import { MethodArgument } from "src/generators/extensions/method-argument";
+import { Reference } from "src/generators/extensions/reference";
+import { BaseTrigger } from "src/generators/triggers/base-trigger";
+import { WorkflowValueDescriptor } from "src/generators/workflow-value-descriptor";
+import { createPythonClassName, toPythonSafeSnakeCase } from "src/utils/casing";
+
+import type { AstNode } from "src/generators/extensions/ast-node";
+import type {
+  ChatMessageTrigger as ChatMessageTriggerType,
+  WorkflowValueDescriptor as WorkflowValueDescriptorType,
+} from "src/types/vellum";
+
+export declare namespace ChatMessageTriggerGenerator {
+  interface Args {
+    workflowContext: BaseTrigger.Args<ChatMessageTriggerType>["workflowContext"];
+    trigger: ChatMessageTriggerType;
+  }
+}
+
+export class ChatMessageTrigger extends BaseTrigger<ChatMessageTriggerType> {
+  protected generateClassName(): string {
+    const label = this.trigger.displayData?.label || "ChatMessageTrigger";
+    return createPythonClassName(label, {
+      force: true,
+    });
+  }
+
+  protected getModuleName(): string {
+    const label = this.trigger.displayData?.label || "chat_message";
+    return toPythonSafeSnakeCase(label);
+  }
+
+  protected getBaseTriggerClassName(): string {
+    return "ChatMessageTrigger";
+  }
+
+  protected getTriggerClassBody(): AstNode[] {
+    const body: AstNode[] = [];
+
+    // Add attribute fields
+    body.push(...this.createAttributeFields());
+
+    // Create Config class if execConfig.output is present
+    const execConfig = this.trigger.execConfig;
+    if (execConfig?.output) {
+      body.push(this.createConfigClass(execConfig.output));
+    }
+
+    return body;
+  }
+
+  private createConfigClass(
+    output: NonNullable<ChatMessageTriggerType["execConfig"]>["output"]
+  ): AstNode {
+    const configClass = new Class({
+      name: "Config",
+      extends_: [
+        new Reference({
+          name: "ChatMessageTrigger",
+          modulePath: VELLUM_WORKFLOW_TRIGGERS_MODULE_PATH,
+          attribute: ["Config"],
+        }),
+      ],
+    });
+
+    if (output) {
+      const outputField = this.createOutputField(output);
+      if (outputField) {
+        configClass.add(outputField);
+      }
+    }
+
+    return configClass;
+  }
+
+  private createOutputField(
+    output: WorkflowValueDescriptorType
+  ): AstNode | undefined {
+    const workflowValueDescriptor = new WorkflowValueDescriptor({
+      workflowContext: this.workflowContext,
+      workflowValueDescriptor: output,
+    });
+
+    const lazyReferenceValue = new ClassInstantiation({
+      classReference: new Reference({
+        name: "LazyReference",
+        modulePath: [
+          ...this.workflowContext.sdkModulePathNames.WORKFLOWS_MODULE_PATH,
+          "references",
+        ],
+      }),
+      arguments_: [
+        new MethodArgument({
+          value: python.lambda({
+            body: workflowValueDescriptor,
+          }),
+        }),
+      ],
+    });
+
+    this.inheritReferences(lazyReferenceValue);
+
+    return new Field({
+      name: "output",
+      initializer: lazyReferenceValue,
+    });
+  }
+}

ee/codegen/src/project.ts

@@ -61,6 +61,7 @@ import { PromptDeploymentNode } from "src/generators/nodes/prompt-deployment-nod
 import { SearchNode } from "src/generators/nodes/search-node";
 import { SubworkflowDeploymentNode } from "src/generators/nodes/subworkflow-deployment-node";
 import { TemplatingNode } from "src/generators/nodes/templating-node";
+import { ChatMessageTrigger } from "src/generators/triggers/chat-message-trigger";
 import { IntegrationTrigger } from "src/generators/triggers/integration-trigger";
 import { ScheduledTrigger } from "src/generators/triggers/scheduled-trigger";
 import { WorkflowSandboxFile } from "src/generators/workflow-sandbox-file";
@@ -926,6 +927,12 @@ ${errors.slice(0, 3).map((err) => {
           trigger,
         });
         triggerPromises.push(scheduledTrigger.persist());
+      } else if (trigger.type === "CHAT_MESSAGE") {
+        const chatMessageTrigger = new ChatMessageTrigger({
+          workflowContext: this.workflowContext,
+          trigger,
+        });
+        triggerPromises.push(chatMessageTrigger.persist());
       }
     });