♻️ Refactor and add retries

next/__tests__/message-service.test.ts

@@ -8,7 +8,6 @@ describe("sendErrorMessage", () => {
   beforeEach(() => {
     renderMessage = jest.fn((message: Message) => ({}));
     instance = new MessageService(renderMessage);
-    instance.setIsRunning(true);
   });
 
   it("should handle Axios errors", () => {

next/src/components/console/ChatMessage.tsx

@@ -63,7 +63,7 @@ const getMessagePrefix = (message: Message) => {
   } else if (getTaskStatus(message) === TASK_STATUS_COMPLETED) {
     return `Executing: ${message.value}`;
   } else if (getTaskStatus(message) === TASK_STATUS_FINAL) {
-    return `Finished: ${message.value}`;
+    return `Finished:`;
   }
   return "";
 };

next/src/pages/index.tsx

@@ -40,7 +40,7 @@ const Home: NextPage = () => {
 
   const setAgent = useAgentStore.use.setAgent();
   const isAgentStopped = useAgentStore.use.isAgentStopped();
-  const updateIsAgentStopped = useAgentStore.use.updateIsAgentStopped();
+  const setIsAgentStopped = useAgentStore.use.setIsAgentStopped();
 
   const agent = useAgentStore.use.agent();
 
@@ -96,7 +96,7 @@ const Home: NextPage = () => {
     resetAllMessageSlices();
     resetAllTaskSlices();
     newAgent?.run().then(console.log).catch(console.error);
-    updateIsAgentStopped();
+    setIsAgentStopped(false);
   };
 
   const handleKeyPress = (
@@ -110,7 +110,7 @@ const Home: NextPage = () => {
 
   const handleStopAgent = () => {
     agent?.manuallyStopAgent();
-    updateIsAgentStopped();
+    setIsAgentStopped(true);
   };
 
   const handleVisibleWindowClick = (visibleWindow: "Chat" | "Tasks") => {

next/src/services/agent/agent-api.ts

@@ -10,12 +10,10 @@ type ApiProps = Pick<RequestBody, "model_settings" | "goal"> & {
 
 export class AgentApi {
   readonly props: ApiProps;
-  readonly onError: (e: unknown) => never;
   runId: string | undefined;
 
-  constructor(apiProps: ApiProps, onError: (e: unknown) => never) {
+  constructor(apiProps: ApiProps) {
     this.props = apiProps;
-    this.onError = onError;
   }
 
   async getInitialTasks(): Promise<string[]> {
@@ -51,7 +49,6 @@ export class AgentApi {
     url: string,
     data: Omit<RequestBody, "goal" | "model_settings" | "run_id">
   ) {
-    useAgentStore.getState().setIsAgentThinking(true);
     const requestBody: RequestBody = {
       model_settings: this.props.model_settings,
       goal: this.props.goal,
@@ -60,6 +57,7 @@ export class AgentApi {
     };
 
     try {
+      useAgentStore.getState().setIsAgentThinking(true);
       const { run_id, ...data } = await apiUtils.post<T & { run_id: string }>(
         url,
         requestBody,
@@ -68,10 +66,22 @@ export class AgentApi {
 
       if (this.runId === undefined) this.runId = run_id;
       return data;
-    } catch (e) {
-      this.onError(e);
     } finally {
       useAgentStore.getState().setIsAgentThinking(false);
     }
   }
 }
+
+export async function withRetries(
+  fn: () => Promise<void>,
+  onError: (error: unknown) => Promise<boolean>, // Function to handle the error and return whether to continue
+  retries = 3
+): Promise<void> {
+  for (let i = 1; i < retries + 1; i++) {
+    try {
+      return await fn();
+    } catch (error) {
+      if ((await onError(error)) || i === retries) return;
+    }
+  }
+}

next/src/services/agent/agent-work/agent-work.ts

@@ -0,0 +1,6 @@
+export default interface AgentWork {
+  run: () => Promise<void>;
+  conclude: () => Promise<void>;
+  next: () => AgentWork | undefined;
+  onError: (e: unknown) => boolean;
+}

next/src/services/agent/agent-work/analyze-task-work.ts

@@ -0,0 +1,36 @@
+import type { Analysis } from "../analysis";
+import type { Task } from "../../../types/task";
+import type AutonomousAgent from "../autonomous-agent";
+import type AgentWork from "./agent-work";
+import ExecuteTaskWork from "./execute-task-work";
+
+export default class AnalyzeTaskWork implements AgentWork {
+  analysis: Analysis | undefined = undefined;
+
+  constructor(private parent: AutonomousAgent, private task: Task) {}
+
+  run = async () => {
+    this.parent.messageService.startTaskMessage(this.task);
+    this.task = this.parent.model.updateTaskStatus(this.task, "executing");
+    this.analysis = await this.parent.$api.analyzeTask(this.task.value);
+  };
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  conclude = async () => {
+    if (this.analysis) {
+      this.parent.messageService.sendAnalysisMessage(this.analysis);
+    } else {
+      this.parent.messageService.skipTaskMessage(this.task);
+    }
+  };
+
+  next = () => {
+    if (!this.analysis) return undefined;
+    return new ExecuteTaskWork(this.parent, this.task, this.analysis);
+  };
+
+  onError = (e: unknown): boolean => {
+    this.parent.messageService.sendErrorMessage(e);
+    return false;
+  };
+}

next/src/services/agent/agent-work/create-task-work.ts

@@ -0,0 +1,31 @@
+import type { Task } from "../../../types/task";
+import type AgentWork from "./agent-work";
+import type AutonomousAgent from "../autonomous-agent";
+
+export default class CreateTaskWork implements AgentWork {
+  taskValues: string[] = [];
+
+  constructor(private parent: AutonomousAgent, private task: Task, private result: string) {}
+
+  run = async () => {
+    this.taskValues = await this.parent.$api.getAdditionalTasks(
+      {
+        current: this.task.value,
+        remaining: this.parent.model.getRemainingTasks().map((task) => task.value),
+        completed: this.parent.model.getCompletedTasks(),
+      },
+      this.result
+    );
+  };
+
+  conclude = async () => {
+    const TIMEOUT_LONG = 1000;
+    await this.parent.createTasks(this.taskValues);
+    await new Promise((r) => setTimeout(r, TIMEOUT_LONG));
+  };
+
+  next = () => undefined;
+
+  // Ignore errors and simply avoid creating more tasks
+  onError = (): boolean => true;
+}

next/src/services/agent/agent-work/execute-task-work.ts

@@ -0,0 +1,60 @@
+import type { Task } from "../../../types/task";
+import type { Analysis } from "../analysis";
+import type { Message } from "../../../types/message";
+import { v1 } from "uuid";
+import { streamText } from "../../stream-utils";
+import { toApiModelSettings } from "../../../utils/interfaces";
+import type AgentWork from "./agent-work";
+import type AutonomousAgent from "../autonomous-agent";
+import CreateTaskWork from "./create-task-work";
+
+export default class ExecuteTaskWork implements AgentWork {
+  result = "";
+
+  constructor(private parent: AutonomousAgent, private task: Task, private analysis: Analysis) {}
+
+  run = async () => {
+    const executionMessage: Message = {
+      ...this.task,
+      id: v1(),
+      status: "completed",
+      info: "Loading...",
+    };
+    this.parent.messageService.sendMessage({ ...executionMessage, status: "completed" });
+
+    // TODO: this should be moved to the api layer
+    await streamText(
+      "/api/agent/execute",
+      {
+        run_id: this.parent.$api.runId,
+        goal: this.parent.model.getGoal(),
+        task: this.task.value,
+        analysis: this.analysis,
+        model_settings: toApiModelSettings(this.parent.modelSettings),
+      },
+      this.parent.$api.props.session?.accessToken || "",
+      () => {
+        executionMessage.info = "";
+      },
+      (text) => {
+        executionMessage.info += text;
+        this.parent.messageService.updateMessage(executionMessage);
+      },
+      () => !this.parent.isRunning
+    );
+    this.result = executionMessage.info || "";
+  };
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  conclude = async () => {
+    this.parent.model.updateTaskStatus(this.task, "completed");
+    this.parent.messageService.sendMessage({ ...this.task, status: "final" });
+  };
+
+  next = () => (this.result ? new CreateTaskWork(this.parent, this.task, "") : undefined);
+
+  onError = (e: unknown): boolean => {
+    this.parent.messageService.sendErrorMessage(e);
+    return false;
+  };
+}

next/src/services/agent/agent-work/start-task-work.ts

@@ -0,0 +1,24 @@
+import type AutonomousAgent from "../autonomous-agent";
+import type AgentWork from "./agent-work";
+
+export default class StartGoalWork implements AgentWork {
+  tasksValues: string[] = [];
+
+  constructor(private parent: AutonomousAgent) {}
+
+  run = async () => {
+    this.parent.messageService.sendGoalMessage(this.parent.model.getGoal());
+    this.tasksValues = await this.parent.$api.getInitialTasks();
+  };
+
+  conclude = async () => {
+    await this.parent.createTasks(this.tasksValues);
+  };
+
+  onError = (e: unknown): boolean => {
+    this.parent.messageService.sendErrorMessage(e);
+    return false;
+  };
+
+  next = () => undefined;
+}