fix: Add job data anonymization before passing for context enrichment (…

…#346)

* Revert "Revert "feat: Retrieve past function results for reinforcement learning (#334)" (#343)"

This reverts commit 5557f13.

* chore: Add job data anonymization in formatJobsContext function

* update

* chore: Remove unused import and fix formatting in jobs.test.ts

* chore: Standardize type casing in test and run files

* fix: Update test expectations for job status XML tag

* fix: Correct case sensitivity in test assertions for job context formatting

control-plane/src/modules/jobs/jobs.test.ts

@@ -2,7 +2,13 @@ import { ulid } from "ulid";
 import { packer } from "../packer";
 import { upsertServiceDefinition } from "../service-definitions";
 import { createOwner } from "../test/util";
-import { createJob, pollJobs, getJob, requestApproval, submitApproval } from "./jobs";
+import {
+  createJob,
+  pollJobs,
+  getJob,
+  requestApproval,
+  submitApproval,
+} from "./jobs";
 import {
   acknowledgeJob,
   persistJobResult,
@@ -412,7 +418,7 @@ describe("pollJobs", () => {
         .fill(0)
         .map(async () => {
           await new Promise((resolve) =>
-            setTimeout(resolve, Math.random() * 10),
+            setTimeout(resolve, Math.random() * 10)
           );
           return poll();
         }),
@@ -443,10 +449,8 @@ describe("submitApproval", () => {
       },
       owner,
     });
-
-  })
+  });
   it("should mark job as approved", async () => {
-
     const result = await createJob({
       targetFn: mockTargetFn,
       targetArgs: mockTargetArgs,
@@ -460,7 +464,7 @@ describe("submitApproval", () => {
     await requestApproval({
       clusterId: owner.clusterId,
       jobId: result.id,
-    })
+    });
 
     const retreivedJob1 = await getJob({
       jobId: result.id,
@@ -473,7 +477,7 @@ describe("submitApproval", () => {
       clusterId: owner.clusterId,
       call: retreivedJob1!,
       approved: true,
-    })
+    });
 
     const retreivedJob2 = await getJob({
       jobId: result.id,
@@ -489,7 +493,7 @@ describe("submitApproval", () => {
       clusterId: owner.clusterId,
       call: retreivedJob1!,
       approved: false,
-    })
+    });
 
     const retreivedJob3 = await getJob({
       jobId: result.id,
@@ -502,7 +506,6 @@ describe("submitApproval", () => {
   });
 
   it("should mark job as denied", async () => {
-
     const result = await createJob({
       targetFn: mockTargetFn,
       targetArgs: mockTargetArgs,
@@ -516,7 +519,7 @@ describe("submitApproval", () => {
     await requestApproval({
       clusterId: owner.clusterId,
       jobId: result.id,
-    })
+    });
 
     const retreivedJob1 = await getJob({
       jobId: result.id,
@@ -529,7 +532,7 @@ describe("submitApproval", () => {
       clusterId: owner.clusterId,
       call: retreivedJob1!,
       approved: false,
-    })
+    });
 
     const retreivedJob2 = await getJob({
       jobId: result.id,
@@ -545,7 +548,7 @@ describe("submitApproval", () => {
       clusterId: owner.clusterId,
       call: retreivedJob1!,
       approved: true,
-    })
+    });
 
     const retreivedJob3 = await getJob({
       jobId: result.id,
@@ -555,6 +558,5 @@ describe("submitApproval", () => {
     expect(retreivedJob3!.approved).toBe(false);
     expect(retreivedJob3!.status).toBe("success");
     expect(retreivedJob3!.resultType).toBe("rejection");
-
   });
 });

control-plane/src/modules/jobs/jobs.ts

@@ -1,4 +1,4 @@
-import { and, eq, gt, isNull, lte, sql } from "drizzle-orm";
+import { and, desc, eq, gt, isNull, lte, sql } from "drizzle-orm";
 import { env } from "../../utilities/env";
 import { JobPollTimeoutError, NotFoundError } from "../../utilities/errors";
 import { getBlobsForJobs } from "../blobs";
@@ -137,6 +137,38 @@ export const getJob = async ({
   };
 };
 
+export const getLatestJobsResultedByFunctionName = async ({
+  clusterId,
+  service,
+  functionName,
+  limit,
+  resultType,
+}: {
+  clusterId: string;
+  service: string;
+  functionName: string;
+  limit: number;
+  resultType: ResultType;
+}) => {
+  return data.db
+    .select({
+      result: data.jobs.result,
+      resultType: data.jobs.result_type,
+      targetArgs: data.jobs.target_args,
+    })
+    .from(data.jobs)
+    .where(
+      and(
+        eq(data.jobs.cluster_id, clusterId),
+        eq(data.jobs.target_fn, functionName),
+        eq(data.jobs.service, service),
+        eq(data.jobs.result_type, resultType),
+      ),
+    )
+    .orderBy(desc(data.jobs.created_at))
+    .limit(limit);
+};
+
 export const getJobsForWorkflow = async ({
   clusterId,
   runId,

control-plane/src/modules/tool-metadata.ts

@@ -75,7 +75,6 @@ export async function getToolMetadata(
   const [result] = await db
     .select({
       additionalContext: toolMetadata.user_defined_context,
-      resultKeys: toolMetadata.result_keys,
     })
     .from(toolMetadata)
     .where(

control-plane/src/modules/workflows/agent/nodes/model-call.ts

@@ -2,10 +2,7 @@ import { ReleventToolLookup } from "../agent";
 import { toAnthropicMessages } from "../../workflow-messages";
 import { logger } from "../../../observability/logger";
 import { WorkflowAgentState, WorkflowAgentStateMessage } from "../state";
-import {
-  addAttributes,
-  withSpan,
-} from "../../../observability/tracer";
+import { addAttributes, withSpan } from "../../../observability/tracer";
 import { AgentError } from "../../../../utilities/errors";
 import { ulid } from "ulid";
 
@@ -74,8 +71,7 @@ const _handleModelCall = async (
     "If there is nothing left to do, return 'done' and provide the final result.",
     "If you encounter invocation errors (e.g., incorrect tool name, missing input), retry based on the error message.",
     "When possible, return multiple invocations to trigger them in parallel.",
-    "Once all tasks have been completed, return the final result as a structured object.",
-    "Provide concise and clear responses. Use **bold** to highlight important words.",
+    "Provide concise and clear responses.",
     state.additionalContext,
     "<TOOLS_SCHEMAS>",
     schemaString,
@@ -263,7 +259,6 @@ const _handleModelCall = async (
     };
   }
 
-
   return {
     messages: [
       {

control-plane/src/modules/workflows/agent/nodes/model-output.ts

@@ -1,4 +1,3 @@
-
 import { JsonSchema7ObjectType } from "zod-to-json-schema";
 import { AgentTool } from "../tool";
 import { workflows } from "../../../data";
@@ -8,8 +7,7 @@ import { WorkflowAgentState } from "../state";
 type ModelInvocationOutput = {
   toolName: string;
   input: unknown;
-
-}
+};
 
 export type ModelOutput = {
   invocations?: ModelInvocationOutput[];
@@ -18,18 +16,17 @@ export type ModelOutput = {
   message?: string;
   done?: boolean;
   issue?: string;
-}
+};
 
 export const buildModelSchema = ({
   state,
   relevantSchemas,
-  resultSchema
+  resultSchema,
 }: {
   state: WorkflowAgentState;
   relevantSchemas: AgentTool[];
   resultSchema?: InferSelectModel<typeof workflows>["result_schema"];
-  }) => {
-
+}) => {
   // Build the toolName enum
   const toolNameEnum = [
     ...relevantSchemas.map((tool) => tool.name),
@@ -48,7 +45,7 @@ export const buildModelSchema = ({
       issue: {
         type: "string",
         description:
-          "Describe any issues you have encountered in this step. Specifically related to the tools you are using.",
+          "Describe any issues you have encountered in this step. Specifically related to the tools you are using. If none, keep this field empty.",
       },
     },
   };

control-plane/src/modules/workflows/agent/run.test.ts

@@ -1,5 +1,5 @@
 import { Run } from "../workflows";
-import { buildMockTools, findRelevantTools } from "./run";
+import { buildMockTools, findRelevantTools, formatJobsContext } from "./run";
 import { upsertServiceDefinition } from "../../service-definitions";
 import { createOwner } from "../../test/util";
 import { ulid } from "ulid";
@@ -51,10 +51,10 @@ describe("findRelevantTools", () => {
     });
 
     expect(tools.map((tool) => tool.name)).toContain(
-      "testService_someFunction",
+      "testService_someFunction"
     );
     expect(tools.map((tool) => tool.name)).not.toContain(
-      "testService_someOtherFunction",
+      "testService_someOtherFunction"
     );
   });
 });
@@ -139,3 +139,65 @@ describe("buildMockTools", () => {
     });
   });
 });
+
+describe("formatJobsContext", () => {
+  it("should return empty string for empty jobs array", () => {
+    const result = formatJobsContext([], "success");
+    expect(result).toBe("");
+  });
+
+  it("should format successful jobs correctly", () => {
+    const jobs = [
+      {
+        targetArgs: JSON.stringify({ param1: "test", param2: 123 }),
+        result: JSON.stringify({ status: "ok", data: "result" }),
+      },
+      {
+        targetArgs: JSON.stringify({ param3: true }),
+        result: JSON.stringify({ status: "ok", count: 42 }),
+      },
+    ];
+
+    const result = formatJobsContext(jobs, "success");
+
+    // Verify structure and anonymization
+    expect(result).toContain('<previous_jobs status="success">');
+    expect(result).toContain("<input>");
+    expect(result).toContain("<output>");
+    expect(result).toContain('"param1":"<string>"');
+    expect(result).toContain('"param2":"<number>"');
+    expect(result).toContain('"param3":"<boolean>"');
+    expect(result).toContain('"status":"<string>"');
+    expect(result).toContain('"count":"<number>"');
+  });
+
+  it("should handle null results", () => {
+    const jobs = [
+      {
+        targetArgs: JSON.stringify({ test: "value" }),
+        result: null,
+      },
+    ];
+
+    const result = formatJobsContext(jobs, "failed");
+
+    expect(result).toContain('<previous_jobs status="failed">');
+    expect(result).toContain("<input>");
+    expect(result).toContain("<output>");
+    expect(result).toContain('"test":"<string>"');
+  });
+
+  it("should anonymize arrays", () => {
+    const result = formatJobsContext(
+      [
+        {
+          targetArgs: JSON.stringify([1, 2, 3]),
+          result: JSON.stringify([4, 5, 6]),
+        },
+      ],
+      "success"
+    );
+    expect(result).toContain(`<input>[\"<number>\"]</input>`);
+    expect(result).toContain(`<output>[\"<number>\"]</output>`);
+  });
+});