@iqai/adk 0.3.7 → 0.4.1

package/CHANGELOG.md

@@ -1,5 +1,22 @@
 # @iqai/adk
 
+## 0.4.1
+
+### Patch Changes
+
+- 1b00e47: Add export to VertexAiRagMemoryService
+
+## 0.4.0
+
+### Minor Changes
+
+- c538a1a: feat(adk): align before/after model callback signatures with runtime (single object arg) and wire before/after tool callbacks into tool execution.
+
+  - beforeModelCallback/afterModelCallback now receive `{ callbackContext, llmRequest|llmResponse }` to match runtime invocation; removes need for casts in examples.
+  - beforeToolCallback/afterToolCallback are now invoked around tool execution; allow argument mutation and result override.
+  - tracing updated to include final args and the produced event.
+  - minor lint/style cleanups in flows.
+
 ## 0.3.7
 
 ### Patch Changes

package/dist/index.d.mts

@@ -1252,7 +1252,10 @@ type ToolUnion = BaseTool | ((...args: any[]) => any);
 /**
  * Single before model callback type
  */
-type SingleBeforeModelCallback = (callbackContext: CallbackContext, llmRequest: LlmRequest) => LlmResponse | null | Promise<LlmResponse | null>;
+type SingleBeforeModelCallback = (args: {
+    callbackContext: CallbackContext;
+    llmRequest: LlmRequest;
+}) => LlmResponse | null | Promise<LlmResponse | null>;
 /**
  * Before model callback type (single or array)
  */
@@ -1260,7 +1263,10 @@ type BeforeModelCallback = SingleBeforeModelCallback | SingleBeforeModelCallback
 /**
  * Single after model callback type
  */
-type SingleAfterModelCallback = (callbackContext: CallbackContext, llmResponse: LlmResponse) => LlmResponse | null | Promise<LlmResponse | null>;
+type SingleAfterModelCallback = (args: {
+    callbackContext: CallbackContext;
+    llmResponse: LlmResponse;
+}) => LlmResponse | null | Promise<LlmResponse | null>;
 /**
  * After model callback type (single or array)
  */
@@ -1517,7 +1523,7 @@ declare class LlmAgent<T extends BaseLlm = BaseLlm> extends BaseAgent {
      * The resolved tools field as a list of BaseTool based on the context
      * This method is only for use by Agent Development Kit
      */
-    canonicalTools(ctx?: ReadonlyContext): Promise<BaseTool[]>;
+    canonicalTools(_ctx?: ReadonlyContext): Promise<BaseTool[]>;
     /**
      * Gets the canonical before model callbacks as an array
      */
@@ -4416,14 +4422,46 @@ declare class InMemoryMemoryService implements BaseMemoryService {
     clear(): void;
 }
 
+/**
+ * A memory service that uses Vertex AI RAG for storage and retrieval.
+ */
+declare class VertexAiRagMemoryService implements BaseMemoryService {
+    private _vertexRagStore;
+    /**
+     * Initializes a VertexAiRagMemoryService.
+     *
+     * @param ragCorpus The name of the Vertex AI RAG corpus to use. Format:
+     *   `projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}`
+     *   or `{rag_corpus_id}`
+     * @param similarityTopK The number of contexts to retrieve.
+     * @param vectorDistanceThreshold Only returns contexts with vector distance
+     *   smaller than the threshold.
+     */
+    constructor(ragCorpus?: string, similarityTopK?: number, vectorDistanceThreshold?: number);
+    /**
+     * Adds a session to the memory service
+     */
+    addSessionToMemory(session: Session): Promise<void>;
+    /**
+     * Searches for sessions that match the query using rag.retrieval_query
+     */
+    searchMemory(options: {
+        appName: string;
+        userId: string;
+        query: string;
+    }): Promise<SearchMemoryResponse>;
+}
+
 /**
  * Memory Services for the Agent Development Kit
  */
 
 type index$4_InMemoryMemoryService = InMemoryMemoryService;
 declare const index$4_InMemoryMemoryService: typeof InMemoryMemoryService;
+type index$4_VertexAiRagMemoryService = VertexAiRagMemoryService;
+declare const index$4_VertexAiRagMemoryService: typeof VertexAiRagMemoryService;
 declare namespace index$4 {
-  export { index$4_InMemoryMemoryService as InMemoryMemoryService };
+  export { index$4_InMemoryMemoryService as InMemoryMemoryService, index$4_VertexAiRagMemoryService as VertexAiRagMemoryService };
 }
 
 /**
@@ -5569,4 +5607,4 @@ declare const traceLlmCall: (invocationContext: InvocationContext, eventId: stri
 
 declare const VERSION = "0.1.0";
 
-export { AF_FUNCTION_CALL_ID_PREFIX, type AfterAgentCallback, type AfterModelCallback, type AfterToolCallback, LlmAgent as Agent, AgentBuilder, type AgentBuilderConfig, type AgentBuilderWithSchema, AgentEvaluator, AgentTool, type AgentToolConfig, type AgentType, index$5 as Agents, AiSdkLlm, AnthropicLlm, ApiKeyCredential, ApiKeyScheme, AuthConfig, AuthCredential, AuthCredentialType, AuthHandler, AuthScheme, AuthSchemeType, AuthTool, type AuthToolArguments, AutoFlow, BaseAgent, type BaseAgentType, BaseCodeExecutor, type BaseCodeExecutorConfig, BaseLLMConnection, BaseLlm, BaseLlmFlow, BaseLlmRequestProcessor, BaseLlmResponseProcessor, type BaseMemoryService, BasePlanner, BaseSessionService, BaseTool, BasicAuthCredential, BearerTokenCredential, type BeforeAgentCallback, type BeforeModelCallback, type BeforeToolCallback, type BuildFunctionDeclarationOptions, type BuiltAgent, BuiltInCodeExecutor, BuiltInPlanner, CallbackContext, type CodeExecutionInput, type CodeExecutionResult, CodeExecutionUtils, CodeExecutorContext, type CreateToolConfig, type CreateToolConfigWithSchema, type CreateToolConfigWithoutSchema, DatabaseSessionService, EnhancedAuthConfig, type EnhancedRunner, type EvalCase, type EvalCaseResult, type EvalMetric, type EvalMetricResult, type EvalMetricResultPerInvocation, EvalResult, type EvalSet, type EvalSetResult, EvalStatus, type EvaluateConfig, index as Evaluation, type EvaluationResult, Evaluator, Event, EventActions, index$2 as Events, ExitLoopTool, type File, FileOperationsTool, FinalResponseMatchV2Evaluator, index$1 as Flows, type FullMessage, FunctionTool, GcsArtifactService, type GetSessionConfig, GetUserChoiceTool, GoogleLlm, GoogleSearch, HttpRequestTool, HttpScheme, InMemoryArtifactService, InMemoryMemoryService, InMemoryRunner, InMemorySessionService, type InstructionProvider, type IntermediateData, type Interval, type Invocation, InvocationContext, type JudgeModelOptions, LLMRegistry, LangGraphAgent, type LangGraphAgentConfig, type LangGraphNode, type ListSessionsResponse, LlmAgent, type LlmAgentConfig, LlmCallsLimitExceededError, type LlmModel, type LlmModelConfig, LlmRequest, LlmResponse, LoadArtifactsTool, LoadMemoryTool, LocalEvalService, LoopAgent, type LoopAgentConfig, McpAbi, McpAtp, McpBamm, McpCoinGecko, type McpConfig, McpDiscord, McpError, McpErrorType, McpFilesystem, McpFraxlend, McpGeneric, McpIqWiki, McpMemory, McpNearAgent, McpNearIntents, McpOdos, McpSamplingHandler, type McpSamplingRequest, type McpSamplingResponse, type McpServerConfig, McpTelegram, McpToolset, type McpTransportType, McpUpbit, index$4 as Memory, type MessagePart, type MetricInfo, type MetricValueInfo, index$6 as Models, type MultiAgentResponse, OAuth2Credential, OAuth2Scheme, type OAuthFlow, type OAuthFlows, OpenAiLlm, OpenIdConnectScheme, ParallelAgent, type ParallelAgentConfig, type PerInvocationResult, PlanReActPlanner, PrebuiltMetrics, REQUEST_EUC_FUNCTION_CALL_NAME, ReadonlyContext, RougeEvaluator, RunConfig, Runner, type RunnerAskReturn, SafetyEvaluatorV1, type SamplingHandler, type SearchMemoryResponse, SequentialAgent, type SequentialAgentConfig, type Session, type SessionInput, type SessionOptions, index$3 as Sessions, type SingleAfterModelCallback, type SingleAfterToolCallback, type SingleAgentCallback, type SingleBeforeModelCallback, type SingleBeforeToolCallback, SingleFlow, State, StreamingMode, type TelemetryConfig, TelemetryService, type ThinkingConfig, type ToolConfig, ToolContext, type ToolUnion, index$7 as Tools, TrajectoryEvaluator, TransferToAgentTool, UserInteractionTool, VERSION, VertexAiSessionService, _findFunctionCallEventIfLastEventIsFunctionResponse, adkToMcpToolType, requestProcessor$2 as agentTransferRequestProcessor, requestProcessor$6 as basicRequestProcessor, buildFunctionDeclaration, requestProcessor as codeExecutionRequestProcessor, responseProcessor as codeExecutionResponseProcessor, requestProcessor$3 as contentRequestProcessor, convertMcpToolToBaseTool, createAuthToolArguments, createBranchContextForSubAgent, createDatabaseSessionService, createFunctionTool, createMysqlSessionService, createPostgresSessionService, createSamplingHandler, createSqliteSessionService, createTool, generateAuthEvent, generateClientFunctionCallId, getLongRunningFunctionCalls, getMcpTools, handleFunctionCallsAsync, handleFunctionCallsLive, requestProcessor$5 as identityRequestProcessor, initializeTelemetry, injectSessionState, requestProcessor$4 as instructionsRequestProcessor, isEnhancedAuthConfig, jsonSchemaToDeclaration, mcpSchemaToParameters, mergeAgentRun, mergeParallelFunctionResponseEvents, newInvocationContextId, requestProcessor$1 as nlPlanningRequestProcessor, responseProcessor$1 as nlPlanningResponseProcessor, normalizeJsonSchema, populateClientFunctionCallId, registerProviders, removeClientFunctionCallId, requestProcessor$7 as requestProcessor, shutdownTelemetry, telemetryService, traceLlmCall, traceToolCall, tracer };
+export { AF_FUNCTION_CALL_ID_PREFIX, type AfterAgentCallback, type AfterModelCallback, type AfterToolCallback, LlmAgent as Agent, AgentBuilder, type AgentBuilderConfig, type AgentBuilderWithSchema, AgentEvaluator, AgentTool, type AgentToolConfig, type AgentType, index$5 as Agents, AiSdkLlm, AnthropicLlm, ApiKeyCredential, ApiKeyScheme, AuthConfig, AuthCredential, AuthCredentialType, AuthHandler, AuthScheme, AuthSchemeType, AuthTool, type AuthToolArguments, AutoFlow, BaseAgent, type BaseAgentType, BaseCodeExecutor, type BaseCodeExecutorConfig, BaseLLMConnection, BaseLlm, BaseLlmFlow, BaseLlmRequestProcessor, BaseLlmResponseProcessor, type BaseMemoryService, BasePlanner, BaseSessionService, BaseTool, BasicAuthCredential, BearerTokenCredential, type BeforeAgentCallback, type BeforeModelCallback, type BeforeToolCallback, type BuildFunctionDeclarationOptions, type BuiltAgent, BuiltInCodeExecutor, BuiltInPlanner, CallbackContext, type CodeExecutionInput, type CodeExecutionResult, CodeExecutionUtils, CodeExecutorContext, type CreateToolConfig, type CreateToolConfigWithSchema, type CreateToolConfigWithoutSchema, DatabaseSessionService, EnhancedAuthConfig, type EnhancedRunner, type EvalCase, type EvalCaseResult, type EvalMetric, type EvalMetricResult, type EvalMetricResultPerInvocation, EvalResult, type EvalSet, type EvalSetResult, EvalStatus, type EvaluateConfig, index as Evaluation, type EvaluationResult, Evaluator, Event, EventActions, index$2 as Events, ExitLoopTool, type File, FileOperationsTool, FinalResponseMatchV2Evaluator, index$1 as Flows, type FullMessage, FunctionTool, GcsArtifactService, type GetSessionConfig, GetUserChoiceTool, GoogleLlm, GoogleSearch, HttpRequestTool, HttpScheme, InMemoryArtifactService, InMemoryMemoryService, InMemoryRunner, InMemorySessionService, type InstructionProvider, type IntermediateData, type Interval, type Invocation, InvocationContext, type JudgeModelOptions, LLMRegistry, LangGraphAgent, type LangGraphAgentConfig, type LangGraphNode, type ListSessionsResponse, LlmAgent, type LlmAgentConfig, LlmCallsLimitExceededError, type LlmModel, type LlmModelConfig, LlmRequest, LlmResponse, LoadArtifactsTool, LoadMemoryTool, LocalEvalService, LoopAgent, type LoopAgentConfig, McpAbi, McpAtp, McpBamm, McpCoinGecko, type McpConfig, McpDiscord, McpError, McpErrorType, McpFilesystem, McpFraxlend, McpGeneric, McpIqWiki, McpMemory, McpNearAgent, McpNearIntents, McpOdos, McpSamplingHandler, type McpSamplingRequest, type McpSamplingResponse, type McpServerConfig, McpTelegram, McpToolset, type McpTransportType, McpUpbit, index$4 as Memory, type MessagePart, type MetricInfo, type MetricValueInfo, index$6 as Models, type MultiAgentResponse, OAuth2Credential, OAuth2Scheme, type OAuthFlow, type OAuthFlows, OpenAiLlm, OpenIdConnectScheme, ParallelAgent, type ParallelAgentConfig, type PerInvocationResult, PlanReActPlanner, PrebuiltMetrics, REQUEST_EUC_FUNCTION_CALL_NAME, ReadonlyContext, RougeEvaluator, RunConfig, Runner, type RunnerAskReturn, SafetyEvaluatorV1, type SamplingHandler, type SearchMemoryResponse, SequentialAgent, type SequentialAgentConfig, type Session, type SessionInput, type SessionOptions, index$3 as Sessions, type SingleAfterModelCallback, type SingleAfterToolCallback, type SingleAgentCallback, type SingleBeforeModelCallback, type SingleBeforeToolCallback, SingleFlow, State, StreamingMode, type TelemetryConfig, TelemetryService, type ThinkingConfig, type ToolConfig, ToolContext, type ToolUnion, index$7 as Tools, TrajectoryEvaluator, TransferToAgentTool, UserInteractionTool, VERSION, VertexAiRagMemoryService, VertexAiSessionService, _findFunctionCallEventIfLastEventIsFunctionResponse, adkToMcpToolType, requestProcessor$2 as agentTransferRequestProcessor, requestProcessor$6 as basicRequestProcessor, buildFunctionDeclaration, requestProcessor as codeExecutionRequestProcessor, responseProcessor as codeExecutionResponseProcessor, requestProcessor$3 as contentRequestProcessor, convertMcpToolToBaseTool, createAuthToolArguments, createBranchContextForSubAgent, createDatabaseSessionService, createFunctionTool, createMysqlSessionService, createPostgresSessionService, createSamplingHandler, createSqliteSessionService, createTool, generateAuthEvent, generateClientFunctionCallId, getLongRunningFunctionCalls, getMcpTools, handleFunctionCallsAsync, handleFunctionCallsLive, requestProcessor$5 as identityRequestProcessor, initializeTelemetry, injectSessionState, requestProcessor$4 as instructionsRequestProcessor, isEnhancedAuthConfig, jsonSchemaToDeclaration, mcpSchemaToParameters, mergeAgentRun, mergeParallelFunctionResponseEvents, newInvocationContextId, requestProcessor$1 as nlPlanningRequestProcessor, responseProcessor$1 as nlPlanningResponseProcessor, normalizeJsonSchema, populateClientFunctionCallId, registerProviders, removeClientFunctionCallId, requestProcessor$7 as requestProcessor, shutdownTelemetry, telemetryService, traceLlmCall, traceToolCall, tracer };

package/dist/index.d.ts

@@ -1252,7 +1252,10 @@ type ToolUnion = BaseTool | ((...args: any[]) => any);
 /**
  * Single before model callback type
  */
-type SingleBeforeModelCallback = (callbackContext: CallbackContext, llmRequest: LlmRequest) => LlmResponse | null | Promise<LlmResponse | null>;
+type SingleBeforeModelCallback = (args: {
+    callbackContext: CallbackContext;
+    llmRequest: LlmRequest;
+}) => LlmResponse | null | Promise<LlmResponse | null>;
 /**
  * Before model callback type (single or array)
  */
@@ -1260,7 +1263,10 @@ type BeforeModelCallback = SingleBeforeModelCallback | SingleBeforeModelCallback
 /**
  * Single after model callback type
  */
-type SingleAfterModelCallback = (callbackContext: CallbackContext, llmResponse: LlmResponse) => LlmResponse | null | Promise<LlmResponse | null>;
+type SingleAfterModelCallback = (args: {
+    callbackContext: CallbackContext;
+    llmResponse: LlmResponse;
+}) => LlmResponse | null | Promise<LlmResponse | null>;
 /**
  * After model callback type (single or array)
  */
@@ -1517,7 +1523,7 @@ declare class LlmAgent<T extends BaseLlm = BaseLlm> extends BaseAgent {
      * The resolved tools field as a list of BaseTool based on the context
      * This method is only for use by Agent Development Kit
      */
-    canonicalTools(ctx?: ReadonlyContext): Promise<BaseTool[]>;
+    canonicalTools(_ctx?: ReadonlyContext): Promise<BaseTool[]>;
     /**
      * Gets the canonical before model callbacks as an array
      */
@@ -4416,14 +4422,46 @@ declare class InMemoryMemoryService implements BaseMemoryService {
     clear(): void;
 }
 
+/**
+ * A memory service that uses Vertex AI RAG for storage and retrieval.
+ */
+declare class VertexAiRagMemoryService implements BaseMemoryService {
+    private _vertexRagStore;
+    /**
+     * Initializes a VertexAiRagMemoryService.
+     *
+     * @param ragCorpus The name of the Vertex AI RAG corpus to use. Format:
+     *   `projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}`
+     *   or `{rag_corpus_id}`
+     * @param similarityTopK The number of contexts to retrieve.
+     * @param vectorDistanceThreshold Only returns contexts with vector distance
+     *   smaller than the threshold.
+     */
+    constructor(ragCorpus?: string, similarityTopK?: number, vectorDistanceThreshold?: number);
+    /**
+     * Adds a session to the memory service
+     */
+    addSessionToMemory(session: Session): Promise<void>;
+    /**
+     * Searches for sessions that match the query using rag.retrieval_query
+     */
+    searchMemory(options: {
+        appName: string;
+        userId: string;
+        query: string;
+    }): Promise<SearchMemoryResponse>;
+}
+
 /**
  * Memory Services for the Agent Development Kit
  */
 
 type index$4_InMemoryMemoryService = InMemoryMemoryService;
 declare const index$4_InMemoryMemoryService: typeof InMemoryMemoryService;
+type index$4_VertexAiRagMemoryService = VertexAiRagMemoryService;
+declare const index$4_VertexAiRagMemoryService: typeof VertexAiRagMemoryService;
 declare namespace index$4 {
-  export { index$4_InMemoryMemoryService as InMemoryMemoryService };
+  export { index$4_InMemoryMemoryService as InMemoryMemoryService, index$4_VertexAiRagMemoryService as VertexAiRagMemoryService };
 }
 
 /**
@@ -5569,4 +5607,4 @@ declare const traceLlmCall: (invocationContext: InvocationContext, eventId: stri
 
 declare const VERSION = "0.1.0";
 
-export { AF_FUNCTION_CALL_ID_PREFIX, type AfterAgentCallback, type AfterModelCallback, type AfterToolCallback, LlmAgent as Agent, AgentBuilder, type AgentBuilderConfig, type AgentBuilderWithSchema, AgentEvaluator, AgentTool, type AgentToolConfig, type AgentType, index$5 as Agents, AiSdkLlm, AnthropicLlm, ApiKeyCredential, ApiKeyScheme, AuthConfig, AuthCredential, AuthCredentialType, AuthHandler, AuthScheme, AuthSchemeType, AuthTool, type AuthToolArguments, AutoFlow, BaseAgent, type BaseAgentType, BaseCodeExecutor, type BaseCodeExecutorConfig, BaseLLMConnection, BaseLlm, BaseLlmFlow, BaseLlmRequestProcessor, BaseLlmResponseProcessor, type BaseMemoryService, BasePlanner, BaseSessionService, BaseTool, BasicAuthCredential, BearerTokenCredential, type BeforeAgentCallback, type BeforeModelCallback, type BeforeToolCallback, type BuildFunctionDeclarationOptions, type BuiltAgent, BuiltInCodeExecutor, BuiltInPlanner, CallbackContext, type CodeExecutionInput, type CodeExecutionResult, CodeExecutionUtils, CodeExecutorContext, type CreateToolConfig, type CreateToolConfigWithSchema, type CreateToolConfigWithoutSchema, DatabaseSessionService, EnhancedAuthConfig, type EnhancedRunner, type EvalCase, type EvalCaseResult, type EvalMetric, type EvalMetricResult, type EvalMetricResultPerInvocation, EvalResult, type EvalSet, type EvalSetResult, EvalStatus, type EvaluateConfig, index as Evaluation, type EvaluationResult, Evaluator, Event, EventActions, index$2 as Events, ExitLoopTool, type File, FileOperationsTool, FinalResponseMatchV2Evaluator, index$1 as Flows, type FullMessage, FunctionTool, GcsArtifactService, type GetSessionConfig, GetUserChoiceTool, GoogleLlm, GoogleSearch, HttpRequestTool, HttpScheme, InMemoryArtifactService, InMemoryMemoryService, InMemoryRunner, InMemorySessionService, type InstructionProvider, type IntermediateData, type Interval, type Invocation, InvocationContext, type JudgeModelOptions, LLMRegistry, LangGraphAgent, type LangGraphAgentConfig, type LangGraphNode, type ListSessionsResponse, LlmAgent, type LlmAgentConfig, LlmCallsLimitExceededError, type LlmModel, type LlmModelConfig, LlmRequest, LlmResponse, LoadArtifactsTool, LoadMemoryTool, LocalEvalService, LoopAgent, type LoopAgentConfig, McpAbi, McpAtp, McpBamm, McpCoinGecko, type McpConfig, McpDiscord, McpError, McpErrorType, McpFilesystem, McpFraxlend, McpGeneric, McpIqWiki, McpMemory, McpNearAgent, McpNearIntents, McpOdos, McpSamplingHandler, type McpSamplingRequest, type McpSamplingResponse, type McpServerConfig, McpTelegram, McpToolset, type McpTransportType, McpUpbit, index$4 as Memory, type MessagePart, type MetricInfo, type MetricValueInfo, index$6 as Models, type MultiAgentResponse, OAuth2Credential, OAuth2Scheme, type OAuthFlow, type OAuthFlows, OpenAiLlm, OpenIdConnectScheme, ParallelAgent, type ParallelAgentConfig, type PerInvocationResult, PlanReActPlanner, PrebuiltMetrics, REQUEST_EUC_FUNCTION_CALL_NAME, ReadonlyContext, RougeEvaluator, RunConfig, Runner, type RunnerAskReturn, SafetyEvaluatorV1, type SamplingHandler, type SearchMemoryResponse, SequentialAgent, type SequentialAgentConfig, type Session, type SessionInput, type SessionOptions, index$3 as Sessions, type SingleAfterModelCallback, type SingleAfterToolCallback, type SingleAgentCallback, type SingleBeforeModelCallback, type SingleBeforeToolCallback, SingleFlow, State, StreamingMode, type TelemetryConfig, TelemetryService, type ThinkingConfig, type ToolConfig, ToolContext, type ToolUnion, index$7 as Tools, TrajectoryEvaluator, TransferToAgentTool, UserInteractionTool, VERSION, VertexAiSessionService, _findFunctionCallEventIfLastEventIsFunctionResponse, adkToMcpToolType, requestProcessor$2 as agentTransferRequestProcessor, requestProcessor$6 as basicRequestProcessor, buildFunctionDeclaration, requestProcessor as codeExecutionRequestProcessor, responseProcessor as codeExecutionResponseProcessor, requestProcessor$3 as contentRequestProcessor, convertMcpToolToBaseTool, createAuthToolArguments, createBranchContextForSubAgent, createDatabaseSessionService, createFunctionTool, createMysqlSessionService, createPostgresSessionService, createSamplingHandler, createSqliteSessionService, createTool, generateAuthEvent, generateClientFunctionCallId, getLongRunningFunctionCalls, getMcpTools, handleFunctionCallsAsync, handleFunctionCallsLive, requestProcessor$5 as identityRequestProcessor, initializeTelemetry, injectSessionState, requestProcessor$4 as instructionsRequestProcessor, isEnhancedAuthConfig, jsonSchemaToDeclaration, mcpSchemaToParameters, mergeAgentRun, mergeParallelFunctionResponseEvents, newInvocationContextId, requestProcessor$1 as nlPlanningRequestProcessor, responseProcessor$1 as nlPlanningResponseProcessor, normalizeJsonSchema, populateClientFunctionCallId, registerProviders, removeClientFunctionCallId, requestProcessor$7 as requestProcessor, shutdownTelemetry, telemetryService, traceLlmCall, traceToolCall, tracer };
+export { AF_FUNCTION_CALL_ID_PREFIX, type AfterAgentCallback, type AfterModelCallback, type AfterToolCallback, LlmAgent as Agent, AgentBuilder, type AgentBuilderConfig, type AgentBuilderWithSchema, AgentEvaluator, AgentTool, type AgentToolConfig, type AgentType, index$5 as Agents, AiSdkLlm, AnthropicLlm, ApiKeyCredential, ApiKeyScheme, AuthConfig, AuthCredential, AuthCredentialType, AuthHandler, AuthScheme, AuthSchemeType, AuthTool, type AuthToolArguments, AutoFlow, BaseAgent, type BaseAgentType, BaseCodeExecutor, type BaseCodeExecutorConfig, BaseLLMConnection, BaseLlm, BaseLlmFlow, BaseLlmRequestProcessor, BaseLlmResponseProcessor, type BaseMemoryService, BasePlanner, BaseSessionService, BaseTool, BasicAuthCredential, BearerTokenCredential, type BeforeAgentCallback, type BeforeModelCallback, type BeforeToolCallback, type BuildFunctionDeclarationOptions, type BuiltAgent, BuiltInCodeExecutor, BuiltInPlanner, CallbackContext, type CodeExecutionInput, type CodeExecutionResult, CodeExecutionUtils, CodeExecutorContext, type CreateToolConfig, type CreateToolConfigWithSchema, type CreateToolConfigWithoutSchema, DatabaseSessionService, EnhancedAuthConfig, type EnhancedRunner, type EvalCase, type EvalCaseResult, type EvalMetric, type EvalMetricResult, type EvalMetricResultPerInvocation, EvalResult, type EvalSet, type EvalSetResult, EvalStatus, type EvaluateConfig, index as Evaluation, type EvaluationResult, Evaluator, Event, EventActions, index$2 as Events, ExitLoopTool, type File, FileOperationsTool, FinalResponseMatchV2Evaluator, index$1 as Flows, type FullMessage, FunctionTool, GcsArtifactService, type GetSessionConfig, GetUserChoiceTool, GoogleLlm, GoogleSearch, HttpRequestTool, HttpScheme, InMemoryArtifactService, InMemoryMemoryService, InMemoryRunner, InMemorySessionService, type InstructionProvider, type IntermediateData, type Interval, type Invocation, InvocationContext, type JudgeModelOptions, LLMRegistry, LangGraphAgent, type LangGraphAgentConfig, type LangGraphNode, type ListSessionsResponse, LlmAgent, type LlmAgentConfig, LlmCallsLimitExceededError, type LlmModel, type LlmModelConfig, LlmRequest, LlmResponse, LoadArtifactsTool, LoadMemoryTool, LocalEvalService, LoopAgent, type LoopAgentConfig, McpAbi, McpAtp, McpBamm, McpCoinGecko, type McpConfig, McpDiscord, McpError, McpErrorType, McpFilesystem, McpFraxlend, McpGeneric, McpIqWiki, McpMemory, McpNearAgent, McpNearIntents, McpOdos, McpSamplingHandler, type McpSamplingRequest, type McpSamplingResponse, type McpServerConfig, McpTelegram, McpToolset, type McpTransportType, McpUpbit, index$4 as Memory, type MessagePart, type MetricInfo, type MetricValueInfo, index$6 as Models, type MultiAgentResponse, OAuth2Credential, OAuth2Scheme, type OAuthFlow, type OAuthFlows, OpenAiLlm, OpenIdConnectScheme, ParallelAgent, type ParallelAgentConfig, type PerInvocationResult, PlanReActPlanner, PrebuiltMetrics, REQUEST_EUC_FUNCTION_CALL_NAME, ReadonlyContext, RougeEvaluator, RunConfig, Runner, type RunnerAskReturn, SafetyEvaluatorV1, type SamplingHandler, type SearchMemoryResponse, SequentialAgent, type SequentialAgentConfig, type Session, type SessionInput, type SessionOptions, index$3 as Sessions, type SingleAfterModelCallback, type SingleAfterToolCallback, type SingleAgentCallback, type SingleBeforeModelCallback, type SingleBeforeToolCallback, SingleFlow, State, StreamingMode, type TelemetryConfig, TelemetryService, type ThinkingConfig, type ToolConfig, ToolContext, type ToolUnion, index$7 as Tools, TrajectoryEvaluator, TransferToAgentTool, UserInteractionTool, VERSION, VertexAiRagMemoryService, VertexAiSessionService, _findFunctionCallEventIfLastEventIsFunctionResponse, adkToMcpToolType, requestProcessor$2 as agentTransferRequestProcessor, requestProcessor$6 as basicRequestProcessor, buildFunctionDeclaration, requestProcessor as codeExecutionRequestProcessor, responseProcessor as codeExecutionResponseProcessor, requestProcessor$3 as contentRequestProcessor, convertMcpToolToBaseTool, createAuthToolArguments, createBranchContextForSubAgent, createDatabaseSessionService, createFunctionTool, createMysqlSessionService, createPostgresSessionService, createSamplingHandler, createSqliteSessionService, createTool, generateAuthEvent, generateClientFunctionCallId, getLongRunningFunctionCalls, getMcpTools, handleFunctionCallsAsync, handleFunctionCallsLive, requestProcessor$5 as identityRequestProcessor, initializeTelemetry, injectSessionState, requestProcessor$4 as instructionsRequestProcessor, isEnhancedAuthConfig, jsonSchemaToDeclaration, mcpSchemaToParameters, mergeAgentRun, mergeParallelFunctionResponseEvents, newInvocationContextId, requestProcessor$1 as nlPlanningRequestProcessor, responseProcessor$1 as nlPlanningResponseProcessor, normalizeJsonSchema, populateClientFunctionCallId, registerProviders, removeClientFunctionCallId, requestProcessor$7 as requestProcessor, shutdownTelemetry, telemetryService, traceLlmCall, traceToolCall, tracer };

package/dist/index.js

@@ -6531,7 +6531,6 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
     }
     const { tool, toolContext } = getToolAndContext(
       invocationContext,
-      functionCallEvent,
       functionCall,
       toolsDict
     );
@@ -6541,10 +6540,44 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
     const spanContext = _api.trace.setSpan(_api.context.active(), span);
     try {
       const functionResponse = await _api.context.with(spanContext, async () => {
-        const result = await callToolAsync(tool, functionArgs, toolContext);
+        const argsForTool = { ...functionArgs };
+        if (isLlmAgent2(agent)) {
+          for (const cb of agent.canonicalBeforeToolCallbacks) {
+            const maybeOverride = await cb(tool, argsForTool, toolContext);
+            if (maybeOverride !== null && maybeOverride !== void 0) {
+              const overriddenEvent = buildResponseEvent(
+                tool,
+                maybeOverride,
+                toolContext,
+                invocationContext
+              );
+              telemetryService.traceToolCall(
+                tool,
+                argsForTool,
+                overriddenEvent
+              );
+              return { result: maybeOverride, event: overriddenEvent };
+            }
+          }
+        }
+        let result = await callToolAsync(tool, argsForTool, toolContext);
         if (tool.isLongRunning && !result) {
           return null;
         }
+        if (isLlmAgent2(agent)) {
+          for (const cb of agent.canonicalAfterToolCallbacks) {
+            const maybeModified = await cb(
+              tool,
+              argsForTool,
+              toolContext,
+              result
+            );
+            if (maybeModified !== null && maybeModified !== void 0) {
+              result = maybeModified;
+              break;
+            }
+          }
+        }
         const functionResponseEvent = buildResponseEvent(
           tool,
           result,
@@ -6553,7 +6586,7 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
         );
         telemetryService.traceToolCall(
           tool,
-          functionArgs,
+          argsForTool,
           functionResponseEvent
         );
         return { result, event: functionResponseEvent };
@@ -6583,7 +6616,7 @@ async function handleFunctionCallsLive(invocationContext, functionCallEvent, too
     toolsDict
   );
 }
-function getToolAndContext(invocationContext, functionCallEvent, functionCall, toolsDict) {
+function getToolAndContext(invocationContext, functionCall, toolsDict) {
   if (!(functionCall.name in toolsDict)) {
     throw new Error(
       `Function ${functionCall.name} is not found in the tools_dict.`
@@ -6700,12 +6733,10 @@ var BaseLlmFlow = (_class25 = class {constructor() { _class25.prototype.__init43
   }
   async *_runOneStepAsync(invocationContext) {
     const llmRequest = new LlmRequest();
-    let preprocessEventCount = 0;
     for await (const event of this._preprocessAsync(
       invocationContext,
       llmRequest
     )) {
-      preprocessEventCount++;
       yield event;
     }
     if (invocationContext.endInvocation) {
@@ -6718,13 +6749,11 @@ var BaseLlmFlow = (_class25 = class {constructor() { _class25.prototype.__init43
       author: invocationContext.agent.name,
       branch: invocationContext.branch
     });
-    let llmResponseCount = 0;
     for await (const llmResponse of this._callLlmAsync(
       invocationContext,
       llmRequest,
       modelResponseEvent
     )) {
-      llmResponseCount++;
       for await (const event of this._postprocessAsync(
         invocationContext,
         llmRequest,
@@ -9241,7 +9270,7 @@ var LlmAgent = (_class27 = class _LlmAgent extends BaseAgent {
    * The resolved tools field as a list of BaseTool based on the context
    * This method is only for use by Agent Development Kit
    */
-  async canonicalTools(ctx) {
+  async canonicalTools(_ctx) {
     const resolvedTools = [];
     for (const toolUnion of this.tools) {
       if (typeof toolUnion === "function") {
@@ -11266,9 +11295,183 @@ var AgentBuilder = (_class33 = class _AgentBuilder {
 // src/memory/index.ts
 var memory_exports = {};
 __export(memory_exports, {
-  InMemoryMemoryService: () => InMemoryMemoryService
+  InMemoryMemoryService: () => InMemoryMemoryService,
+  VertexAiRagMemoryService: () => VertexAiRagMemoryService
 });
 
+// src/memory/vertex-ai-rag-memory-service.ts
+
+var _fs = require('fs');
+var _os = require('os');
+
+var rag = {
+  async upload_file(options) {
+    console.log("Mock upload_file:", options);
+  },
+  async retrieval_query(options) {
+    console.log("Mock retrieval_query:", options);
+    return { contexts: { contexts: [] } };
+  }
+};
+function _mergeEventLists(eventLists) {
+  const merged = [];
+  while (eventLists.length > 0) {
+    const current = eventLists.shift();
+    const currentTs = new Set(current.map((event) => event.timestamp));
+    let mergeFound = true;
+    while (mergeFound) {
+      mergeFound = false;
+      const remaining = [];
+      for (const other of eventLists) {
+        const otherTs = new Set(other.map((event) => event.timestamp));
+        const hasOverlap = Array.from(currentTs).some((ts) => otherTs.has(ts));
+        if (hasOverlap) {
+          const newEvents = other.filter((e) => !currentTs.has(e.timestamp));
+          current.push(...newEvents);
+          newEvents.forEach((e) => currentTs.add(e.timestamp));
+          mergeFound = true;
+        } else {
+          remaining.push(other);
+        }
+      }
+      eventLists.splice(0, eventLists.length, ...remaining);
+    }
+    merged.push(current);
+  }
+  return merged;
+}
+var VertexAiRagMemoryService = class {
+  
+  /**
+   * Initializes a VertexAiRagMemoryService.
+   *
+   * @param ragCorpus The name of the Vertex AI RAG corpus to use. Format:
+   *   `projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}`
+   *   or `{rag_corpus_id}`
+   * @param similarityTopK The number of contexts to retrieve.
+   * @param vectorDistanceThreshold Only returns contexts with vector distance
+   *   smaller than the threshold.
+   */
+  constructor(ragCorpus, similarityTopK, vectorDistanceThreshold = 10) {
+    this._vertexRagStore = {
+      rag_resources: ragCorpus ? [{ rag_corpus: ragCorpus }] : [],
+      similarity_top_k: similarityTopK,
+      vector_distance_threshold: vectorDistanceThreshold
+    };
+  }
+  /**
+   * Adds a session to the memory service
+   */
+  async addSessionToMemory(session) {
+    const tempFileName = `temp_${_crypto.randomUUID.call(void 0, )}.txt`;
+    const tempFilePath = _path.join.call(void 0, _os.tmpdir.call(void 0, ), tempFileName);
+    try {
+      const outputLines = [];
+      for (const event of session.events) {
+        if (!event.content || !event.content.parts) {
+          continue;
+        }
+        const textParts = event.content.parts.filter((part) => part.text).map((part) => part.text.replace(/\n/g, " "));
+        if (textParts.length > 0) {
+          outputLines.push(
+            JSON.stringify({
+              author: event.author,
+              timestamp: event.timestamp,
+              text: textParts.join(".")
+            })
+          );
+        }
+      }
+      const outputString = outputLines.join("\n");
+      _fs.writeFileSync.call(void 0, tempFilePath, outputString, "utf8");
+      if (!this._vertexRagStore.rag_resources || this._vertexRagStore.rag_resources.length === 0) {
+        throw new Error("Rag resources must be set.");
+      }
+      for (const ragResource of this._vertexRagStore.rag_resources) {
+        await rag.upload_file({
+          corpus_name: ragResource.rag_corpus,
+          path: tempFilePath,
+          // This is the temp workaround as upload file does not support
+          // adding metadata, thus use display_name to store the session info.
+          display_name: `${session.appName}.${session.userId}.${session.id}`
+        });
+      }
+    } finally {
+      try {
+        _fs.unlinkSync.call(void 0, tempFilePath);
+      } catch (error) {
+        console.warn("Failed to delete temporary file:", tempFilePath, error);
+      }
+    }
+  }
+  /**
+   * Searches for sessions that match the query using rag.retrieval_query
+   */
+  async searchMemory(options) {
+    const { appName, userId, query } = options;
+    const response = await rag.retrieval_query({
+      text: query,
+      rag_resources: this._vertexRagStore.rag_resources,
+      rag_corpora: this._vertexRagStore.rag_corpora,
+      similarity_top_k: this._vertexRagStore.similarity_top_k,
+      vector_distance_threshold: this._vertexRagStore.vector_distance_threshold
+    });
+    const memoryResults = [];
+    const sessionEventsMap = /* @__PURE__ */ new Map();
+    for (const context4 of response.contexts.contexts) {
+      if (!context4.source_display_name.startsWith(`${appName}.${userId}.`)) {
+        continue;
+      }
+      const sessionId = context4.source_display_name.split(".").pop();
+      const events = [];
+      if (context4.text) {
+        const lines = context4.text.split("\n");
+        for (const line of lines) {
+          const trimmedLine = line.trim();
+          if (!trimmedLine) {
+            continue;
+          }
+          try {
+            const eventData = JSON.parse(trimmedLine);
+            const author = eventData.author || "";
+            const timestamp = Number.parseFloat(eventData.timestamp || "0");
+            const text = eventData.text || "";
+            const content = {
+              parts: [{ text }]
+            };
+            const event = {
+              author,
+              timestamp,
+              content
+            };
+            events.push(event);
+          } catch (e7) {
+          }
+        }
+      }
+      if (sessionEventsMap.has(sessionId)) {
+        sessionEventsMap.get(sessionId).push(events);
+      } else {
+        sessionEventsMap.set(sessionId, [events]);
+      }
+    }
+    for (const [sessionId, eventLists] of sessionEventsMap.entries()) {
+      const mergedEventLists = _mergeEventLists(eventLists);
+      for (const events of mergedEventLists) {
+        const sortedEvents = events.sort((a, b) => a.timestamp - b.timestamp).filter((event) => event.content);
+        memoryResults.push(
+          ...sortedEvents.map((event) => ({
+            author: event.author,
+            content: event.content,
+            timestamp: formatTimestamp(event.timestamp)
+          }))
+        );
+      }
+    }
+    return { memories: memoryResults };
+  }
+};
+
 // src/sessions/index.ts
 var sessions_exports = {};
 __export(sessions_exports, {
@@ -11674,7 +11877,7 @@ var DatabaseSessionService = (_class34 = class extends BaseSessionService {
     if (!jsonString) return defaultValue;
     try {
       return JSON.parse(jsonString);
-    } catch (e7) {
+    } catch (e8) {
       return defaultValue;
     }
   }
@@ -13700,4 +13903,5 @@ var VERSION = "0.1.0";
 
 
 
-exports.AF_FUNCTION_CALL_ID_PREFIX = AF_FUNCTION_CALL_ID_PREFIX; exports.Agent = LlmAgent; exports.AgentBuilder = AgentBuilder; exports.AgentEvaluator = AgentEvaluator; exports.AgentTool = AgentTool; exports.Agents = agents_exports; exports.AiSdkLlm = AiSdkLlm; exports.AnthropicLlm = AnthropicLlm; exports.ApiKeyCredential = ApiKeyCredential; exports.ApiKeyScheme = ApiKeyScheme; exports.AuthConfig = AuthConfig; exports.AuthCredential = AuthCredential; exports.AuthCredentialType = AuthCredentialType; exports.AuthHandler = AuthHandler; exports.AuthScheme = AuthScheme; exports.AuthSchemeType = AuthSchemeType; exports.AuthTool = AuthTool; exports.AutoFlow = AutoFlow; exports.BaseAgent = BaseAgent; exports.BaseCodeExecutor = BaseCodeExecutor; exports.BaseLLMConnection = BaseLLMConnection; exports.BaseLlm = BaseLlm; exports.BaseLlmFlow = BaseLlmFlow; exports.BaseLlmRequestProcessor = BaseLlmRequestProcessor; exports.BaseLlmResponseProcessor = BaseLlmResponseProcessor; exports.BasePlanner = BasePlanner; exports.BaseSessionService = BaseSessionService; exports.BaseTool = BaseTool; exports.BasicAuthCredential = BasicAuthCredential; exports.BearerTokenCredential = BearerTokenCredential; exports.BuiltInCodeExecutor = BuiltInCodeExecutor; exports.BuiltInPlanner = BuiltInPlanner; exports.CallbackContext = CallbackContext; exports.CodeExecutionUtils = CodeExecutionUtils; exports.CodeExecutorContext = CodeExecutorContext; exports.DatabaseSessionService = DatabaseSessionService; exports.EnhancedAuthConfig = EnhancedAuthConfig; exports.EvalResult = EvalResult; exports.EvalStatus = EvalStatus; exports.Evaluation = evaluation_exports; exports.Evaluator = Evaluator; exports.Event = Event; exports.EventActions = EventActions; exports.Events = events_exports; exports.ExitLoopTool = ExitLoopTool; exports.FileOperationsTool = FileOperationsTool; exports.FinalResponseMatchV2Evaluator = FinalResponseMatchV2Evaluator; exports.Flows = flows_exports; exports.FunctionTool = FunctionTool; exports.GcsArtifactService = GcsArtifactService; exports.GetUserChoiceTool = GetUserChoiceTool; exports.GoogleLlm = GoogleLlm; exports.GoogleSearch = GoogleSearch; exports.HttpRequestTool = HttpRequestTool; exports.HttpScheme = HttpScheme; exports.InMemoryArtifactService = InMemoryArtifactService; exports.InMemoryMemoryService = InMemoryMemoryService; exports.InMemoryRunner = InMemoryRunner; exports.InMemorySessionService = InMemorySessionService; exports.InvocationContext = InvocationContext; exports.LLMRegistry = LLMRegistry; exports.LangGraphAgent = LangGraphAgent; exports.LlmAgent = LlmAgent; exports.LlmCallsLimitExceededError = LlmCallsLimitExceededError; exports.LlmRequest = LlmRequest; exports.LlmResponse = LlmResponse; exports.LoadArtifactsTool = LoadArtifactsTool; exports.LoadMemoryTool = LoadMemoryTool; exports.LocalEvalService = LocalEvalService; exports.LoopAgent = LoopAgent; exports.McpAbi = McpAbi; exports.McpAtp = McpAtp; exports.McpBamm = McpBamm; exports.McpCoinGecko = McpCoinGecko; exports.McpDiscord = McpDiscord; exports.McpError = McpError; exports.McpErrorType = McpErrorType; exports.McpFilesystem = McpFilesystem; exports.McpFraxlend = McpFraxlend; exports.McpGeneric = McpGeneric; exports.McpIqWiki = McpIqWiki; exports.McpMemory = McpMemory; exports.McpNearAgent = McpNearAgent; exports.McpNearIntents = McpNearIntents; exports.McpOdos = McpOdos; exports.McpSamplingHandler = McpSamplingHandler; exports.McpTelegram = McpTelegram; exports.McpToolset = McpToolset; exports.McpUpbit = McpUpbit; exports.Memory = memory_exports; exports.Models = models_exports; exports.OAuth2Credential = OAuth2Credential; exports.OAuth2Scheme = OAuth2Scheme; exports.OpenAiLlm = OpenAiLlm; exports.OpenIdConnectScheme = OpenIdConnectScheme; exports.ParallelAgent = ParallelAgent; exports.PlanReActPlanner = PlanReActPlanner; exports.PrebuiltMetrics = PrebuiltMetrics; exports.REQUEST_EUC_FUNCTION_CALL_NAME = REQUEST_EUC_FUNCTION_CALL_NAME; exports.ReadonlyContext = ReadonlyContext; exports.RougeEvaluator = RougeEvaluator; exports.RunConfig = RunConfig; exports.Runner = Runner; exports.SafetyEvaluatorV1 = SafetyEvaluatorV1; exports.SequentialAgent = SequentialAgent; exports.Sessions = sessions_exports; exports.SingleFlow = SingleFlow; exports.State = State; exports.StreamingMode = StreamingMode; exports.TelemetryService = TelemetryService; exports.ToolContext = ToolContext; exports.Tools = tools_exports; exports.TrajectoryEvaluator = TrajectoryEvaluator; exports.TransferToAgentTool = TransferToAgentTool; exports.UserInteractionTool = UserInteractionTool; exports.VERSION = VERSION; exports.VertexAiSessionService = VertexAiSessionService; exports._findFunctionCallEventIfLastEventIsFunctionResponse = _findFunctionCallEventIfLastEventIsFunctionResponse; exports.adkToMcpToolType = adkToMcpToolType; exports.agentTransferRequestProcessor = requestProcessor8; exports.basicRequestProcessor = requestProcessor2; exports.buildFunctionDeclaration = buildFunctionDeclaration; exports.codeExecutionRequestProcessor = requestProcessor3; exports.codeExecutionResponseProcessor = responseProcessor; exports.contentRequestProcessor = requestProcessor4; exports.convertMcpToolToBaseTool = convertMcpToolToBaseTool; exports.createAuthToolArguments = createAuthToolArguments; exports.createBranchContextForSubAgent = createBranchContextForSubAgent; exports.createDatabaseSessionService = createDatabaseSessionService; exports.createFunctionTool = createFunctionTool; exports.createMysqlSessionService = createMysqlSessionService; exports.createPostgresSessionService = createPostgresSessionService; exports.createSamplingHandler = createSamplingHandler; exports.createSqliteSessionService = createSqliteSessionService; exports.createTool = createTool; exports.generateAuthEvent = generateAuthEvent; exports.generateClientFunctionCallId = generateClientFunctionCallId; exports.getLongRunningFunctionCalls = getLongRunningFunctionCalls; exports.getMcpTools = getMcpTools; exports.handleFunctionCallsAsync = handleFunctionCallsAsync; exports.handleFunctionCallsLive = handleFunctionCallsLive; exports.identityRequestProcessor = requestProcessor5; exports.initializeTelemetry = initializeTelemetry; exports.injectSessionState = injectSessionState; exports.instructionsRequestProcessor = requestProcessor6; exports.isEnhancedAuthConfig = isEnhancedAuthConfig; exports.jsonSchemaToDeclaration = jsonSchemaToDeclaration; exports.mcpSchemaToParameters = mcpSchemaToParameters; exports.mergeAgentRun = mergeAgentRun; exports.mergeParallelFunctionResponseEvents = mergeParallelFunctionResponseEvents; exports.newInvocationContextId = newInvocationContextId; exports.nlPlanningRequestProcessor = requestProcessor7; exports.nlPlanningResponseProcessor = responseProcessor2; exports.normalizeJsonSchema = normalizeJsonSchema; exports.populateClientFunctionCallId = populateClientFunctionCallId; exports.registerProviders = registerProviders; exports.removeClientFunctionCallId = removeClientFunctionCallId; exports.requestProcessor = requestProcessor; exports.shutdownTelemetry = shutdownTelemetry; exports.telemetryService = telemetryService; exports.traceLlmCall = traceLlmCall; exports.traceToolCall = traceToolCall; exports.tracer = tracer;
+
+exports.AF_FUNCTION_CALL_ID_PREFIX = AF_FUNCTION_CALL_ID_PREFIX; exports.Agent = LlmAgent; exports.AgentBuilder = AgentBuilder; exports.AgentEvaluator = AgentEvaluator; exports.AgentTool = AgentTool; exports.Agents = agents_exports; exports.AiSdkLlm = AiSdkLlm; exports.AnthropicLlm = AnthropicLlm; exports.ApiKeyCredential = ApiKeyCredential; exports.ApiKeyScheme = ApiKeyScheme; exports.AuthConfig = AuthConfig; exports.AuthCredential = AuthCredential; exports.AuthCredentialType = AuthCredentialType; exports.AuthHandler = AuthHandler; exports.AuthScheme = AuthScheme; exports.AuthSchemeType = AuthSchemeType; exports.AuthTool = AuthTool; exports.AutoFlow = AutoFlow; exports.BaseAgent = BaseAgent; exports.BaseCodeExecutor = BaseCodeExecutor; exports.BaseLLMConnection = BaseLLMConnection; exports.BaseLlm = BaseLlm; exports.BaseLlmFlow = BaseLlmFlow; exports.BaseLlmRequestProcessor = BaseLlmRequestProcessor; exports.BaseLlmResponseProcessor = BaseLlmResponseProcessor; exports.BasePlanner = BasePlanner; exports.BaseSessionService = BaseSessionService; exports.BaseTool = BaseTool; exports.BasicAuthCredential = BasicAuthCredential; exports.BearerTokenCredential = BearerTokenCredential; exports.BuiltInCodeExecutor = BuiltInCodeExecutor; exports.BuiltInPlanner = BuiltInPlanner; exports.CallbackContext = CallbackContext; exports.CodeExecutionUtils = CodeExecutionUtils; exports.CodeExecutorContext = CodeExecutorContext; exports.DatabaseSessionService = DatabaseSessionService; exports.EnhancedAuthConfig = EnhancedAuthConfig; exports.EvalResult = EvalResult; exports.EvalStatus = EvalStatus; exports.Evaluation = evaluation_exports; exports.Evaluator = Evaluator; exports.Event = Event; exports.EventActions = EventActions; exports.Events = events_exports; exports.ExitLoopTool = ExitLoopTool; exports.FileOperationsTool = FileOperationsTool; exports.FinalResponseMatchV2Evaluator = FinalResponseMatchV2Evaluator; exports.Flows = flows_exports; exports.FunctionTool = FunctionTool; exports.GcsArtifactService = GcsArtifactService; exports.GetUserChoiceTool = GetUserChoiceTool; exports.GoogleLlm = GoogleLlm; exports.GoogleSearch = GoogleSearch; exports.HttpRequestTool = HttpRequestTool; exports.HttpScheme = HttpScheme; exports.InMemoryArtifactService = InMemoryArtifactService; exports.InMemoryMemoryService = InMemoryMemoryService; exports.InMemoryRunner = InMemoryRunner; exports.InMemorySessionService = InMemorySessionService; exports.InvocationContext = InvocationContext; exports.LLMRegistry = LLMRegistry; exports.LangGraphAgent = LangGraphAgent; exports.LlmAgent = LlmAgent; exports.LlmCallsLimitExceededError = LlmCallsLimitExceededError; exports.LlmRequest = LlmRequest; exports.LlmResponse = LlmResponse; exports.LoadArtifactsTool = LoadArtifactsTool; exports.LoadMemoryTool = LoadMemoryTool; exports.LocalEvalService = LocalEvalService; exports.LoopAgent = LoopAgent; exports.McpAbi = McpAbi; exports.McpAtp = McpAtp; exports.McpBamm = McpBamm; exports.McpCoinGecko = McpCoinGecko; exports.McpDiscord = McpDiscord; exports.McpError = McpError; exports.McpErrorType = McpErrorType; exports.McpFilesystem = McpFilesystem; exports.McpFraxlend = McpFraxlend; exports.McpGeneric = McpGeneric; exports.McpIqWiki = McpIqWiki; exports.McpMemory = McpMemory; exports.McpNearAgent = McpNearAgent; exports.McpNearIntents = McpNearIntents; exports.McpOdos = McpOdos; exports.McpSamplingHandler = McpSamplingHandler; exports.McpTelegram = McpTelegram; exports.McpToolset = McpToolset; exports.McpUpbit = McpUpbit; exports.Memory = memory_exports; exports.Models = models_exports; exports.OAuth2Credential = OAuth2Credential; exports.OAuth2Scheme = OAuth2Scheme; exports.OpenAiLlm = OpenAiLlm; exports.OpenIdConnectScheme = OpenIdConnectScheme; exports.ParallelAgent = ParallelAgent; exports.PlanReActPlanner = PlanReActPlanner; exports.PrebuiltMetrics = PrebuiltMetrics; exports.REQUEST_EUC_FUNCTION_CALL_NAME = REQUEST_EUC_FUNCTION_CALL_NAME; exports.ReadonlyContext = ReadonlyContext; exports.RougeEvaluator = RougeEvaluator; exports.RunConfig = RunConfig; exports.Runner = Runner; exports.SafetyEvaluatorV1 = SafetyEvaluatorV1; exports.SequentialAgent = SequentialAgent; exports.Sessions = sessions_exports; exports.SingleFlow = SingleFlow; exports.State = State; exports.StreamingMode = StreamingMode; exports.TelemetryService = TelemetryService; exports.ToolContext = ToolContext; exports.Tools = tools_exports; exports.TrajectoryEvaluator = TrajectoryEvaluator; exports.TransferToAgentTool = TransferToAgentTool; exports.UserInteractionTool = UserInteractionTool; exports.VERSION = VERSION; exports.VertexAiRagMemoryService = VertexAiRagMemoryService; exports.VertexAiSessionService = VertexAiSessionService; exports._findFunctionCallEventIfLastEventIsFunctionResponse = _findFunctionCallEventIfLastEventIsFunctionResponse; exports.adkToMcpToolType = adkToMcpToolType; exports.agentTransferRequestProcessor = requestProcessor8; exports.basicRequestProcessor = requestProcessor2; exports.buildFunctionDeclaration = buildFunctionDeclaration; exports.codeExecutionRequestProcessor = requestProcessor3; exports.codeExecutionResponseProcessor = responseProcessor; exports.contentRequestProcessor = requestProcessor4; exports.convertMcpToolToBaseTool = convertMcpToolToBaseTool; exports.createAuthToolArguments = createAuthToolArguments; exports.createBranchContextForSubAgent = createBranchContextForSubAgent; exports.createDatabaseSessionService = createDatabaseSessionService; exports.createFunctionTool = createFunctionTool; exports.createMysqlSessionService = createMysqlSessionService; exports.createPostgresSessionService = createPostgresSessionService; exports.createSamplingHandler = createSamplingHandler; exports.createSqliteSessionService = createSqliteSessionService; exports.createTool = createTool; exports.generateAuthEvent = generateAuthEvent; exports.generateClientFunctionCallId = generateClientFunctionCallId; exports.getLongRunningFunctionCalls = getLongRunningFunctionCalls; exports.getMcpTools = getMcpTools; exports.handleFunctionCallsAsync = handleFunctionCallsAsync; exports.handleFunctionCallsLive = handleFunctionCallsLive; exports.identityRequestProcessor = requestProcessor5; exports.initializeTelemetry = initializeTelemetry; exports.injectSessionState = injectSessionState; exports.instructionsRequestProcessor = requestProcessor6; exports.isEnhancedAuthConfig = isEnhancedAuthConfig; exports.jsonSchemaToDeclaration = jsonSchemaToDeclaration; exports.mcpSchemaToParameters = mcpSchemaToParameters; exports.mergeAgentRun = mergeAgentRun; exports.mergeParallelFunctionResponseEvents = mergeParallelFunctionResponseEvents; exports.newInvocationContextId = newInvocationContextId; exports.nlPlanningRequestProcessor = requestProcessor7; exports.nlPlanningResponseProcessor = responseProcessor2; exports.normalizeJsonSchema = normalizeJsonSchema; exports.populateClientFunctionCallId = populateClientFunctionCallId; exports.registerProviders = registerProviders; exports.removeClientFunctionCallId = removeClientFunctionCallId; exports.requestProcessor = requestProcessor; exports.shutdownTelemetry = shutdownTelemetry; exports.telemetryService = telemetryService; exports.traceLlmCall = traceLlmCall; exports.traceToolCall = traceToolCall; exports.tracer = tracer;

package/dist/index.mjs

@@ -6531,7 +6531,6 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
     }
     const { tool, toolContext } = getToolAndContext(
       invocationContext,
-      functionCallEvent,
       functionCall,
       toolsDict
     );
@@ -6541,10 +6540,44 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
     const spanContext = trace2.setSpan(context2.active(), span);
     try {
       const functionResponse = await context2.with(spanContext, async () => {
-        const result = await callToolAsync(tool, functionArgs, toolContext);
+        const argsForTool = { ...functionArgs };
+        if (isLlmAgent2(agent)) {
+          for (const cb of agent.canonicalBeforeToolCallbacks) {
+            const maybeOverride = await cb(tool, argsForTool, toolContext);
+            if (maybeOverride !== null && maybeOverride !== void 0) {
+              const overriddenEvent = buildResponseEvent(
+                tool,
+                maybeOverride,
+                toolContext,
+                invocationContext
+              );
+              telemetryService.traceToolCall(
+                tool,
+                argsForTool,
+                overriddenEvent
+              );
+              return { result: maybeOverride, event: overriddenEvent };
+            }
+          }
+        }
+        let result = await callToolAsync(tool, argsForTool, toolContext);
         if (tool.isLongRunning && !result) {
           return null;
         }
+        if (isLlmAgent2(agent)) {
+          for (const cb of agent.canonicalAfterToolCallbacks) {
+            const maybeModified = await cb(
+              tool,
+              argsForTool,
+              toolContext,
+              result
+            );
+            if (maybeModified !== null && maybeModified !== void 0) {
+              result = maybeModified;
+              break;
+            }
+          }
+        }
         const functionResponseEvent = buildResponseEvent(
           tool,
           result,
@@ -6553,7 +6586,7 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
         );
         telemetryService.traceToolCall(
           tool,
-          functionArgs,
+          argsForTool,
           functionResponseEvent
         );
         return { result, event: functionResponseEvent };
@@ -6583,7 +6616,7 @@ async function handleFunctionCallsLive(invocationContext, functionCallEvent, too
     toolsDict
   );
 }
-function getToolAndContext(invocationContext, functionCallEvent, functionCall, toolsDict) {
+function getToolAndContext(invocationContext, functionCall, toolsDict) {
   if (!(functionCall.name in toolsDict)) {
     throw new Error(
       `Function ${functionCall.name} is not found in the tools_dict.`
@@ -6700,12 +6733,10 @@ var BaseLlmFlow = class {
   }
   async *_runOneStepAsync(invocationContext) {
     const llmRequest = new LlmRequest();
-    let preprocessEventCount = 0;
     for await (const event of this._preprocessAsync(
       invocationContext,
       llmRequest
     )) {
-      preprocessEventCount++;
       yield event;
     }
     if (invocationContext.endInvocation) {
@@ -6718,13 +6749,11 @@ var BaseLlmFlow = class {
       author: invocationContext.agent.name,
       branch: invocationContext.branch
     });
-    let llmResponseCount = 0;
     for await (const llmResponse of this._callLlmAsync(
       invocationContext,
       llmRequest,
       modelResponseEvent
     )) {
-      llmResponseCount++;
       for await (const event of this._postprocessAsync(
         invocationContext,
         llmRequest,
@@ -9241,7 +9270,7 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
    * The resolved tools field as a list of BaseTool based on the context
    * This method is only for use by Agent Development Kit
    */
-  async canonicalTools(ctx) {
+  async canonicalTools(_ctx) {
     const resolvedTools = [];
     for (const toolUnion of this.tools) {
       if (typeof toolUnion === "function") {
@@ -11266,9 +11295,183 @@ var AgentBuilder = class _AgentBuilder {
 // src/memory/index.ts
 var memory_exports = {};
 __export(memory_exports, {
-  InMemoryMemoryService: () => InMemoryMemoryService
+  InMemoryMemoryService: () => InMemoryMemoryService,
+  VertexAiRagMemoryService: () => VertexAiRagMemoryService
 });
 
+// src/memory/vertex-ai-rag-memory-service.ts
+import { randomUUID as randomUUID2 } from "crypto";
+import { unlinkSync, writeFileSync } from "fs";
+import { tmpdir } from "os";
+import { join } from "path";
+var rag = {
+  async upload_file(options) {
+    console.log("Mock upload_file:", options);
+  },
+  async retrieval_query(options) {
+    console.log("Mock retrieval_query:", options);
+    return { contexts: { contexts: [] } };
+  }
+};
+function _mergeEventLists(eventLists) {
+  const merged = [];
+  while (eventLists.length > 0) {
+    const current = eventLists.shift();
+    const currentTs = new Set(current.map((event) => event.timestamp));
+    let mergeFound = true;
+    while (mergeFound) {
+      mergeFound = false;
+      const remaining = [];
+      for (const other of eventLists) {
+        const otherTs = new Set(other.map((event) => event.timestamp));
+        const hasOverlap = Array.from(currentTs).some((ts) => otherTs.has(ts));
+        if (hasOverlap) {
+          const newEvents = other.filter((e) => !currentTs.has(e.timestamp));
+          current.push(...newEvents);
+          newEvents.forEach((e) => currentTs.add(e.timestamp));
+          mergeFound = true;
+        } else {
+          remaining.push(other);
+        }
+      }
+      eventLists.splice(0, eventLists.length, ...remaining);
+    }
+    merged.push(current);
+  }
+  return merged;
+}
+var VertexAiRagMemoryService = class {
+  _vertexRagStore;
+  /**
+   * Initializes a VertexAiRagMemoryService.
+   *
+   * @param ragCorpus The name of the Vertex AI RAG corpus to use. Format:
+   *   `projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}`
+   *   or `{rag_corpus_id}`
+   * @param similarityTopK The number of contexts to retrieve.
+   * @param vectorDistanceThreshold Only returns contexts with vector distance
+   *   smaller than the threshold.
+   */
+  constructor(ragCorpus, similarityTopK, vectorDistanceThreshold = 10) {
+    this._vertexRagStore = {
+      rag_resources: ragCorpus ? [{ rag_corpus: ragCorpus }] : [],
+      similarity_top_k: similarityTopK,
+      vector_distance_threshold: vectorDistanceThreshold
+    };
+  }
+  /**
+   * Adds a session to the memory service
+   */
+  async addSessionToMemory(session) {
+    const tempFileName = `temp_${randomUUID2()}.txt`;
+    const tempFilePath = join(tmpdir(), tempFileName);
+    try {
+      const outputLines = [];
+      for (const event of session.events) {
+        if (!event.content || !event.content.parts) {
+          continue;
+        }
+        const textParts = event.content.parts.filter((part) => part.text).map((part) => part.text.replace(/\n/g, " "));
+        if (textParts.length > 0) {
+          outputLines.push(
+            JSON.stringify({
+              author: event.author,
+              timestamp: event.timestamp,
+              text: textParts.join(".")
+            })
+          );
+        }
+      }
+      const outputString = outputLines.join("\n");
+      writeFileSync(tempFilePath, outputString, "utf8");
+      if (!this._vertexRagStore.rag_resources || this._vertexRagStore.rag_resources.length === 0) {
+        throw new Error("Rag resources must be set.");
+      }
+      for (const ragResource of this._vertexRagStore.rag_resources) {
+        await rag.upload_file({
+          corpus_name: ragResource.rag_corpus,
+          path: tempFilePath,
+          // This is the temp workaround as upload file does not support
+          // adding metadata, thus use display_name to store the session info.
+          display_name: `${session.appName}.${session.userId}.${session.id}`
+        });
+      }
+    } finally {
+      try {
+        unlinkSync(tempFilePath);
+      } catch (error) {
+        console.warn("Failed to delete temporary file:", tempFilePath, error);
+      }
+    }
+  }
+  /**
+   * Searches for sessions that match the query using rag.retrieval_query
+   */
+  async searchMemory(options) {
+    const { appName, userId, query } = options;
+    const response = await rag.retrieval_query({
+      text: query,
+      rag_resources: this._vertexRagStore.rag_resources,
+      rag_corpora: this._vertexRagStore.rag_corpora,
+      similarity_top_k: this._vertexRagStore.similarity_top_k,
+      vector_distance_threshold: this._vertexRagStore.vector_distance_threshold
+    });
+    const memoryResults = [];
+    const sessionEventsMap = /* @__PURE__ */ new Map();
+    for (const context4 of response.contexts.contexts) {
+      if (!context4.source_display_name.startsWith(`${appName}.${userId}.`)) {
+        continue;
+      }
+      const sessionId = context4.source_display_name.split(".").pop();
+      const events = [];
+      if (context4.text) {
+        const lines = context4.text.split("\n");
+        for (const line of lines) {
+          const trimmedLine = line.trim();
+          if (!trimmedLine) {
+            continue;
+          }
+          try {
+            const eventData = JSON.parse(trimmedLine);
+            const author = eventData.author || "";
+            const timestamp = Number.parseFloat(eventData.timestamp || "0");
+            const text = eventData.text || "";
+            const content = {
+              parts: [{ text }]
+            };
+            const event = {
+              author,
+              timestamp,
+              content
+            };
+            events.push(event);
+          } catch {
+          }
+        }
+      }
+      if (sessionEventsMap.has(sessionId)) {
+        sessionEventsMap.get(sessionId).push(events);
+      } else {
+        sessionEventsMap.set(sessionId, [events]);
+      }
+    }
+    for (const [sessionId, eventLists] of sessionEventsMap.entries()) {
+      const mergedEventLists = _mergeEventLists(eventLists);
+      for (const events of mergedEventLists) {
+        const sortedEvents = events.sort((a, b) => a.timestamp - b.timestamp).filter((event) => event.content);
+        memoryResults.push(
+          ...sortedEvents.map((event) => ({
+            author: event.author,
+            content: event.content,
+            timestamp: formatTimestamp(event.timestamp)
+          }))
+        );
+      }
+    }
+    return { memories: memoryResults };
+  }
+};
+
 // src/sessions/index.ts
 var sessions_exports = {};
 __export(sessions_exports, {
@@ -13653,6 +13856,7 @@ export {
   TransferToAgentTool,
   UserInteractionTool,
   VERSION,
+  VertexAiRagMemoryService,
   VertexAiSessionService,
   _findFunctionCallEventIfLastEventIsFunctionResponse,
   adkToMcpToolType,

package/package.json

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