@iqai/adk 0.1.20 → 0.1.21

package/dist/index.mjs

@@ -1353,7 +1353,7 @@ var BaseLlm = class {
    * @param llmRequest LlmRequest, the request to send to the LLM.
    * @returns BaseLLMConnection, the connection to the LLM.
    */
-  connect(llmRequest) {
+  connect(_llmRequest) {
     throw new Error(`Live connection is not supported for ${this.model}.`);
   }
 };
@@ -8088,6 +8088,9 @@ var IdentityLlmRequestProcessor = class extends BaseLlmRequestProcessor {
 };
 var requestProcessor5 = new IdentityLlmRequestProcessor();
 
+// src/flows/llm-flows/instructions.ts
+import { zodToJsonSchema as zodToJsonSchema2 } from "zod-to-json-schema";
+
 // src/utils/instructions-utils.ts
 async function injectSessionState(template, readonlyContext) {
   const invocationContext = readonlyContext._invocationContext;
@@ -8202,6 +8205,22 @@ var InstructionsLlmRequestProcessor = class extends BaseLlmRequestProcessor {
       }
       llmRequest.appendInstructions([instruction]);
     }
+    if (agent.outputSchema) {
+      try {
+        const raw = zodToJsonSchema2(agent.outputSchema, {
+          target: "jsonSchema7",
+          $refStrategy: "none"
+        });
+        const { $schema, ...json } = raw || {};
+        llmRequest.appendInstructions([
+          "You must respond with application/json that validates against this JSON Schema:",
+          "```json",
+          JSON.stringify(json, null, 2),
+          "```"
+        ]);
+      } catch {
+      }
+    }
     for await (const _ of []) {
       yield _;
     }
@@ -8493,12 +8512,87 @@ function removeThoughtFromRequest(llmRequest) {
 var requestProcessor7 = new NlPlanningRequestProcessor();
 var responseProcessor2 = new NlPlanningResponseProcessor();
 
+// src/flows/llm-flows/output-schema.ts
+init_logger();
+var OutputSchemaResponseProcessor = class extends BaseLlmResponseProcessor {
+  logger = new Logger({ name: "OutputSchemaResponseProcessor" });
+  async *runAsync(invocationContext, llmResponse) {
+    if (!llmResponse || !llmResponse.content || !llmResponse.content.parts || llmResponse.content.parts.length === 0) {
+      return;
+    }
+    const agent = invocationContext.agent;
+    if (!("outputSchema" in agent) || !agent.outputSchema) {
+      return;
+    }
+    let textContent = llmResponse.content.parts.map((part) => {
+      if (part && typeof part === "object" && "text" in part) {
+        return part.text || "";
+      }
+      return "";
+    }).join("");
+    if (!textContent.trim()) {
+      return;
+    }
+    try {
+      const parsed = JSON.parse(textContent);
+      const validated = agent.outputSchema.parse(parsed);
+      textContent = JSON.stringify(validated, null, 2);
+      llmResponse.content.parts = llmResponse.content.parts.map((part) => {
+        if (part && typeof part === "object" && "text" in part) {
+          return {
+            ...part,
+            text: textContent
+          };
+        }
+        return part;
+      });
+      this.logger.debug("Output schema validation successful", {
+        agent: agent.name,
+        originalLength: textContent.length,
+        validatedKeys: Object.keys(validated)
+      });
+    } catch (error) {
+      const errorMessage = error instanceof Error ? error.message : String(error);
+      const detailedError = `Output schema validation failed for agent '${agent.name}': ${errorMessage}`;
+      this.logger.error(detailedError, {
+        agent: agent.name,
+        responseContent: textContent.substring(0, 200) + (textContent.length > 200 ? "..." : ""),
+        error: errorMessage
+      });
+      llmResponse.errorCode = "OUTPUT_SCHEMA_VALIDATION_FAILED";
+      llmResponse.errorMessage = detailedError;
+      llmResponse.error = new Error(detailedError);
+      const errorEvent = new Event({
+        id: Event.newId(),
+        invocationId: invocationContext.invocationId,
+        author: agent.name,
+        branch: invocationContext.branch,
+        content: {
+          role: "assistant",
+          parts: [
+            {
+              text: `Error: ${detailedError}`
+            }
+          ]
+        }
+      });
+      errorEvent.errorCode = "OUTPUT_SCHEMA_VALIDATION_FAILED";
+      errorEvent.errorMessage = detailedError;
+      errorEvent.error = new Error(detailedError);
+      yield errorEvent;
+    }
+  }
+};
+var responseProcessor3 = new OutputSchemaResponseProcessor();
+
 // src/flows/llm-flows/shared-memory.ts
 var SharedMemoryRequestProcessor = class extends BaseLlmRequestProcessor {
   async *runAsync(invocationContext, llmRequest) {
     const memoryService = invocationContext.memoryService;
     if (!memoryService) return;
-    const lastUserEvent = invocationContext.session.events.findLast((e) => e.author === "user" && e.content?.parts?.length);
+    const lastUserEvent = invocationContext.session.events.findLast(
+      (e) => e.author === "user" && e.content?.parts?.length
+    );
     if (!lastUserEvent) return;
     const query = (lastUserEvent.content.parts ?? []).map((p) => p.text || "").join(" ");
     const results = await memoryService.searchMemory({
@@ -8557,8 +8651,10 @@ var SingleFlow = class extends BaseLlmFlow {
     this.responseProcessors.push(
       responseProcessor2,
       // Phase 5: NL Planning
+      responseProcessor3,
+      // Phase 6: Output Schema validation and parsing - validates response against agent's output schema
       responseProcessor
-      // Phase 5: Code Execution (placeholder)
+      // Phase 7: Code Execution (placeholder)
     );
     this.logger.debug("SingleFlow initialized with processors");
   }
@@ -8731,12 +8827,26 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
    * The input schema when agent is used as a tool
    */
   inputSchema;
-  // Schema type - depends on specific implementation
   /**
    * The output schema when agent replies
    */
   outputSchema;
-  // Schema type - depends on specific implementation
+  /**
+   * Callback or list of callbacks to be called before calling the LLM
+   */
+  beforeModelCallback;
+  /**
+   * Callback or list of callbacks to be called after calling the LLM
+   */
+  afterModelCallback;
+  /**
+   * Callback or list of callbacks to be called before calling a tool
+   */
+  beforeToolCallback;
+  /**
+   * Callback or list of callbacks to be called after calling a tool
+   */
+  afterToolCallback;
   logger = new Logger({ name: "LlmAgent" });
   /**
    * Constructor for LlmAgent
@@ -8767,6 +8877,11 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
     this.generateContentConfig = config.generateContentConfig;
     this.inputSchema = config.inputSchema;
     this.outputSchema = config.outputSchema;
+    this.beforeModelCallback = config.beforeModelCallback;
+    this.afterModelCallback = config.afterModelCallback;
+    this.beforeToolCallback = config.beforeToolCallback;
+    this.afterToolCallback = config.afterToolCallback;
+    this.validateOutputSchemaConfig();
   }
   /**
    * The resolved model field as BaseLLM
@@ -8783,13 +8898,15 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
       return new AiSdkLlm(this.model);
     }
     let ancestorAgent = this.parentAgent;
-    while (ancestorAgent !== null) {
+    while (ancestorAgent !== null && ancestorAgent !== void 0) {
       if (ancestorAgent instanceof _LlmAgent) {
         return ancestorAgent.canonicalModel;
       }
       ancestorAgent = ancestorAgent.parentAgent;
     }
-    throw new Error(`No model found for ${this.name}.`);
+    throw new Error(
+      `No model found for agent "${this.name}". Please specify a model directly on this agent using the 'model' property`
+    );
   }
   /**
    * The resolved instruction field to construct instruction for this agent
@@ -8829,6 +8946,80 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
     }
     return resolvedTools;
   }
+  /**
+   * Gets the canonical before model callbacks as an array
+   */
+  get canonicalBeforeModelCallbacks() {
+    if (!this.beforeModelCallback) {
+      return [];
+    }
+    if (Array.isArray(this.beforeModelCallback)) {
+      return this.beforeModelCallback;
+    }
+    return [this.beforeModelCallback];
+  }
+  /**
+   * Gets the canonical after model callbacks as an array
+   */
+  get canonicalAfterModelCallbacks() {
+    if (!this.afterModelCallback) {
+      return [];
+    }
+    if (Array.isArray(this.afterModelCallback)) {
+      return this.afterModelCallback;
+    }
+    return [this.afterModelCallback];
+  }
+  /**
+   * Gets the canonical before tool callbacks as an array
+   */
+  get canonicalBeforeToolCallbacks() {
+    if (!this.beforeToolCallback) {
+      return [];
+    }
+    if (Array.isArray(this.beforeToolCallback)) {
+      return this.beforeToolCallback;
+    }
+    return [this.beforeToolCallback];
+  }
+  /**
+   * Gets the canonical after tool callbacks as an array
+   */
+  get canonicalAfterToolCallbacks() {
+    if (!this.afterToolCallback) {
+      return [];
+    }
+    if (Array.isArray(this.afterToolCallback)) {
+      return this.afterToolCallback;
+    }
+    return [this.afterToolCallback];
+  }
+  /**
+   * Validates output schema configuration
+   * This matches the Python implementation's __check_output_schema
+   */
+  validateOutputSchemaConfig() {
+    if (!this.outputSchema) {
+      return;
+    }
+    if (!this.disallowTransferToParent || !this.disallowTransferToPeers) {
+      this.logger.warn(
+        `Invalid config for agent ${this.name}: output_schema cannot co-exist with agent transfer configurations. Setting disallow_transfer_to_parent=true, disallow_transfer_to_peers=true`
+      );
+      this.disallowTransferToParent = true;
+      this.disallowTransferToPeers = true;
+    }
+    if (this.subAgents && this.subAgents.length > 0) {
+      throw new Error(
+        `Invalid config for agent ${this.name}: if output_schema is set, sub_agents must be empty to disable agent transfer.`
+      );
+    }
+    if (this.tools && this.tools.length > 0) {
+      throw new Error(
+        `Invalid config for agent ${this.name}: if output_schema is set, tools must be empty`
+      );
+    }
+  }
   /**
    * Gets the appropriate LLM flow for this agent
    * This matches the Python implementation's _llm_flow property
@@ -8844,8 +9035,28 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
    * This matches the Python implementation's __maybe_save_output_to_state
    */
   maybeSaveOutputToState(event) {
+    if (event.author !== this.name) {
+      this.logger.debug(
+        `Skipping output save for agent ${this.name}: event authored by ${event.author}`
+      );
+      return;
+    }
     if (this.outputKey && event.isFinalResponse() && event.content?.parts) {
-      const result = event.content.parts.map((part) => part.text || "").join("");
+      let result = event.content.parts.map((part) => part.text || "").join("");
+      if (this.outputSchema) {
+        if (!result.trim()) {
+          return;
+        }
+        try {
+          const parsed = JSON.parse(result);
+          result = this.outputSchema.parse(parsed);
+        } catch (error) {
+          this.logger.error("Failed to validate output with schema:", error);
+          throw new Error(
+            `Output validation failed: ${error instanceof Error ? error.message : String(error)}`
+          );
+        }
+      }
       if (result) {
         if (!event.actions.stateDelta) {
           event.actions.stateDelta = {};
@@ -10188,6 +10399,14 @@ var AgentBuilder = class _AgentBuilder {
     this.config.instruction = instruction;
     return this;
   }
+  withInputSchema(schema) {
+    this.config.inputSchema = schema;
+    return this;
+  }
+  withOutputSchema(schema) {
+    this.config.outputSchema = schema;
+    return this;
+  }
   /**
    * Add tools to the agent
    * @param tools Tools to add to the agent
@@ -10394,6 +10613,14 @@ var AgentBuilder = class _AgentBuilder {
     }
     return { agent, runner, session };
   }
+  /**
+   * Type-safe build method for agents with output schemas
+   * Provides better type inference for the ask method return type
+   */
+  async buildWithSchema() {
+    const result = await this.build();
+    return result;
+  }
   /**
    * Quick execution helper - build and run a message
    * @param message Message to send to the agent (string or full message object)
@@ -10428,7 +10655,9 @@ var AgentBuilder = class _AgentBuilder {
           memoryService: this.memoryService,
           artifactService: this.artifactService,
           outputKey: this.config.outputKey,
-          sessionService: this.sessionService
+          sessionService: this.sessionService,
+          inputSchema: this.config.inputSchema,
+          outputSchema: this.config.outputSchema
         });
       }
       case "sequential":
@@ -10487,14 +10716,16 @@ var AgentBuilder = class _AgentBuilder {
     return `app-${this.config.name}`;
   }
   /**
-   * Create enhanced runner with simplified API
+   * Create enhanced runner with simplified API and proper typing
    * @param baseRunner The base runner instance
    * @param session The session instance
    * @returns Enhanced runner with simplified API
    */
   createEnhancedRunner(baseRunner, session) {
     const sessionOptions = this.sessionOptions;
+    const outputSchema = this.config.outputSchema;
     return {
+      __outputSchema: outputSchema,
       async ask(message) {
         const newMessage = typeof message === "string" ? { parts: [{ text: message }] } : typeof message === "object" && "contents" in message ? { parts: message.contents[message.contents.length - 1].parts } : message;
         let response = "";
@@ -10515,6 +10746,18 @@ var AgentBuilder = class _AgentBuilder {
             }
           }
         }
+        if (outputSchema) {
+          try {
+            const parsed = JSON.parse(response);
+            return outputSchema.parse(parsed);
+          } catch (parseError) {
+            try {
+              return outputSchema.parse(response);
+            } catch (validationError) {
+              return response.trim();
+            }
+          }
+        }
         return response.trim();
       },
       runAsync(params) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@iqai/adk",
-  "version": "0.1.20",
+  "version": "0.1.21",
   "description": "Agent Development Kit for TypeScript with multi-provider LLM support",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",