Skip to content

Commit 6f6fd3f

Browse files
yaonyanyaonyan
committed
refactor: convert sampling from option to independent mode plugins
- Add new execution modes: agentic_sampling and agentic_workflow_sampling - Create dedicated mode plugins for sampling (mode-agentic-sampling-plugin.ts, mode-workflow-sampling-plugin.ts) - Create separate registrar functions for sampling modes - Remove sampling parameter from existing mode plugins and registrars - Update ComposeDefinition to use samplingConfig instead of sampling option - Simplify mode-specific logic by eliminating conditional sampling checks - Update examples to use new mode naming convention This refactor improves maintainability by: - Separating concerns: each mode has its own plugin - Eliminating confusing option/mode mixing - Making execution modes explicit and type-safe - Simplifying codebase by removing conditional branching
1 parent 7ae11cc commit 6f6fd3f

16 files changed

+285
-101
lines changed

‎packages/core/examples/sampling/01-basic-composition.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,7 @@ export const toolDefinitions: ComposeDefinition[] = [
1818
{
1919
name: "file-organizer",
2020
options: {
21-
sampling: true,
22-
mode: "agentic_workflow",
21+
mode: "agentic_sampling",
2322
},
2423
description:
2524
`I am a smart file organizer that helps users manage their files efficiently.

‎packages/core/src/executors/agentic/agentic-sampling-registrar.ts‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
import { jsonSchema } from "../../utils/schema.ts";
2+
import type { SamplingConfig } from "../../types.ts";
3+
import { createGoogleCompatibleJSONSchema } from "../../utils/common/provider.ts";
4+
import type { ComposableMCPServer } from "../../compose.ts";
5+
import { CompiledPrompts } from "../../prompts/index.ts";
6+
import { SamplingExecutor } from "../sampling/agentic-sampling-executor.ts";
7+
import type { ExternalTool } from "../sampling/base-sampling-executor.ts";
8+
import { createArgsDefFactory } from "../../factories/args-def-factory.ts";
9+
10+
export interface RegisterAgenticSamplingToolParams {
11+
description: string;
12+
name: string;
13+
allToolNames: string[];
14+
depGroups: Record<string, unknown>;
15+
toolNameToDetailList: [string, unknown][];
16+
samplingConfig?: SamplingConfig;
17+
}
18+
19+
export function registerAgenticSamplingTool(
20+
server: ComposableMCPServer,
21+
{
22+
description,
23+
name,
24+
allToolNames,
25+
depGroups,
26+
toolNameToDetailList,
27+
samplingConfig,
28+
}: RegisterAgenticSamplingToolParams,
29+
) {
30+
const createArgsDef = createArgsDefFactory(
31+
name,
32+
allToolNames,
33+
depGroups,
34+
undefined,
35+
undefined,
36+
);
37+
38+
// Create sampling executor
39+
const samplingExecutor = new SamplingExecutor(
40+
name,
41+
description,
42+
allToolNames,
43+
toolNameToDetailList as [string, ExternalTool][],
44+
server,
45+
samplingConfig,
46+
);
47+
48+
// Use sampling-specific prompt
49+
description = CompiledPrompts.samplingExecution({
50+
toolName: name,
51+
description,
52+
toolList: allToolNames.map((name) => `- ${name}`).join("\n"),
53+
});
54+
55+
// Use sampling-specific args definition
56+
const argsDef = createArgsDef.forSampling();
57+
const schema = allToolNames.length > 0
58+
? argsDef
59+
: { type: "object", properties: {} };
60+
61+
server.tool(
62+
name,
63+
description,
64+
jsonSchema<Record<string, unknown>>(
65+
createGoogleCompatibleJSONSchema(schema as Record<string, unknown>),
66+
),
67+
async (args: Record<string, unknown>) => {
68+
return await samplingExecutor.executeSampling(
69+
args,
70+
schema as Record<string, unknown>,
71+
);
72+
},
73+
);
74+
}

‎packages/core/src/executors/agentic/agentic-tool-registrar.ts‎

Lines changed: 9 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@ import { createGoogleCompatibleJSONSchema } from "../../utils/common/provider.ts
44
import type { ComposableMCPServer } from "../../compose.ts";
55
import { CompiledPrompts } from "../../prompts/index.ts";
66
import { AgenticExecutor } from "./agentic-executor.ts";
7-
import { SamplingExecutor } from "../sampling/agentic-sampling-executor.ts";
8-
import type { ExternalTool } from "../sampling/base-sampling-executor.ts";
97
import { createArgsDefFactory } from "../../factories/args-def-factory.ts";
108

119
export function registerAgenticTool(
@@ -16,7 +14,6 @@ export function registerAgenticTool(
1614
allToolNames,
1715
depGroups,
1816
toolNameToDetailList,
19-
sampling = false,
2017
}: RegisterToolParams,
2118
) {
2219
const createArgsDef = createArgsDefFactory(
@@ -27,44 +24,25 @@ export function registerAgenticTool(
2724
undefined,
2825
);
2926

30-
// Determine if sampling mode is enabled and extract config
31-
const isSamplingMode = sampling === true || typeof sampling === "object";
32-
const samplingConfig = typeof sampling === "object" ? sampling : undefined;
33-
34-
// Create executors
27+
// Create executor
3528
const agenticExecutor = new AgenticExecutor(
3629
name,
3730
allToolNames,
3831
toolNameToDetailList,
3932
server,
4033
);
4134

42-
const samplingExecutor = new SamplingExecutor(
43-
name,
35+
description = CompiledPrompts.autonomousExecution({
36+
toolName: name,
4437
description,
45-
allToolNames,
46-
toolNameToDetailList as [string, ExternalTool][],
47-
server,
48-
samplingConfig,
49-
);
50-
51-
description = isSamplingMode
52-
? CompiledPrompts.samplingExecution({
53-
toolName: name,
54-
description,
55-
toolList: allToolNames.map((name) => `- ${name}`).join("\n"),
56-
})
57-
: CompiledPrompts.autonomousExecution({
58-
toolName: name,
59-
description,
60-
});
38+
});
6139

6240
const agenticArgsDef = createArgsDef.forAgentic(
6341
toolNameToDetailList,
6442
false, // not sampling mode
6543
);
6644
const argsDef: Schema<Record<PropertyKey, never>>["jsonSchema"] =
67-
isSamplingMode ? createArgsDef.forSampling() : agenticArgsDef;
45+
agenticArgsDef;
6846
const schema = allToolNames.length > 0
6947
? argsDef
7048
: { type: "object", properties: {} };
@@ -76,18 +54,10 @@ export function registerAgenticTool(
7654
createGoogleCompatibleJSONSchema(schema as Record<string, unknown>),
7755
),
7856
async (args: Record<string, unknown>) => {
79-
// Use appropriate executor based on mode
80-
if (isSamplingMode) {
81-
return await samplingExecutor.executeSampling(
82-
args,
83-
schema as Record<string, unknown>,
84-
);
85-
} else {
86-
return await agenticExecutor.execute(
87-
args,
88-
schema as Record<string, unknown>,
89-
);
90-
}
57+
return await agenticExecutor.execute(
58+
args,
59+
schema as Record<string, unknown>,
60+
);
9161
},
9262
);
9363
}

‎packages/core/src/executors/workflow/workflow-sampling-registrar.ts‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import { jsonSchema } from "../../utils/schema.ts";
2+
import type { SamplingConfig } from "../../types.ts";
3+
import type { MCPCStep } from "../../utils/state.ts";
4+
import { WorkflowState } from "../../utils/state.ts";
5+
import { createGoogleCompatibleJSONSchema } from "../../utils/common/provider.ts";
6+
import type { ComposableMCPServer } from "../../compose.ts";
7+
import { CompiledPrompts } from "../../prompts/index.ts";
8+
import { createArgsDefFactory } from "../../factories/args-def-factory.ts";
9+
import { WorkflowSamplingExecutor } from "../sampling/workflow-sampling-executor.ts";
10+
import type { ExternalTool } from "../sampling/base-sampling-executor.ts";
11+
12+
export interface RegisterWorkflowSamplingToolParams {
13+
description: string;
14+
name: string;
15+
allToolNames: string[];
16+
depGroups: Record<string, unknown>;
17+
toolNameToDetailList: [string, unknown][];
18+
predefinedSteps?: MCPCStep[];
19+
samplingConfig?: SamplingConfig;
20+
ensureStepActions?: string[];
21+
toolNameToIdMapping?: Map<string, string>;
22+
}
23+
24+
export function registerWorkflowSamplingTool(
25+
server: ComposableMCPServer,
26+
{
27+
description,
28+
name,
29+
allToolNames,
30+
depGroups,
31+
toolNameToDetailList,
32+
predefinedSteps,
33+
samplingConfig,
34+
ensureStepActions,
35+
toolNameToIdMapping: _toolNameToIdMapping,
36+
}: RegisterWorkflowSamplingToolParams,
37+
) {
38+
const createArgsDef = createArgsDefFactory(
39+
name,
40+
allToolNames,
41+
depGroups,
42+
predefinedSteps,
43+
ensureStepActions,
44+
);
45+
46+
// Create sampling executor
47+
const workflowSamplingExecutor = new WorkflowSamplingExecutor(
48+
name,
49+
description,
50+
allToolNames,
51+
toolNameToDetailList as [string, ExternalTool][],
52+
createArgsDef,
53+
server,
54+
predefinedSteps,
55+
samplingConfig,
56+
);
57+
58+
const workflowState = new WorkflowState();
59+
60+
// Use sampling-specific prompt
61+
const baseDescription = CompiledPrompts.samplingExecution({
62+
toolName: name,
63+
description,
64+
toolList: allToolNames.map((name) => `- ${name}`).join("\n"),
65+
});
66+
67+
// Use sampling-specific args definition
68+
const argsDef = createArgsDef.forSampling();
69+
70+
const schema = allToolNames.length > 0
71+
? argsDef
72+
: { type: "object", properties: {} };
73+
74+
server.tool(
75+
name,
76+
baseDescription,
77+
jsonSchema<Record<string, unknown>>(
78+
createGoogleCompatibleJSONSchema(schema as Record<string, unknown>),
79+
),
80+
async (args: Record<string, unknown>) => {
81+
try {
82+
return await workflowSamplingExecutor.executeWorkflowSampling(
83+
args,
84+
schema as Record<string, unknown>,
85+
workflowState,
86+
);
87+
} catch (error) {
88+
workflowState.reset();
89+
return {
90+
content: [
91+
{
92+
type: "text",
93+
text: `Workflow execution error: ${(error as Error).message}`,
94+
},
95+
],
96+
isError: true,
97+
};
98+
}
99+
},
100+
);
101+
}

‎packages/core/src/executors/workflow/workflow-tool-registrar.ts‎

Lines changed: 14 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,6 @@ import { WorkflowExecutor } from "./workflow-executor.ts";
66
import type { ComposableMCPServer } from "../../compose.ts";
77
import { CompiledPrompts } from "../../prompts/index.ts";
88
import { createArgsDefFactory } from "../../factories/args-def-factory.ts";
9-
import { WorkflowSamplingExecutor } from "../sampling/workflow-sampling-executor.ts";
10-
import type { ExternalTool } from "../sampling/base-sampling-executor.ts";
119

1210
export function registerAgenticWorkflowTool(
1311
server: ComposableMCPServer,
@@ -18,7 +16,6 @@ export function registerAgenticWorkflowTool(
1816
depGroups,
1917
toolNameToDetailList,
2018
predefinedSteps,
21-
sampling = false,
2219
ensureStepActions,
2320
toolNameToIdMapping,
2421
}: RegisterWorkflowToolParams,
@@ -31,11 +28,7 @@ export function registerAgenticWorkflowTool(
3128
ensureStepActions,
3229
);
3330

34-
// Determine if sampling mode is enabled and extract config
35-
const isSamplingMode = sampling === true || typeof sampling === "object";
36-
const samplingConfig = typeof sampling === "object" ? sampling : undefined;
37-
38-
// Create executors
31+
// Create executor
3932
const workflowExecutor = new WorkflowExecutor(
4033
name,
4134
allToolNames,
@@ -47,44 +40,26 @@ export function registerAgenticWorkflowTool(
4740
toolNameToIdMapping,
4841
);
4942

50-
const workflowSamplingExecutor = new WorkflowSamplingExecutor(
51-
name,
52-
description,
53-
allToolNames,
54-
toolNameToDetailList as [string, ExternalTool][],
55-
createArgsDef,
56-
server,
57-
predefinedSteps,
58-
samplingConfig,
59-
);
60-
6143
const workflowState = new WorkflowState();
6244

6345
const planningInstructions = predefinedSteps
6446
? "- Set `init: true` (steps are predefined)"
6547
: "- Set `init: true` and define complete `steps` array";
6648

67-
// Generate description based on mode
68-
const baseDescription = isSamplingMode
69-
? CompiledPrompts.samplingExecution({
70-
toolName: name,
71-
description,
72-
toolList: allToolNames.map((name) => `- ${name}`).join("\n"),
73-
})
74-
: CompiledPrompts.workflowExecution({
75-
toolName: name,
76-
description: description,
77-
planningInstructions,
78-
});
49+
// Generate description
50+
const baseDescription = CompiledPrompts.workflowExecution({
51+
toolName: name,
52+
description: description,
53+
planningInstructions,
54+
});
7955

80-
// Generate schema based on mode
81-
const argsDef = isSamplingMode
82-
? createArgsDef.forSampling()
83-
: createArgsDef.forTool();
56+
// Generate schema
57+
const argsDef = createArgsDef.forTool();
8458

85-
const toolDescription = isSamplingMode
86-
? baseDescription
87-
: createArgsDef.forToolDescription(baseDescription, workflowState);
59+
const toolDescription = createArgsDef.forToolDescription(
60+
baseDescription,
61+
workflowState,
62+
);
8863

8964
server.tool(
9065
name,
@@ -94,16 +69,7 @@ export function registerAgenticWorkflowTool(
9469
),
9570
async (args: Record<string, unknown>) => {
9671
try {
97-
// Use appropriate executor based on mode
98-
if (isSamplingMode) {
99-
return await workflowSamplingExecutor.executeWorkflowSampling(
100-
args as Record<string, unknown>,
101-
argsDef as Record<string, unknown>,
102-
workflowState,
103-
);
104-
} else {
105-
return await workflowExecutor.execute(args, workflowState);
106-
}
72+
return await workflowExecutor.execute(args, workflowState);
10773
} catch (error) {
10874
workflowState.reset();
10975
return {

‎packages/core/src/plugin-types.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ export interface AgentToolRegistrationContext {
134134
hiddenToolNames: string[];
135135
options: {
136136
mode?: string;
137-
sampling?: boolean | { maxIterations?: number; summarize?: boolean };
137+
samplingConfig?: { maxIterations?: number; summarize?: boolean };
138138
steps?: Array<{ description: string; actions: string[] }>;
139139
ensureStepActions?: string[];
140140
[key: string]: unknown;

0 commit comments

Comments
 (0)