Refactor report tools to include runtime context in parameters for create and edit functions

packages/ai/src/tools/visualization-tools/create-reports-file-tool.ts

@@ -1,3 +1,4 @@
+import type { RuntimeContext } from '@mastra/core/runtime-context';
 import { createTool } from '@mastra/core/tools';
 import { wrapTraced } from 'braintrust';
 import { z } from 'zod';
@@ -19,10 +20,18 @@ interface CreateReportsOutput {
 
 // Main create reports function
 const createReportsFile = wrapTraced(
-  async (params: CreateReportsParams): Promise<CreateReportsOutput> => {
+  async ({
+    params,
+    runtimeContext,
+  }: {
+    params: CreateReportsParams;
+    runtimeContext: RuntimeContext;
+  }): Promise<CreateReportsOutput> => {
     // Dummy implementation - just return success
     const dummyId = `report_${Date.now()}`;
 
+    runtimeContext.set(dummyId, params.code);
+
     return {
       success: true,
       message: `Successfully created report: ${params.name}`,
@@ -53,7 +62,10 @@ export const createReports = createTool({
       code: z.string(),
     }),
   }),
-  execute: async ({ context }) => {
-    return await createReportsFile(context as CreateReportsParams);
+  execute: async ({ context, runtimeContext }) => {
+    return await createReportsFile({
+      params: context as CreateReportsParams,
+      runtimeContext,
+    });
   },
 });

packages/ai/src/tools/visualization-tools/edit-reports-file-tool.ts

@@ -1,3 +1,4 @@
+import type { RuntimeContext } from '@mastra/core/runtime-context';
 import { createTool } from '@mastra/core/tools';
 import { wrapTraced } from 'braintrust';
 import { z } from 'zod';
@@ -25,9 +26,15 @@ interface EditReportsOutput {
 
 // Main edit reports function
 const editReportsFile = wrapTraced(
-  async (params: EditReportsParams): Promise<EditReportsOutput> => {
+  async ({
+    params,
+    runtimeContext,
+  }: {
+    params: EditReportsParams;
+    runtimeContext: RuntimeContext;
+  }): Promise<EditReportsOutput> => {
     // Dummy implementation - simulate applying edits
-    let simulatedCode = '// Existing report code';
+    let simulatedCode = runtimeContext.get(params.id) as string;
 
     // Simulate applying each edit
     for (const edit of params.edits) {
@@ -40,6 +47,8 @@ const editReportsFile = wrapTraced(
       }
     }
 
+    runtimeContext.set(params.id, simulatedCode);
+
     return {
       success: true,
       message: `Successfully edited report: ${params.name}`,
@@ -80,7 +89,10 @@ export const editReports = createTool({
       code: z.string(),
     }),
   }),
-  execute: async ({ context }) => {
-    return await editReportsFile(context as EditReportsParams);
+  execute: async ({ context, runtimeContext }) => {
+    return await editReportsFile({
+      params: context as EditReportsParams,
+      runtimeContext,
+    });
   },
 });