adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
365 lines (364 loc) • 14.4 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.LlmAgent = void 0;
/**
* Implementation of an agent that uses an LLM flow.
*/
const BaseAgent_1 = require("./BaseAgent");
const InvocationContext_1 = require("./InvocationContext");
const SingleFlow_1 = require("../flows/llm_flows/SingleFlow");
const types_1 = require("../models/types");
const models_1 = require("../models");
const BaseTool_1 = require("../tools/BaseTool");
const FunctionTool_1 = require("../tools/FunctionTool");
const Session_1 = require("../sessions/Session");
const uuid_1 = require("uuid");
const State_1 = require("../sessions/State");
/**
* Convert a tool union to a BaseTool instance
*/
function convertToolUnionToTool(toolUnion) {
if (toolUnion instanceof BaseTool_1.BaseTool) {
return toolUnion;
}
else if (typeof toolUnion === 'function') {
return new FunctionTool_1.FunctionTool(toolUnion);
}
throw new Error('Invalid tool type');
}
/**
* An agent that uses an LLM flow to process requests.
*/
class LlmAgent extends BaseAgent_1.BaseAgent {
/**
* Creates a new LLM agent.
*
* @param options Options for the agent including name
*/
constructor(options) {
if (!options.name) {
throw new Error('Agent name is required');
}
super(options.name, options);
/** The LLM model used by this agent */
this.model = '';
/** The instruction template for the agent */
this.instruction = '';
/** Global instruction for all agents in the tree */
this.globalInstruction = '';
/** Tools available to this agent */
this.tools = [];
/** Whether to disallow transfers to the parent agent */
this.disallowTransferToParent = false;
/** Whether to disallow transfers to peer agents */
this.disallowTransferToPeers = false;
/** Include contents setting */
this.includeContents = 'default';
// Set properties from options
this.customFlow = options.flow;
this.model = options.model || '';
this.instruction = options.instruction || '';
this.globalInstruction = options.globalInstruction || '';
this.tools = options.tools || [];
this.generateContentConfig = options.generateContentConfig;
this.disallowTransferToParent = options.disallowTransferToParent || false;
this.disallowTransferToPeers = options.disallowTransferToPeers || false;
this.includeContents = options.includeContents || 'default';
this.inputSchema = options.inputSchema;
this.outputSchema = options.outputSchema;
this.outputKey = options.outputKey;
this.planner = options.planner;
this.codeExecutor = options.codeExecutor;
this.examples = options.examples;
this.beforeModelCallback = options.beforeModelCallback;
this.afterModelCallback = options.afterModelCallback;
this.beforeToolCallback = options.beforeToolCallback;
this.afterToolCallback = options.afterToolCallback;
// Validate output schema configuration
this.validateOutputSchema();
}
/**
* Creates a new session for this agent
*
* @returns A promise resolving to a new Session object
*/
async createSession(options = {}) {
const messages = [];
// Create a session with the Session class
const session = new Session_1.Session({
id: options.id,
appName: options.appName || 'default-app',
userId: options.userId || 'default-user',
state: options.state || new State_1.State(),
events: options.events || []
});
// Add this agent to the session's agents map
session.agents.set(this.name, this);
// Extend the session with our custom methods for message handling
const extendedSession = session;
// Add the sendMessage method
extendedSession.sendMessage = async (message) => {
// Convert string to Message if needed
const msgObj = typeof message === 'string'
? {
id: (0, uuid_1.v4)(),
role: types_1.MessageRole.USER,
parts: [{ text: message }],
timestamp: new Date(),
text: () => typeof message === 'string' ? message : JSON.stringify(message)
}
: message;
// Add to message history
messages.push(msgObj);
// Create an invocation context
const context = new InvocationContext_1.InvocationContext({
invocationId: (0, uuid_1.v4)(),
session: session,
agent: this,
userContent: {
role: types_1.MessageRole.USER,
parts: msgObj.parts
}
});
// Process the message using the agent
const events = [];
for await (const event of this.invoke(context)) {
events.push(event);
// Also add to session's events
session.events.push(event);
}
// Create a response from the final event
const finalEvent = events[events.length - 1];
if (!finalEvent || !finalEvent.content) {
throw new Error('No response generated by agent');
}
// Create response message
const responseMsg = {
id: (0, uuid_1.v4)(),
role: types_1.MessageRole.ASSISTANT,
parts: finalEvent.content.parts || [],
timestamp: new Date(),
text: () => {
if (!finalEvent.content || !finalEvent.content.parts)
return '';
return finalEvent.content.parts
.filter(part => part.text !== undefined)
.map(part => part.text)
.join('');
}
};
// Add to message history
messages.push(responseMsg);
return responseMsg;
};
// Add the getMessages method
extendedSession.getMessages = () => {
return [...messages];
};
return extendedSession;
}
/**
* Implementation of the agent's async invocation logic.
*
* @param invocationContext The invocation context
* @returns An async generator of events
*/
async *runAsyncImpl(invocationContext) {
// Forward to the LLM flow
for await (const event of this.llmFlow.runAsync(invocationContext)) {
this.maybeSaveOutputToState(event);
yield event;
}
}
/**
* Implementation of the agent's live invocation logic.
*
* @param invocationContext The invocation context
* @returns An async generator of events
*/
async *runLiveImpl(invocationContext) {
// Forward to the LLM flow
for await (const event of this.llmFlow.runLive(invocationContext)) {
this.maybeSaveOutputToState(event);
yield event;
}
if (invocationContext.endInvocation) {
return;
}
}
/**
* Sets the user content for the agent.
*
* @param content The user content
* @param invocationContext The invocation context
*/
setUserContent(content, invocationContext) {
invocationContext.userContent = content;
}
/**
* Gets the resolved model as a BaseLlm.
* This method is only for use by Agent Development Kit.
*/
get canonicalModel() {
if (typeof this.model !== 'string') {
return this.model;
}
else if (this.model) {
return models_1.LlmRegistry.newLlm(this.model);
}
else {
// Find model from ancestors
let ancestorAgent = this.parentAgent;
while (ancestorAgent !== undefined) {
if (ancestorAgent instanceof LlmAgent) {
return ancestorAgent.canonicalModel;
}
ancestorAgent = ancestorAgent.parentAgent;
}
throw new Error(`No model found for ${this.name}`);
}
}
/**
* Gets the resolved instruction for this agent.
* This method is only for use by Agent Development Kit.
*/
canonicalInstruction(ctx) {
if (typeof this.instruction === 'string') {
return this.instruction;
}
else {
return this.instruction(ctx);
}
}
/**
* Gets the resolved global instruction.
* This method is only for use by Agent Development Kit.
*/
canonicalGlobalInstruction(ctx) {
if (typeof this.globalInstruction === 'string') {
return this.globalInstruction;
}
else {
return this.globalInstruction(ctx);
}
}
/**
* Gets the resolved tools as BaseTool instances.
* This method is only for use by Agent Development Kit.
*/
get canonicalTools() {
return this.tools.map(tool => convertToolUnionToTool(tool));
}
/**
* Returns the LLM flow to use for this agent.
*/
get llmFlow() {
// Use custom flow if configured
if (this.customFlow) {
return this.customFlow;
}
// Default flow based on agent transfer settings
if (this.disallowTransferToParent &&
this.disallowTransferToPeers &&
this.subAgents.length === 0) {
return new SingleFlow_1.SingleFlow();
}
// Dynamic import of AutoFlow to break circular dependency
// eslint-disable-next-line @typescript-eslint/no-var-requires
const { AutoFlow } = require('../flows/llm_flows/AutoFlow');
return new AutoFlow();
}
/**
* Saves the model output to state if needed.
*/
maybeSaveOutputToState(event) {
if (this.outputKey &&
event.isFinalResponse() &&
event.content &&
event.content.parts) {
let result = event.content.parts
.filter(part => part.text !== undefined)
.map(part => part.text)
.join('');
if (this.outputSchema) {
try {
// Parse JSON and validate against schema
const parsed = JSON.parse(result);
// Validate against the schema if it's provided
// In TypeScript we don't have Pydantic's model_validate_json,
// so we do basic validation based on the schema type
if (typeof this.outputSchema === 'function') {
// Assuming outputSchema is a constructor function or class
try {
// Try to instantiate using the schema class/constructor
const validated = new this.outputSchema(parsed);
result = validated;
}
catch (validationError) {
console.warn(`Schema validation failed: ${validationError}`);
// Still use the parsed result, even if validation failed
result = parsed;
}
}
else if (typeof this.outputSchema === 'object') {
// Basic property validation if schema is an object with properties
const schemaProps = Object.keys(this.outputSchema.properties || {});
const requiredProps = this.outputSchema.required || [];
// Check required properties
for (const prop of requiredProps) {
if (parsed[prop] === undefined) {
console.warn(`Schema validation failed: missing required property '${prop}'`);
}
}
// Remove properties not in schema if strict
if (this.outputSchema.additionalProperties === false) {
const filteredResult = {};
for (const key of schemaProps) {
if (parsed[key] !== undefined) {
filteredResult[key] = parsed[key];
}
}
result = filteredResult;
}
else {
result = parsed;
}
}
else {
// If we can't determine schema type, just use parsed JSON
result = parsed;
}
}
catch (error) {
console.warn(`Failed to parse output as JSON: ${error}`);
}
}
// Update the event's state delta with the output
event.actions.stateDelta[this.outputKey] = result;
}
}
/**
* Validates the output schema configuration.
*/
validateOutputSchema() {
if (!this.outputSchema) {
return;
}
if (!this.disallowTransferToParent || !this.disallowTransferToPeers) {
console.warn(`Invalid config for agent ${this.name}: output_schema cannot co-exist with ` +
`agent transfer configurations. Setting ` +
`disallowTransferToParent=true, disallowTransferToPeers=true`);
this.disallowTransferToParent = true;
this.disallowTransferToPeers = true;
}
if (this.subAgents.length > 0) {
throw new Error(`Invalid config for agent ${this.name}: if outputSchema is set, ` +
`subAgents must be empty to disable agent transfer.`);
}
if (this.tools.length > 0) {
throw new Error(`Invalid config for agent ${this.name}: if outputSchema is set, ` +
`tools must be empty`);
}
}
}
exports.LlmAgent = LlmAgent;