@zhushanwen/pi-structured-output 0.2.2 → 0.3.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (2) hide show
  1. package/package.json +1 -1
  2. package/src/index.ts +71 -127
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@zhushanwen/pi-structured-output",
3
- "version": "0.2.2",
3
+ "version": "0.3.1",
4
4
  "description": "Structured output tool for Pi — enforces JSON Schema via tool call mechanism with Ajv validation",
5
5
  "type": "module",
6
6
  "main": "index.ts",
package/src/index.ts CHANGED
@@ -1,21 +1,11 @@
1
1
  /**
2
- * Structured Output Extension
2
+ * Structured Output Extension — 全局可用 tool
3
3
  *
4
- * Detects STRUCTURED_OUTPUT_SCHEMA env var on session start, registers a tool
5
- * with Ajv-compiled validation, injects system prompt, and enforces tool usage
6
- * via turn_end + sendUserMessage.
4
+ * 始终注册 `structured-output` tool，AI 可在任何场景下调用。
5
+ * 调用时传入 schema + data，扩展用 Ajv 验证 data 是否符合 schema。
7
6
  *
8
- * Design: FR-1 to FR-5 from spec, FR-4 dual-layer enforcement.
9
- * Reference: Claude Code's SyntheticOutputTool (Ajv + Stop hook).
10
- *
11
- * Key design decisions (borrowed from Claude Code):
12
- * - Schema is injected into BOTH system prompt AND tool description, so LLM
13
- * knows the exact output structure regardless of which signal it reads.
14
- * - Enforcement checks "last call succeeded" (not "was called"), so validation
15
- * failures trigger retries rather than being silently skipped.
16
- * - Retry cap prevents infinite enforcement loops on persistent schema mismatch.
17
- * - WeakMap caches Ajv compile results for repeated calls with the same schema
18
- * object reference (mirrors Claude Code's toolCache pattern).
7
+ * workflow 场景：agent-pool 通过 prompt 指示 AI 调用此 tool 并传入 schema。
8
+ * 普通对话场景：AI 需要返回结构化数据时自行调用。
19
9
  */
20
10
 
21
11
  import Ajv, { type ValidateFunction } from "ajv";
@@ -25,127 +15,81 @@ import { Type } from "@sinclair/typebox";
25
15
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
26
16
  type PiAPI = any;
27
17
 
28
- const ENV_KEY = "STRUCTURED_OUTPUT_SCHEMA";
29
18
  const TOOL_NAME = "structured-output";
30
- const MAX_RETRIES = parseInt(process.env.MAX_STRUCTURED_OUTPUT_RETRIES || "5", 10);
31
-
32
- const ENFORCEMENT_MESSAGE = "你必须调用 structured-output tool 来返回结果。";
33
19
 
34
20
  // ── Ajv WeakMap cache ─────────────────────────────────────────
35
- // Workflow scripts may call agent({schema}) 30-80 times per run.
36
- // Without caching, each call does new Ajv() + validateSchema() + compile().
37
- // WeakMap keyed by schema object reference brings this to near-zero overhead.
21
+ // Repeated calls with the same schema object reference are cached.
38
22
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
39
23
  const ajvCache = new WeakMap<object, ValidateFunction>();
40
24
 
41
25
  function getOrCompileValidator(schema: Record<string, unknown>): ValidateFunction {
42
- const cached = ajvCache.get(schema);
43
- if (cached) return cached;
26
+ const cached = ajvCache.get(schema);
27
+ if (cached) return cached;
44
28
 
45
- const ajv = new Ajv({ strict: false });
46
- const validate = ajv.compile(schema);
47
- ajvCache.set(schema, validate);
48
- return validate;
29
+ const ajv = new Ajv({ strict: false });
30
+ const validate = ajv.compile(schema);
31
+ ajvCache.set(schema, validate);
32
+ return validate;
49
33
  }
50
34
 
51
35
  export default function structuredOutputExtension(pi: PiAPI): void {
52
- const schemaStr = process.env[ENV_KEY];
53
- if (!schemaStr) return;
54
-
55
- // Parse schema
56
- let schema: Record<string, unknown>;
57
- try {
58
- schema = JSON.parse(schemaStr);
59
- } catch {
60
- console.error(`[${TOOL_NAME}] Failed to parse ${ENV_KEY}`);
61
- return;
62
- }
63
-
64
- // Compile with Ajv (cached)
65
- let validate: ValidateFunction;
66
- try {
67
- validate = getOrCompileValidator(schema);
68
- } catch (e) {
69
- console.error(`[${TOOL_NAME}] Invalid JSON Schema:`, (e as Error).message);
70
- return;
71
- }
72
-
73
- // Build prompts with schema embedded
74
- const schemaJsonStr = JSON.stringify(schema, null, 2);
75
- const systemPrompt =
76
- "你必须在完成分析后调用 structured-output tool 来返回结构化结果。" +
77
- "不要在文本回复中输出 JSON，直接调用 structured-output tool。" +
78
- "这是你返回最终结果的唯一方式。\n\n" +
79
- "输出必须严格符合以下 JSON Schema:\n" + schemaJsonStr;
80
-
81
- const toolDescription =
82
- "Return structured output conforming to the JSON Schema. " +
83
- "You MUST call this tool exactly once to return your final result.\n\n" +
84
- "The output must conform to this JSON Schema:\n" + schemaJsonStr;
85
-
86
- // Track structured-output call state per turn
87
- let turnCallCount = 0;
88
- let lastCallSucceeded = false;
89
-
90
- // Register tool
91
- pi.registerTool({
92
- name: TOOL_NAME,
93
- label: "Structured Output",
94
- description: toolDescription,
95
- promptSnippet: "Call structured-output with your final structured answer",
96
- promptGuidelines: [
97
- "You MUST call structured-output as your final action.",
98
- "Do not output JSON in your text response — use this tool instead.",
99
- ],
100
- parameters: Type.Record(Type.String(), Type.Any()),
101
- async execute(_toolCallId: string, params: Record<string, unknown>) {
102
- const valid = validate(params);
103
- if (!valid) {
104
- const errors = validate.errors
105
- ?.map((err) => `${err.instancePath} ${err.message}`)
106
- .join("; ");
107
- throw new Error(`Schema validation failed: ${errors}`);
108
- }
109
- // Mark success for enforcement check
110
- lastCallSucceeded = true;
111
- return {
112
- content: [
113
- { type: "text" as const, text: "Structured output recorded successfully." },
114
- ],
115
- details: params,
116
- };
117
- },
118
- });
119
-
120
- // System prompt injection with embedded schema
121
- pi.on("before_agent_start", async (_event: unknown, ctx: { addSystemInstruction: (s: string) => void }) => {
122
- ctx.addSystemInstruction(systemPrompt);
123
- });
124
-
125
- // Track calls — count for retry cap, success for enforcement
126
- pi.on("tool_execution_start", async (event: { toolName: string }) => {
127
- if (event.toolName === TOOL_NAME) {
128
- turnCallCount++;
129
- }
130
- });
131
-
132
- pi.on("turn_end", async () => {
133
- if (lastCallSucceeded) {
134
- // Reset for next potential turn (shouldn't happen, but defensive)
135
- lastCallSucceeded = false;
136
- turnCallCount = 0;
137
- return;
138
- }
139
-
140
- // Retry cap: stop enforcement after MAX_RETRIES attempts
141
- if (turnCallCount >= MAX_RETRIES) {
142
- console.error(
143
- `[${TOOL_NAME}] Max retries (${MAX_RETRIES}) reached without valid output. Giving up enforcement.`,
144
- );
145
- turnCallCount = 0;
146
- return;
147
- }
148
-
149
- pi.sendUserMessage(ENFORCEMENT_MESSAGE);
150
- });
36
+ pi.registerTool({
37
+ name: TOOL_NAME,
38
+ label: "Structured Output",
39
+ description:
40
+ "Return structured output validated against a JSON Schema. "
41
+ + "Call this tool to produce validated JSON data. "
42
+ + "Pass `schema` (a JSON Schema draft-07 object) and `data` (the value to validate).\n\n"
43
+ + "✅ Correct: schema={type:'object',properties:{name:{type:'string'},age:{type:'number'}},required:['name']}, data={name:'Alice',age:30}\n"
44
+ + "✅ Correct: schema={type:'array',items:{type:'string'}}, data=['a','b','c']\n"
45
+ + "✅ Correct: schema={type:'string',enum:['low','medium','high']}, data='medium'\n\n"
46
+ + "❌ Wrong: putting the answer in text instead of calling this tool\n"
47
+ + "❌ Wrong: data not matching schema (e.g. schema requires number but data is string)\n"
48
+ + "❌ Wrong: schema={type:'object'} with data='hello' (string ≠ object)",
49
+ promptSnippet:
50
+ "Use structured-output to return validated JSON data. "
51
+ + "Pass schema (JSON Schema draft-07) and data (your output). "
52
+ + "Example: {schema:{type:'object',properties:{score:{type:'number'}},required:['score']}, data:{score:8}}",
53
+ promptGuidelines: [
54
+ "schema must be a valid JSON Schema (draft-07). data must conform to it.",
55
+ "Both primitive types (string, number, boolean) and complex types (object, array) are valid schema roots.",
56
+ "Do not output JSON in text — call this tool instead.",
57
+ ],
58
+ parameters: Type.Object({
59
+ schema: Type.Any({
60
+ description: "JSON Schema draft-07 object. Example: {type:'object',properties:{name:{type:'string'}},required:['name']}",
61
+ }),
62
+ data: Type.Any({
63
+ description: "The value to validate against schema. Example: {name:'Alice'}",
64
+ }),
65
+ }),
66
+ async execute(
67
+ _toolCallId: string,
68
+ params: { schema: Record<string, unknown>; data: unknown },
69
+ ) {
70
+ const { schema, data } = params;
71
+
72
+ let validate: ValidateFunction;
73
+ try {
74
+ validate = getOrCompileValidator(schema);
75
+ } catch (e) {
76
+ throw new Error(`Invalid JSON Schema: ${(e as Error).message}`);
77
+ }
78
+
79
+ const valid = validate(data);
80
+ if (!valid) {
81
+ const errors = validate.errors
82
+ ?.map((err) => `${err.instancePath} ${err.message}`)
83
+ .join("; ");
84
+ throw new Error(`Schema validation failed: ${errors}`);
85
+ }
86
+
87
+ return {
88
+ content: [
89
+ { type: "text" as const, text: "Structured output recorded successfully." },
90
+ ],
91
+ details: data as Record<string, unknown>,
92
+ };
93
+ },
94
+ });
151
95
  }