Add codegen support for ChatMessageTrigger

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

@@ -38,6 +38,45 @@ 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.bottom_node import BottomNode
+
+
+class ChatMessage(ChatMessageTrigger):
+    message: str
+
+    class Config(ChatMessageTrigger.Config):
+        output = LazyReference(lambda: BottomNode.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 vellum.workflows.triggers.chat_message import ChatMessageTrigger
+
+from .nodes.bottom_node import BottomNode
+from .nodes.top_node import TopNode
+
+
+class Workflow(BaseWorkflow):
+    graph = {
+        TopNode,
+        ChatMessageTrigger >> BottomNode,
+    }
+"
+`;
+
 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,127 @@
+export default {
+  workflow_raw_data: {
+    nodes: [
+      {
+        id: "entrypoint-node",
+        type: "ENTRYPOINT",
+        data: {
+          label: "Entrypoint",
+          source_handle_id: "entrypoint-source",
+        },
+        inputs: [],
+      },
+      {
+        id: "top-node",
+        type: "GENERIC",
+        label: "Top Node",
+        display_data: null,
+        base: {
+          name: "BaseNode",
+          module: ["vellum", "workflows", "nodes", "bases", "base"],
+        },
+        definition: {
+          name: "TopNode",
+          module: ["testing", "nodes", "top_node"],
+        },
+        trigger: {
+          id: "top-target",
+          merge_behavior: "AWAIT_ATTRIBUTES",
+        },
+        ports: [
+          {
+            id: "top-default-port-id",
+            name: "default",
+            type: "DEFAULT",
+          },
+        ],
+        outputs: [],
+        attributes: [],
+      },
+      {
+        id: "bottom-node",
+        type: "GENERIC",
+        label: "Bottom Node",
+        display_data: null,
+        base: {
+          name: "BaseNode",
+          module: ["vellum", "workflows", "nodes", "bases", "base"],
+        },
+        definition: {
+          name: "BottomNode",
+          module: ["testing", "nodes", "bottom_node"],
+        },
+        trigger: {
+          id: "bottom-target",
+          merge_behavior: "AWAIT_ATTRIBUTES",
+        },
+        ports: [
+          {
+            id: "bottom-default-port-id",
+            name: "default",
+            type: "DEFAULT",
+          },
+        ],
+        outputs: [
+          {
+            id: "bottom-output-id",
+            name: "result",
+            type: "STRING",
+          },
+        ],
+        attributes: [],
+      },
+    ],
+    edges: [
+      {
+        id: "edge-1",
+        source_node_id: "entrypoint-node",
+        source_handle_id: "entrypoint-source",
+        target_node_id: "top-node",
+        target_handle_id: "top-target",
+        type: "DEFAULT",
+      },
+      {
+        id: "edge-2",
+        source_node_id: "chat-message-trigger",
+        source_handle_id: "chat-message-trigger",
+        target_node_id: "bottom-node",
+        target_handle_id: "bottom-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: "bottom-node",
+          node_output_id: "bottom-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: ["vellum", "workflows", "triggers", "chat_message"],
+    });
+  });
 });

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

@@ -0,0 +1,120 @@
+import { python } from "@fern-api/python-ast";
+
+import {
+  OUTPUTS_CLASS_NAME,
+  VELLUM_WORKFLOW_TRIGGERS_MODULE_PATH,
+} from "src/constants";
+import { AccessAttribute } from "src/generators/extensions/access-attribute";
+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 { createPythonClassName, toPythonSafeSnakeCase } from "src/utils/casing";
+
+import type { AstNode } from "src/generators/extensions/ast-node";
+import type { ChatMessageTrigger as ChatMessageTriggerType } 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 && output.type === "NODE_OUTPUT") {
+      const nodeContext = this.workflowContext.findNodeContext(output.nodeId);
+
+      if (nodeContext) {
+        const nodeOutputName = nodeContext.getNodeOutputNameById(
+          output.nodeOutputId
+        );
+
+        if (nodeOutputName) {
+          const lazyReferenceValue = new ClassInstantiation({
+            classReference: new Reference({
+              name: "LazyReference",
+              modulePath: [
+                ...this.workflowContext.sdkModulePathNames
+                  .WORKFLOWS_MODULE_PATH,
+                "references",
+              ],
+            }),
+            arguments_: [
+              new MethodArgument({
+                value: python.lambda({
+                  body: new AccessAttribute({
+                    lhs: new Reference({
+                      name: nodeContext.nodeClassName,
+                      modulePath: nodeContext.nodeModulePath,
+                    }),
+                    rhs: new Reference({
+                      name: `${OUTPUTS_CLASS_NAME}.${nodeOutputName}`,
+                      modulePath: [],
+                    }),
+                  }),
+                }),
+              }),
+            ],
+          });
+
+          configClass.add(
+            new Field({
+              name: "output",
+              initializer: lazyReferenceValue,
+            })
+          );
+        }
+      }
+    }
+
+    return configClass;
+  }
+}

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" && trigger.execConfig) {
+        const chatMessageTrigger = new ChatMessageTrigger({
+          workflowContext: this.workflowContext,
+          trigger,
+        });
+        triggerPromises.push(chatMessageTrigger.persist());
       }
     });
 

ee/codegen/src/serializers/vellum.ts

@@ -2435,17 +2435,54 @@ export declare namespace IntegrationTriggerSerializer {
   }
 }
 
+const ChatMessageTriggerExecConfigSerializer = objectSchema({
+  output: WorkflowValueDescriptorSerializer.optional(),
+});
+
+export declare namespace ChatMessageTriggerExecConfigSerializer {
+  interface Raw {
+    output?: WorkflowValueDescriptorSerializer.Raw | null;
+  }
+}
+
+const ChatMessageTriggerSerializer = objectSchema({
+  id: stringSchema(),
+  attributes: listSchema(VellumVariableSerializer),
+  definition: CodeResourceDefinitionSerializer.optional(),
+  displayData: propertySchema(
+    "display_data",
+    WorkflowTriggerDisplayDataSerializer.nullable().optional()
+  ),
+  execConfig: propertySchema(
+    "exec_config",
+    ChatMessageTriggerExecConfigSerializer.optional()
+  ),
+});
+
+export declare namespace ChatMessageTriggerSerializer {
+  interface Raw {
+    id: string;
+    type: "CHAT_MESSAGE";
+    attributes: VellumVariableSerializer.Raw[];
+    definition?: CodeResourceDefinitionSerializer.Raw | null;
+    display_data?: WorkflowTriggerDisplayDataSerializer.Raw | null;
+    exec_config?: ChatMessageTriggerExecConfigSerializer.Raw | null;
+  }
+}
+
 export const WorkflowTriggerSerializer = unionSchema("type", {
   MANUAL: ManualTriggerSerializer,
   SCHEDULED: ScheduleTriggerSerializer,
   INTEGRATION: IntegrationTriggerSerializer,
+  CHAT_MESSAGE: ChatMessageTriggerSerializer,
 }) as unknown as Schema<WorkflowTriggerSerializer.Raw, WorkflowTrigger>;
 
 export declare namespace WorkflowTriggerSerializer {
   type Raw =
     | ManualTriggerSerializer.Raw
     | ScheduleTriggerSerializer.Raw
-    | IntegrationTriggerSerializer.Raw;
+    | IntegrationTriggerSerializer.Raw
+    | ChatMessageTriggerSerializer.Raw;
 }
 
 export const WorkflowRawDataSerializer: ObjectSchema<