generalise actions

examples/example.ts

@@ -15,6 +15,8 @@ import * as readline from "readline";
 import { type GoalStatus } from "../packages/core/src/core/goal-manager";
 import chalk from "chalk";
 import { LogLevel } from "../packages/core/src/types";
+import { starknetTransactionAction } from "../packages/core/src/core/actions/starknet-transaction";
+import { graphqlAction } from "../packages/core/src/core/actions/graphql";
 
 async function getCliInput(prompt: string): Promise<string> {
   const rl = readline.createInterface({
@@ -61,6 +63,24 @@ async function main() {
     availableActions: "",
   });
 
+  // Register actions
+  dreams.registerAction("EXECUTE_TRANSACTION", starknetTransactionAction, {
+    description: "Execute a transaction on the Starknet blockchain",
+    example: JSON.stringify({
+      contractAddress: "0x1234567890abcdef",
+      entrypoint: "execute",
+      calldata: [1, 2, 3],
+    }),
+  });
+
+  dreams.registerAction("GRAPHQL_FETCH", graphqlAction, {
+    description: "Fetch data from the Eternum GraphQL API",
+    example: JSON.stringify({
+      query:
+        "query GetRealmInfo { eternumRealmModels(where: { realm_id: 42 }) { edges { node { ... on eternum_Realm { entity_id level } } } }",
+    }),
+  });
+
   // Subscribe to events
   dreams.on("step", (step) => {
     if (step.type === "system") {

packages/core/src/core/actions/graphql.ts

@@ -0,0 +1,10 @@
+import type { ActionHandler } from "../../types";
+import { fetchData } from "../providers";
+
+export const graphqlAction: ActionHandler = async (action, chain) => {
+  const { query, variables } = action.payload || {};
+  const result = await fetchData(query, variables);
+  const resultStr =
+    `query: ` + query + `\n\nresult: ` + JSON.stringify(result, null, 2);
+  return resultStr;
+};

packages/core/src/core/actions/starknet-transaction.ts

@@ -0,0 +1,17 @@
+import type { ActionHandler } from "../../types";
+import { executeStarknetTransaction } from "../providers";
+import type { CoTTransaction } from "../../types";
+
+export const starknetTransactionAction: ActionHandler = async (
+  action,
+  chain
+) => {
+  const result = await executeStarknetTransaction(
+    action.payload as CoTTransaction
+  );
+  return `Transaction executed successfully: ${JSON.stringify(
+    result,
+    null,
+    2
+  )}`;
+};

packages/core/src/core/actions/system-prompt.ts

@@ -0,0 +1,22 @@
+import type { ActionHandler } from "../../types";
+
+import * as readline from "readline";
+
+async function askUser(question: string): Promise<string> {
+  const rl = readline.createInterface({
+    input: process.stdin,
+    output: process.stdout,
+  });
+
+  return new Promise((resolve) => {
+    rl.question(`${question}\nYour response: `, (answer) => {
+      rl.close();
+      resolve(answer);
+    });
+  });
+}
+
+export const systemPromptAction: ActionHandler = async (action, chain) => {
+  const userResponse = await askUser(action.payload.prompt);
+  return userResponse;
+};

packages/core/src/core/chain-of-thought.ts

@@ -1,37 +1,22 @@
 import type { LLMClient } from "./llm-client";
 import type {
+  ActionHandler,
   ChainOfThoughtContext,
   CoTAction,
-  CoTTransaction,
   LLMStructuredResponse,
 } from "../types";
 import { queryValidator } from "./validation";
 import { Logger } from "./logger";
-import { executeStarknetTransaction, fetchData } from "./providers";
 import { EventEmitter } from "events";
 import { GoalManager, type HorizonType, type GoalStatus } from "./goal-manager";
 import { StepManager, type Step, type StepType } from "./step-manager";
 import { LogLevel } from "../types";
+import { injectTags } from "./utils";
+import { systemPromptAction } from "./actions/system-prompt";
 
 // Todo: remove these when we bundle
 import * as fs from "fs";
 import * as path from "path";
-import * as readline from "readline";
-import { injectTags } from "./utils";
-
-async function askUser(question: string): Promise<string> {
-  const rl = readline.createInterface({
-    input: process.stdin,
-    output: process.stdout,
-  });
-
-  return new Promise((resolve) => {
-    rl.question(`${question}\nYour response: `, (answer) => {
-      rl.close();
-      resolve(answer);
-    });
-  });
-}
 
 export class ChainOfThought extends EventEmitter {
   private stepManager: StepManager;
@@ -41,11 +26,17 @@ export class ChainOfThought extends EventEmitter {
   private contextLogPath: string;
   goalManager: GoalManager;
 
+  private actionRegistry = new Map<string, ActionHandler>();
+  private actionExamples = new Map<
+    string,
+    { description: string; example: string }
+  >();
+
   constructor(
     private llmClient: LLMClient,
     initialContext?: ChainOfThoughtContext
   ) {
-    super(); // Initialize EventEmitter
+    super();
     this.stepManager = new StepManager();
     this.context = initialContext ?? {
       worldState: "",
@@ -69,6 +60,12 @@ export class ChainOfThought extends EventEmitter {
     this.logContext("INITIAL_CONTEXT");
 
     this.goalManager = new GoalManager();
+
+    this.registerDefaultActions();
+  }
+
+  private registerDefaultActions() {
+    this.actionRegistry.set("SYSTEM_PROMPT", systemPromptAction);
   }
 
   public async planStrategy(objective: string): Promise<void> {
@@ -667,6 +664,24 @@ export class ChainOfThought extends EventEmitter {
     return this.snapshots;
   }
 
+  /**
+   * Allows registration of a custom action type and handler.
+   * e.g. chainOfThought.registerAction("SEND_EMAIL", sendEmailHandler);
+   */
+  public registerAction(
+    type: string,
+    handler: ActionHandler,
+    example?: { description: string; example: string }
+  ): void {
+    this.logger.debug("registerAction", "Registering custom action", { type });
+    this.actionRegistry.set(type, handler);
+
+    // Store the example snippet if provided
+    if (example) {
+      this.actionExamples.set(type, example);
+    }
+  }
+
   /**
    * A central method to handle CoT actions that might be triggered by the LLM.
    * @param action The action to be executed.
@@ -675,37 +690,26 @@ export class ChainOfThought extends EventEmitter {
     this.logger.debug("executeAction", "Executing action", { action });
     this.emit("action:start", action);
 
-    // Add action step at the start
+    // Add action step
     const actionStep = this.addStep(
       `Executing action: ${action.type}`,
       "action",
       ["action-execution"],
       { action }
     );
 
-    console.log("log", action);
-
     try {
-      const result = await (async () => {
-        switch (action.type) {
-          case "SYSTEM_PROMPT":
-            // Handle system prompt by asking user
-            const userResponse = await askUser(action.payload.prompt);
-            return userResponse;
-
-          case "GRAPHQL_FETCH":
-            return await this.graphqlFetchAction(action.payload);
-
-          case "EXECUTE_TRANSACTION":
-            return this.runTransaction(action.payload as CoTTransaction);
-
-          default:
-            this.logger.warn("executeAction", "Unknown action type", {
-              actionType: action.type,
-            });
-            return "Unknown action type: " + action.type;
-        }
-      })();
+      // Lookup the handler in our registry
+      const handler = this.actionRegistry.get(action.type);
+
+      if (!handler) {
+        // No handler found
+        const errorMsg = `No handler registered for action type "${action.type}"`;
+        throw new Error(errorMsg);
+      }
+
+      // If found, run it
+      const result = await handler(action, this);
 
       // Update the action step with the result
       this.stepManager.updateStep(actionStep.id, {
@@ -732,60 +736,11 @@ export class ChainOfThought extends EventEmitter {
         content: `Action failed: ${error}`,
         meta: { ...actionStep.meta, error },
       });
-
       this.emit("action:error", { action, error });
       throw error;
     }
   }
 
-  private async graphqlFetchAction(
-    payload?: Record<string, any>
-  ): Promise<string> {
-    this.logger.debug("graphqlFetchAction", "Executing GraphQL fetch", {
-      payload,
-    });
-
-    // Example of expected fields in the payload
-    const { query, variables } = payload || {};
-
-    const result = await fetchData(query, variables);
-
-    const resultStr =
-      `query: ` + query + `\n\nresult: ` + JSON.stringify(result, null, 2);
-
-    return resultStr;
-  }
-
-  /**
-   * Execute a "transaction" that can modify the chain of thought, the context,
-   * or even the game state. The specifics are up to your design.
-   */
-  private async runTransaction(transaction: CoTTransaction): Promise<string> {
-    this.logger.debug("runTransaction", "Running transaction", { transaction });
-
-    // Add step describing the transaction
-
-    const result = await executeStarknetTransaction(transaction);
-
-    const resultStr = `Transaction executed successfully: ${JSON.stringify(
-      result,
-      null,
-      2
-    )}`;
-
-    this.addStep(
-      `Running transaction: ${transaction.contractAddress}`,
-      "action",
-      ["transaction"],
-      {
-        transactionData: transaction.calldata,
-        result: resultStr,
-      }
-    );
-
-    return resultStr;
-  }
-
   /**
    * Generate a unique ID for steps. In a real environment,
    * you might use a library like uuid to ensure uniqueness.
@@ -891,6 +846,16 @@ export class ChainOfThought extends EventEmitter {
     throw new Error(error);
   }
 
+  /**
+   * Returns a formatted string listing all available actions registered in the action registry
+   */
+  private getAvailableActions(): string {
+    const actions = Array.from(this.actionRegistry.keys());
+    return `Available actions:\n${actions
+      .map((action) => `- ${action}`)
+      .join("\n")}`;
+  }
+
   /**
    * Build a prompt that instructs the LLM to produce structured data.
    * You can adapt the instructions, tone, or style as needed.
@@ -902,6 +867,17 @@ export class ChainOfThought extends EventEmitter {
     // For example, let's just send the last few steps:
     const lastSteps = JSON.stringify(this.stepManager.getSteps());
 
+    const actionExamplesText = Array.from(this.actionExamples.entries())
+      .map(([type, { description, example }]) => {
+        return `
+Action Type: ${type}
+Description: ${description}
+Example JSON:
+${example}
+`;
+      })
+      .join("\n\n");
+
     // Replace any {{tag}} with the provided value or leave unchanged
 
     const prompt = `
@@ -964,32 +940,30 @@ ${this.context.availableActions}
 Return a JSON array where each step contains:
 - plan: A short explanation of what you will do
 - meta: A metadata object with requirements for the step. Find this in the context.
-- actions: A list of actions to be executed. You can either use GRAPHQL_FETCH or EXECUTE_TRANSACTION. You must only use these.
+- actions: A list of actions to be executed. You can either use ${this.getAvailableActions()}. You must only use these.
+
+<AVAILABLE_ACTIONS>
+Below is a list of actions you may use. For each action, 
+the "payload" must follow the indicated structure exactly.
+
+${actionExamplesText}
+</AVAILABLE_ACTIONS>
 
 {{output_format_details}}
 
-Provide a JSON response with the following structure (this is an example only):
+Provide a JSON response with the following structure (this is an example only, use the available actions to determine the structure):
 
 {
   "plan": "A short explanation of what you will do",
       "meta": {
       "requirements": {},
-  "actions": [
-        {
-          type: "GRAPHQL_FETCH",
-          payload: {
-            query: "query GetRealmInfo { eternumRealmModels(where: { realm_id: 42 }) { edges { node { ... on eternum_Realm { entity_id level } } } }",
-          },
-        },
-        {
-          type: "EXECUTE_TRANSACTION",
-          payload: {
-              contractAddress: "0x1234567890abcdef",
-              entrypoint: "execute",
-              calldata: [1, 2, 3],
-          },
-        },
+        "actions": [
+    {
+      "type": "ACTION_TYPE",
+      "payload": { ... }
+    }
   ]
+
 }
 
 </OUTPUT_FORMAT>

packages/core/src/types/index.ts

@@ -1,3 +1,5 @@
+import type { ChainOfThought } from "../core/chain-of-thought";
+
 // Base event type
 export interface BaseEvent {
   type: string;
@@ -193,3 +195,9 @@ export interface LogEntry {
   message: string;
   data?: any;
 }
+
+export type ActionHandler = (
+  action: CoTAction,
+  chain: ChainOfThought,
+  example?: { description: string; example: string }
+) => Promise<string>;