adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
693 lines (692 loc) • 31.1 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.BaseLlmFlow = void 0;
const LlmAgent_1 = require("../../agents/LlmAgent");
const TranscriptionEntry_1 = require("../../agents/TranscriptionEntry");
const Event_1 = require("../../events/Event");
const LlmRequest_1 = require("../../models/LlmRequest");
const CallbackContext_1 = require("../../agents/CallbackContext");
const ToolContext_1 = require("../../tools/ToolContext");
const telemetry = __importStar(require("../../telemetry"));
const RunConfig_1 = require("../../agents/RunConfig");
const LiveRequestQueue_1 = require("../../agents/LiveRequestQueue");
const AudioTranscriber_1 = require("./AudioTranscriber");
const functions = __importStar(require("./functions"));
// Add the close method to LiveRequestQueue prototype
if (typeof LiveRequestQueue_1.LiveRequestQueue.prototype.close !== 'function') {
LiveRequestQueue_1.LiveRequestQueue.prototype.close = function () {
this.sendClose();
};
}
/**
* A basic flow that calls the LLM in a loop until a final response is generated.
* This flow ends when it transfers to another agent.
*/
class BaseLlmFlow {
constructor() {
/**
* List of request processors to run before LLM call
*/
this.requestProcessors = [];
/**
* List of response processors to run after LLM call
*/
this.responseProcessors = [];
}
/**
* Runs the flow using live API.
*
* @param invocationContext The invocation context
* @returns An async generator of events
*/
async *runLive(invocationContext) {
const llmRequest = new LlmRequest_1.LlmRequest();
const eventId = Event_1.Event.newId();
// Preprocess before calling the LLM
yield* this._preprocessAsync(invocationContext, llmRequest);
if (invocationContext.endInvocation) {
return;
}
const llm = this._getLlm(invocationContext);
const llmConnection = await llm.connect(llmRequest);
try {
if (llmRequest.contents) {
// Send conversation history to the model
if (invocationContext.transcriptionCache) {
// Use AudioTranscriber if available
try {
const audioTranscriber = new AudioTranscriber_1.AudioTranscriber();
const contents = audioTranscriber.transcribeFile(invocationContext);
await llmConnection.sendHistory(contents);
invocationContext.transcriptionCache = undefined;
telemetry.traceSendData(invocationContext, eventId, contents);
}
catch (error) {
console.error('Error transcribing audio:', error);
invocationContext.transcriptionCache = undefined;
}
}
else {
await llmConnection.sendHistory(llmRequest.contents);
telemetry.traceSendData(invocationContext, eventId, llmRequest.contents);
}
}
// Start sending task
const sendTask = this._sendToModel(llmConnection, invocationContext);
try {
for await (const event of this._receiveFromModel(llmConnection, eventId, invocationContext, llmRequest)) {
if (!event) {
break;
}
yield event;
// Send back the function response
if (typeof event.getFunctionResponses === 'function' && event.getFunctionResponses().length > 0) {
if (invocationContext.liveRequestQueue && typeof invocationContext.liveRequestQueue.sendContent === 'function' && event.content) {
invocationContext.liveRequestQueue.sendContent(event.content);
}
}
// Store transcription for model responses
if (event.content && event.content.parts &&
event.content.parts.some(part => part.text) &&
!event.partial) {
if (!invocationContext.transcriptionCache) {
invocationContext.transcriptionCache = [];
}
invocationContext.transcriptionCache.push(new TranscriptionEntry_1.TranscriptionEntry({
textContent: JSON.stringify(event.content),
metadata: { type: 'model' }
}));
}
// Check for transfer_to_agent
const hasTransferToAgent = event.content?.parts?.some(part => part.functionResponse?.name === 'transfer_to_agent');
if (hasTransferToAgent) {
await new Promise(resolve => setTimeout(resolve, 1000)); // wait 1 second
// Cancel the tasks that belong to the closed connection
await llmConnection.close();
break;
}
}
}
finally {
// Clean up
try {
await llmConnection.close();
}
catch (error) {
console.error('Error closing connection:', error);
}
}
}
finally {
// Final cleanup
try {
await llmConnection.close();
}
catch (error) {
// Ignore errors during final cleanup
}
}
}
/**
* Sends data to the model in a loop.
*
* @param llmConnection The LLM connection
* @param invocationContext The invocation context
*/
async _sendToModel(llmConnection, invocationContext) {
const liveRequestQueue = invocationContext.liveRequestQueue;
if (!liveRequestQueue) {
return;
}
try {
// Use a variable for the loop condition instead of 'while (true)'
const shouldContinue = true;
while (shouldContinue) {
let liveRequest;
try {
// Use a timeout to allow the event loop to process other events
const timeoutPromise = new Promise((_, reject) => {
setTimeout(() => reject(new Error('Timeout')), 250);
});
liveRequest = await Promise.race([
liveRequestQueue.get(),
timeoutPromise
]).catch(() => null);
if (!liveRequest) {
continue;
}
}
catch (error) {
continue;
}
// Duplicate the live request to all active streams
if (invocationContext.activeStreamingTools) {
for (const streamingTool of Object.values(invocationContext.activeStreamingTools)) {
if (streamingTool.stream) {
streamingTool.stream.send(liveRequest);
}
}
}
// Small delay to yield to event loop
await new Promise(resolve => setTimeout(resolve, 0));
if (liveRequest.close) {
await llmConnection.close();
return;
}
if (liveRequest.blob) {
// Cache audio data for transcription
if (!invocationContext.transcriptionCache) {
invocationContext.transcriptionCache = [];
}
// Store the transcription entry with metadata
invocationContext.transcriptionCache.push(new TranscriptionEntry_1.TranscriptionEntry({
audioData: liveRequest.blob,
metadata: { type: 'user' }
}));
// Create a proper Blob object
const blob = {
data: liveRequest.blob,
mimeType: 'audio/wav'
};
await llmConnection.sendRealtime(blob);
}
if (liveRequest.content) {
await llmConnection.sendContent(liveRequest.content);
}
}
}
catch (error) {
console.error('Error in send task:', error);
}
}
/**
* Receives data from the model and processes events.
*
* @param llmConnection The LLM connection
* @param eventId The event ID
* @param invocationContext The invocation context
* @param llmRequest The LLM request
* @returns An async generator of events
*/
async *_receiveFromModel(llmConnection, eventId, invocationContext, llmRequest) {
if (!invocationContext.liveRequestQueue) {
return;
}
try {
// Run until explicitly returned or an error occurs
const isReceiving = true;
while (isReceiving) {
for await (const llmResponse of llmConnection.receive()) {
const modelResponseEvent = new Event_1.Event({
id: Event_1.Event.newId(),
invocationId: invocationContext.invocationId,
author: invocationContext.agent.name,
content: llmResponse.content,
});
yield* this._postprocessLive(invocationContext, llmRequest, llmResponse, modelResponseEvent);
if (invocationContext.endInvocation) {
return;
}
// Give opportunity for other tasks to run
await new Promise(resolve => setTimeout(resolve, 0));
}
}
}
catch (error) {
// Handle connection errors - use a specific check similar to Python's ConnectionClosedOK
// The Python version uses 'except ConnectionClosedOK:' which is a normal connection closure
if (error?.name === 'ConnectionClosedOK' ||
error?.code === 'CONNECTION_CLOSED' ||
error?.message?.includes('connection closed')) {
// Normal connection close, just return
return;
}
console.error('Error in receive task:', error);
}
}
/**
* Runs the flow asynchronously.
*
* @param invocationContext The invocation context
* @returns An async generator of events
*/
async *runAsync(invocationContext) {
while (true) {
let lastEvent;
for await (const event of this._runOneStepAsync(invocationContext)) {
lastEvent = event;
yield event;
}
if (!lastEvent || lastEvent.isFinalResponse()) {
break;
}
}
}
/**
* Runs one step of the flow asynchronously.
*
* @param invocationContext The invocation context
* @returns An async generator of events
*/
async *_runOneStepAsync(invocationContext) {
// Create request object first - matching Python implementation
const llmRequest = new LlmRequest_1.LlmRequest();
// Preprocess before calling the LLM - yield any events from preprocessing
yield* this._preprocessAsync(invocationContext, llmRequest);
if (invocationContext.endInvocation) {
return;
}
// Create model response event after preprocessing, like in Python
const modelResponseEvent = new Event_1.Event({
id: Event_1.Event.newId(),
invocationId: invocationContext.invocationId,
author: invocationContext.agent.name,
branch: invocationContext.branch,
});
// Call LLM and immediately process each response (like in Python)
for await (const llmResponse of this._callLlmAsync(invocationContext, llmRequest, modelResponseEvent)) {
// Process each LLM response as it comes in
yield* this._postprocessAsync(invocationContext, llmRequest, llmResponse, modelResponseEvent);
}
}
/**
* Preprocesses the request before calling the LLM.
*
* @param invocationContext The invocation context
* @param llmRequest The LLM request
* @returns An async generator of events
*/
async *_preprocessAsync(invocationContext, llmRequest) {
const agent = invocationContext.agent;
// Log starting state
// Make sure the agent is an LlmAgent
if (!(agent instanceof LlmAgent_1.LlmAgent)) {
return;
}
// Run all request processors
for (const processor of this.requestProcessors) {
yield* processor.runAsync(invocationContext, llmRequest);
if (invocationContext.endInvocation) {
return;
}
}
// Run processors for tools
if (agent.canonicalTools) {
for (const tool of agent.canonicalTools) {
// Create a new ToolContext directly with the invocation context
// This matches the Python implementation: tool_context = ToolContext(invocation_context)
const toolContext = new ToolContext_1.ToolContext(invocationContext);
await tool.processLlmRequest({
toolContext,
llmRequest
});
}
}
}
/**
* Postprocesses the response after the LLM call.
*
* @param invocationContext The invocation context
* @param llmRequest The LLM request
* @param llmResponse The LLM response
* @param modelResponseEvent The model response event
* @returns An async generator of events
*/
async *_postprocessAsync(invocationContext, llmRequest, llmResponse, modelResponseEvent) {
// Process the response with response processors
yield* this._postprocessRunProcessorsAsync(invocationContext, llmResponse);
if (invocationContext.endInvocation) {
return;
}
// Skip the model response event if there is no content and no error code
// This is needed for the code executor to trigger another loop
if (!llmResponse.content && !llmResponse.errorCode && !llmResponse.interrupted) {
return;
}
// Generate the finalized model response event
const finalEvent = this._finalizeModelResponseEvent(llmRequest, llmResponse, modelResponseEvent);
// Yield the event first
yield finalEvent;
// Handle any function calls in the response
if (finalEvent && finalEvent.getFunctionCalls().length > 0) {
yield* this._postprocessHandleFunctionCallsAsync(invocationContext, finalEvent, llmRequest);
}
}
/**
* Postprocesses the response for the live API.
*
* @param invocationContext The invocation context
* @param llmRequest The LLM request
* @param llmResponse The LLM response
* @param modelResponseEvent The model response event
* @returns An async generator of events
*/
async *_postprocessLive(invocationContext, llmRequest, llmResponse, modelResponseEvent) {
// Process the response with response processors
yield* this._postprocessRunProcessorsAsync(invocationContext, llmResponse);
// Skip the model response event if there is no content and no error code and no turn_complete
if (!llmResponse.content && !llmResponse.errorCode &&
!llmResponse.interrupted && !llmResponse.turnComplete) {
return;
}
// Generate the finalized model response event
const finalEvent = this._finalizeModelResponseEvent(llmRequest, llmResponse, modelResponseEvent);
if (!finalEvent) {
return;
}
// For live mode, yield the event immediately
yield finalEvent;
// Handle function calls
if (finalEvent.getFunctionCalls().length > 0) {
try {
// Get tools dictionary if available
const toolsDict = 'getToolsDict' in llmRequest ? llmRequest.getToolsDict() : {};
const functionResponseEvent = await functions.handleFunctionCallsLive(invocationContext, finalEvent, toolsDict);
if (functionResponseEvent) {
// Check for auth event
const authEvent = functions.generateAuthEvent(invocationContext, functionResponseEvent);
if (authEvent) {
yield authEvent;
}
yield functionResponseEvent;
// Check for transfer_to_agent
const transferToAgent = functionResponseEvent.actions?.transferToAgent;
if (transferToAgent && typeof invocationContext.agent.runLive === 'function') {
const agentToRun = this._getAgentToRun(invocationContext, transferToAgent);
if (typeof agentToRun.runLive === 'function') {
yield* agentToRun.runLive(invocationContext);
}
}
}
}
catch (error) {
console.error('Error handling function calls live:', error);
}
}
}
/**
* Runs the response processors.
*
* @param invocationContext The invocation context
* @param llmResponse The LLM response
* @returns An async generator of events
*/
async *_postprocessRunProcessorsAsync(invocationContext, llmResponse) {
for (const processor of this.responseProcessors) {
yield* processor.runAsync(invocationContext, llmResponse);
if (invocationContext.endInvocation) {
return;
}
}
}
/**
* Handles function calls in the response.
*
* @param invocationContext The invocation context
* @param functionCallEvent The function call event
* @param llmRequest The LLM request
* @returns An async generator of events
*/
async *_postprocessHandleFunctionCallsAsync(invocationContext, functionCallEvent, llmRequest) {
try {
// Handle function calls asynchronously
// Get tools dictionary if available
const toolsDict = 'getToolsDict' in llmRequest ? llmRequest.getToolsDict() : {};
const functionResponseEvent = await functions.handleFunctionCallsAsync(invocationContext, functionCallEvent, toolsDict);
if (functionResponseEvent) {
// Check for auth event
const authEvent = functions.generateAuthEvent(invocationContext, functionResponseEvent);
if (authEvent) {
yield authEvent;
}
yield functionResponseEvent;
// Check for transfer_to_agent
const transferToAgent = functionResponseEvent.actions?.transferToAgent;
if (transferToAgent && typeof invocationContext.agent.runAsync === 'function') {
const agentToRun = this._getAgentToRun(invocationContext, transferToAgent);
if (typeof agentToRun.runAsync === 'function') {
yield* agentToRun.runAsync(invocationContext);
return;
}
}
// If not transferring to an agent, continue the flow with another step
// This matches Python implementation of recursively continuing the flow
yield* this._runOneStepAsync(invocationContext);
}
}
catch (error) {
console.error('Error handling function calls:', error);
}
}
/**
* Gets the agent to run for a transfer_to_agent function call.
*
* @param invocationContext The invocation context
* @param transferToAgent The agent to transfer to
* @returns The agent to run
*/
_getAgentToRun(invocationContext, transferToAgent) {
// Check if rootAgent is available
const rootAgent = invocationContext.agent.rootAgent;
if (rootAgent) {
// Check if findAgent method exists
if (typeof rootAgent.findAgent === 'function') {
const agentToRun = rootAgent.findAgent(transferToAgent);
if (agentToRun) {
return agentToRun;
}
}
}
// Fallback: try to get from session
if (transferToAgent && typeof transferToAgent === 'object' && 'agent_name' in transferToAgent) {
const agentName = transferToAgent.agent_name;
if (invocationContext.session && typeof invocationContext.session.getAgent === 'function') {
const agent = invocationContext.session.getAgent(agentName);
if (agent) {
return agent;
}
}
}
throw new Error(`Agent ${transferToAgent} not found`);
}
/**
* Calls the LLM asynchronously.
*
* @param invocationContext The invocation context
* @param llmRequest The LLM request
* @param modelResponseEvent The model response event
* @returns An async generator of LLM responses
*/
async *_callLlmAsync(invocationContext, llmRequest, modelResponseEvent) {
// Run before_model_callback if it exists
const callbackResponse = this._handleBeforeModelCallback(invocationContext, llmRequest, modelResponseEvent);
if (callbackResponse) {
yield callbackResponse;
return;
}
const llm = this._getLlm(invocationContext);
// Start tracing span for LLM call
const tracingSpan = telemetry.tracer.startAsCurrentSpan('call_llm');
try {
if (invocationContext.runConfig?.supportCfc) {
invocationContext.liveRequestQueue = invocationContext.liveRequestQueue || new LiveRequestQueue_1.LiveRequestQueue();
for await (const llmResponse of this.runLive(invocationContext)) {
// Run after_model_callback if it exists
const alteredLlmResponse = this._handleAfterModelCallback(invocationContext, llmResponse, modelResponseEvent);
// Only yield partial response in SSE streaming mode
if (!invocationContext.runConfig?.streamingMode ||
invocationContext.runConfig.streamingMode === RunConfig_1.StreamingMode.SSE ||
!llmResponse.partial) {
yield alteredLlmResponse || llmResponse;
}
if (llmResponse.turnComplete) {
invocationContext.liveRequestQueue.close();
}
}
}
else {
// Check if we can make this llm call
// If the current call pushes the counter beyond the max set value,
// then the execution is stopped right here
if (typeof invocationContext.incrementLlmCallCount === 'function') {
invocationContext.incrementLlmCallCount();
}
for await (const llmResponse of llm.generateContentAsync(llmRequest, invocationContext.runConfig?.streamingMode === RunConfig_1.StreamingMode.SSE)) {
// Trace LLM call
telemetry.traceCallLlm(invocationContext, modelResponseEvent.id, llmRequest, llmResponse);
// Run after_model_callback if it exists
const alteredLlmResponse = this._handleAfterModelCallback(invocationContext, llmResponse, modelResponseEvent);
yield alteredLlmResponse || llmResponse;
}
}
}
finally {
// End tracing span
tracingSpan.end();
}
}
/**
* Handles the before model callback.
*
* @param invocationContext The invocation context
* @param llmRequest The LLM request
* @param modelResponseEvent The model response event
* @returns The callback response or undefined
*/
_handleBeforeModelCallback(invocationContext, llmRequest, modelResponseEvent) {
const agent = invocationContext.agent;
if (!(agent instanceof LlmAgent_1.LlmAgent) || !agent.beforeModelCallback) {
return undefined;
}
const callbackContext = new CallbackContext_1.CallbackContext(invocationContext, modelResponseEvent.actions);
return agent.beforeModelCallback(callbackContext, llmRequest);
}
/**
* Handles the after model callback.
*
* @param invocationContext The invocation context
* @param llmResponse The LLM response
* @param modelResponseEvent The model response event
* @returns The altered LLM response or undefined
*/
_handleAfterModelCallback(invocationContext, llmResponse, modelResponseEvent) {
const agent = invocationContext.agent;
if (!(agent instanceof LlmAgent_1.LlmAgent) || !agent.afterModelCallback) {
return undefined;
}
const callbackContext = new CallbackContext_1.CallbackContext(invocationContext, modelResponseEvent.actions);
return agent.afterModelCallback(callbackContext, llmResponse);
}
/**
* Finalizes the model response event.
*
* @param llmRequest The LLM request
* @param llmResponse The LLM response
* @param modelResponseEvent The model response event
* @returns The finalized event
*/
_finalizeModelResponseEvent(llmRequest, llmResponse, modelResponseEvent) {
// In Python, this is done via model_validate which merges properties
// Let's first ensure content parts are properly filtered, similar to Python
if (llmResponse.content && llmResponse.content.parts) {
llmResponse.content.parts = llmResponse.content.parts.filter(part => {
// Keep parts with valid text
if (part.text !== undefined && part.text !== null) {
return true;
}
// Keep parts with valid function calls
if (part.functionCall && part.functionCall.name) {
return true;
}
// Keep parts with valid function responses
if (part.functionResponse && part.functionResponse.name) {
return true;
}
// If we reached here, this part doesn't have valid required fields
return false;
});
}
// Create a new event with properties from both sources
// This simulates Python's model_validate approach
const finalEvent = new Event_1.Event({
...modelResponseEvent, // Spread existing event properties
// Add properties from llmResponse that aren't undefined
content: llmResponse.content || modelResponseEvent.content,
partial: llmResponse.partial !== undefined ? llmResponse.partial : modelResponseEvent.partial,
turnComplete: llmResponse.turnComplete !== undefined ? llmResponse.turnComplete : modelResponseEvent.turnComplete,
errorCode: llmResponse.errorCode || modelResponseEvent.errorCode,
errorMessage: llmResponse.errorMessage || modelResponseEvent.errorMessage,
interrupted: llmResponse.interrupted !== undefined ? llmResponse.interrupted : modelResponseEvent.interrupted,
customMetadata: llmResponse.customMetadata || modelResponseEvent.customMetadata
});
// Process function calls if present
if (finalEvent.content) {
const functionCalls = finalEvent.getFunctionCalls();
if (functionCalls.length > 0) {
functions.populateClientFunctionCallId(finalEvent);
// Get tools dictionary if available
const toolsDict = 'getToolsDict' in llmRequest ? llmRequest.getToolsDict() : {};
finalEvent.longRunningToolIds = functions.getLongRunningFunctionCalls(functionCalls, toolsDict);
}
}
return finalEvent;
}
/**
* Gets the LLM from the invocation context.
*
* @param invocationContext The invocation context
* @returns The LLM
*/
_getLlm(invocationContext) {
// First try to get LLM from invocation context
if (invocationContext.llm) {
return invocationContext.llm;
}
// If not in context, get from agent (matching Python implementation)
const agent = invocationContext.agent;
if (agent instanceof LlmAgent_1.LlmAgent) {
// Python directly accesses agent.canonical_model without extra checks
const llm = agent.canonicalModel;
if (llm) {
// Cache in context for future use
invocationContext.llm = llm;
return llm;
}
}
throw new Error('LLM not found in invocation context or agent');
}
}
exports.BaseLlmFlow = BaseLlmFlow;