@fleetagent/pi-coding-agent 0.0.2 → 0.0.4

package/dist/core/agent-session.js

@@ -14,7 +14,7 @@
  */
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { basename, dirname } from "node:path";
-import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, modelsAreEqual, resetApiProviders, streamSimple, } from "@fleetagent/pi-ai";
+import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, modelsAreEqual, resetApiProviders, streamSimple, validateToolArguments, } from "@fleetagent/pi-ai";
 import { theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
 import { resolvePath } from "../utils/paths.js";
@@ -27,6 +27,7 @@ import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
+import { STRUCTURED_RESPONSE_INTERNAL_CUSTOM_TYPE } from "./messages.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { createSyntheticSourceInfo } from "./source-info.js";
@@ -54,6 +55,41 @@ export function parseSkillBlock(text) {
 // ============================================================================
 /** Standard thinking levels */
 const THINKING_LEVELS = ["off", "minimal", "low", "medium", "high"];
+const DEFAULT_STRUCTURED_RESPONSE_TOOL_NAME = "structured_output";
+const DEFAULT_STRUCTURED_RESPONSE_CORRECTIONS = 2;
+function isRecord(value) {
+    return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+function getAssistantText(message) {
+    return message.content
+        .filter((block) => block.type === "text")
+        .map((block) => block.text)
+        .join("\n");
+}
+function extractJsonCandidates(text) {
+    const candidates = [];
+    const trimmed = text.trim();
+    if (trimmed) {
+        try {
+            candidates.push(JSON.parse(trimmed));
+        }
+        catch {
+            // Try fenced JSON blocks below.
+        }
+    }
+    const fencePattern = /```(?:json)?\s*([\s\S]*?)```/gi;
+    let match = fencePattern.exec(text);
+    while (match) {
+        try {
+            candidates.push(JSON.parse(match[1].trim()));
+        }
+        catch {
+            // Ignore invalid fenced blocks.
+        }
+        match = fencePattern.exec(text);
+    }
+    return candidates;
+}
 // ============================================================================
 // AgentSession Class
 // ============================================================================
@@ -674,6 +710,186 @@ export class AgentSession {
         }
         return await this._checkCompaction(msg);
     }
+    async getStructuredResponse(options) {
+        if (this.isStreaming) {
+            throw new Error("Agent is already processing. Wait for completion before requesting structured output.");
+        }
+        if (!this.model) {
+            throw new Error(formatNoModelSelectedMessage());
+        }
+        if (!this._modelRegistry.hasConfiguredAuth(this.model)) {
+            throw new Error(formatNoApiKeyFoundMessage(this.model.provider));
+        }
+        const schemaName = options.name ?? DEFAULT_STRUCTURED_RESPONSE_TOOL_NAME;
+        const tool = {
+            name: schemaName,
+            description: options.description ??
+                "Return the requested structured response. Call this tool exactly once with arguments matching the schema.",
+            parameters: options.schema,
+        };
+        const lastAssistant = this._findLastAssistantMessage();
+        if (!lastAssistant) {
+            throw new Error("No assistant response is available to structure.");
+        }
+        const direct = this._tryParseStructuredAssistantText(tool, lastAssistant);
+        if (direct.ok) {
+            this._appendStructuredInternalEntry("result", schemaName, 0, "Validated structured response from assistant JSON.", {
+                stage: "result",
+                schemaName,
+                attempt: 0,
+                source: "json",
+            });
+            return { output: direct.output, attempts: 0, source: "json", message: lastAssistant };
+        }
+        const maxCorrections = options.maxCorrections ?? DEFAULT_STRUCTURED_RESPONSE_CORRECTIONS;
+        const maxAttempts = maxCorrections + 1;
+        const { apiKey, headers } = await this._getRequiredRequestAuth(this.model);
+        const messages = this._buildStructuredResponseMessages(options.scope ?? "latest", lastAssistant, schemaName);
+        let lastError = direct.error;
+        for (let attempt = 1; attempt <= maxAttempts; attempt++) {
+            const requestText = attempt === 1
+                ? `Extract the structured response by calling ${schemaName} exactly once.`
+                : `Correct the structured response by calling ${schemaName} exactly once.\n\nValidation error:\n${lastError}`;
+            const userMessage = {
+                role: "user",
+                content: [{ type: "text", text: requestText }],
+                timestamp: Date.now(),
+            };
+            messages.push(userMessage);
+            this._appendStructuredInternalEntry("request", schemaName, attempt, requestText, {
+                stage: "request",
+                schemaName,
+                attempt,
+                validationError: attempt === 1 ? undefined : lastError,
+            });
+            const responseStream = await this.agent.streamFn(this.model, {
+                systemPrompt: "You are a structured data extraction assistant. Do not answer in prose. Use the provided tool to return the structured data.",
+                messages,
+                tools: [tool],
+            }, {
+                apiKey,
+                headers,
+                sessionId: this.agent.sessionId,
+                transport: this.agent.transport,
+                thinkingBudgets: this.agent.thinkingBudgets,
+                maxRetryDelayMs: this.agent.maxRetryDelayMs,
+            });
+            const assistantMessage = await responseStream.result();
+            messages.push(assistantMessage);
+            this._appendStructuredInternalEntry("assistant", schemaName, attempt, this._formatStructuredAssistantLog(assistantMessage), {
+                stage: "assistant",
+                schemaName,
+                attempt,
+            });
+            const toolCall = assistantMessage.content.find((block) => block.type === "toolCall" && block.name === schemaName);
+            if (toolCall) {
+                const validation = this._validateStructuredArguments(tool, toolCall.arguments);
+                if (validation.ok) {
+                    this._appendStructuredInternalEntry("result", schemaName, attempt, "Validated structured tool response.", {
+                        stage: "result",
+                        schemaName,
+                        attempt,
+                        source: "tool",
+                    });
+                    return { output: validation.output, attempts: attempt, source: "tool", message: assistantMessage };
+                }
+                lastError = validation.error;
+                const toolResult = {
+                    role: "toolResult",
+                    toolCallId: toolCall.id,
+                    toolName: toolCall.name,
+                    content: [{ type: "text", text: validation.error }],
+                    isError: true,
+                    timestamp: Date.now(),
+                };
+                messages.push(toolResult);
+                this._appendStructuredInternalEntry("tool_result", schemaName, attempt, validation.error, {
+                    stage: "tool_result",
+                    schemaName,
+                    attempt,
+                    validationError: validation.error,
+                });
+                continue;
+            }
+            const textValidation = this._tryParseStructuredAssistantText(tool, assistantMessage);
+            if (textValidation.ok) {
+                this._appendStructuredInternalEntry("result", schemaName, attempt, "Validated structured JSON response.", {
+                    stage: "result",
+                    schemaName,
+                    attempt,
+                    source: "json",
+                });
+                return { output: textValidation.output, attempts: attempt, source: "json", message: assistantMessage };
+            }
+            lastError = textValidation.error;
+        }
+        throw new Error(`Structured response validation failed after ${maxAttempts} attempt(s):\n${lastError}`);
+    }
+    _buildStructuredResponseMessages(scope, lastAssistant, schemaName) {
+        if (scope === "conversation") {
+            return this.agent.state.messages.filter((message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult");
+        }
+        return [
+            {
+                role: "user",
+                content: [
+                    {
+                        type: "text",
+                        text: `Extract structured data from the latest assistant response below. Call ${schemaName} exactly once.\n\n` +
+                            `<assistant_response>\n${getAssistantText(lastAssistant)}\n</assistant_response>`,
+                    },
+                ],
+                timestamp: Date.now(),
+            },
+        ];
+    }
+    _tryParseStructuredAssistantText(tool, message) {
+        const text = getAssistantText(message);
+        for (const candidate of extractJsonCandidates(text)) {
+            if (!isRecord(candidate)) {
+                continue;
+            }
+            const validation = this._validateStructuredArguments(tool, candidate);
+            if (validation.ok) {
+                return validation;
+            }
+        }
+        return { ok: false, error: "Assistant response did not contain valid JSON matching the requested schema." };
+    }
+    _validateStructuredArguments(tool, arguments_) {
+        try {
+            const output = validateToolArguments(tool, {
+                type: "toolCall",
+                id: "structured-response-validation",
+                name: tool.name,
+                arguments: arguments_,
+            });
+            return { ok: true, output };
+        }
+        catch (error) {
+            return { ok: false, error: error instanceof Error ? error.message : String(error) };
+        }
+    }
+    _appendStructuredInternalEntry(stage, schemaName, attempt, content, details) {
+        this.session.appendCustomMessageEntry(STRUCTURED_RESPONSE_INTERNAL_CUSTOM_TYPE, content, false, {
+            ...details,
+            stage,
+            schemaName,
+            attempt,
+        });
+    }
+    _formatStructuredAssistantLog(message) {
+        const parts = message.content.map((block) => {
+            if (block.type === "text") {
+                return block.text;
+            }
+            if (block.type === "thinking") {
+                return "[thinking omitted]";
+            }
+            return `tool:${block.name} ${JSON.stringify(block.arguments)}`;
+        });
+        return parts.join("\n");
+    }
     /**
      * Send a prompt to the agent.
      * - Handles extension commands (registered via pi.registerCommand) immediately, even during streaming