@aigne/core 1.4.1-9 → 1.5.1-0

package/lib/esm/execution-engine/context.d.ts

@@ -1,11 +1,31 @@
-import type EventEmitter from "node:events";
-import type { Agent, FunctionAgentFn, Message } from "../agents/agent.js";
+import EventEmitter from "node:events";
+import { Agent, type FunctionAgentFn, type Message } from "../agents/agent.js";
+import { UserAgent } from "../agents/user-agent.js";
 import type { ChatModel } from "../models/chat-model.js";
-import type { MessagePayload, MessageQueueListener, Unsubscribe } from "./message-queue.js";
+import { type MessagePayload, MessageQueue, type MessageQueueListener, type Unsubscribe } from "./message-queue.js";
 export type Runnable<I extends Message = Message, O extends Message = Message> = Agent<I, O> | FunctionAgentFn;
+export interface ContextUsage {
+    promptTokens: number;
+    completionTokens: number;
+    agentCalls: number;
+}
+export interface ContextLimits {
+    maxTokens?: number;
+    maxAgentCalls?: number;
+    timeout?: number;
+}
 export interface Context extends EventEmitter {
     model?: ChatModel;
     tools?: Agent[];
+    usage: ContextUsage;
+    limits?: ContextLimits;
+    status?: "normal" | "timeout";
+    /**
+     * Create a user agent to consistently call an agent
+     * @param agent Agent to call
+     * @returns User agent
+     */
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>): UserAgent<I, O>;
     /**
      * Call an agent with a message
      * @param agent Agent to call
@@ -21,11 +41,11 @@ export interface Context extends EventEmitter {
      * @returns the output of the agent and the final active agent
      */
     call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options: {
-        returnActiveAgent?: true;
+        returnActiveAgent: true;
     }): Promise<[O, Runnable]>;
     call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options?: {
         returnActiveAgent?: boolean;
-    }): Promise<O | [O, Runnable]>;
+    }): UserAgent<I, O> | Promise<O | [O, Runnable]>;
     /**
      * Publish a message to a topic, the engine will call the listeners of the topic
      * @param topic topic name, or an array of topic names
@@ -38,3 +58,36 @@ export interface Context extends EventEmitter {
     subscribe(topic: string, listener?: MessageQueueListener): Unsubscribe | Promise<MessagePayload>;
     unsubscribe(topic: string, listener: MessageQueueListener): void;
 }
+export declare class ExecutionContext extends EventEmitter implements Context {
+    private readonly engine?;
+    constructor(engine?: {
+        model?: ChatModel;
+        tools?: Agent[];
+        limits?: ContextLimits;
+        messageQueue?: MessageQueue;
+        emit?: EventEmitter["emit"];
+    } | undefined);
+    private messageQueue;
+    get model(): ChatModel | undefined;
+    get tools(): Agent<Message, Message>[] | undefined;
+    usage: ContextUsage;
+    limits?: ContextLimits | undefined;
+    private abortController;
+    private timer?;
+    private initTimeout;
+    get status(): "normal" | "timeout";
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>): UserAgent<I, O>;
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string): Promise<O>;
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options: {
+        returnActiveAgent: true;
+    }): Promise<[O, Runnable]>;
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message?: I | string, options?: {
+        returnActiveAgent?: boolean;
+    }): UserAgent<I, O> | Promise<O | [O, Runnable]>;
+    publish(topic: string | string[], message: Message | string, from?: Agent): void;
+    subscribe(topic: string, listener?: undefined): Promise<MessagePayload>;
+    subscribe(topic: string, listener: MessageQueueListener): Unsubscribe;
+    unsubscribe(topic: string, listener: MessageQueueListener): void;
+    emit(eventName: string | symbol, ...args: unknown[]): boolean;
+    private callAgent;
+}

package/lib/esm/execution-engine/context.js

@@ -1 +1,164 @@
-export {};
+import EventEmitter from "node:events";
+import { z } from "zod";
+import { Agent } from "../agents/agent.js";
+import { isTransferAgentOutput, transferAgentOutputKey } from "../agents/types.js";
+import { UserAgent } from "../agents/user-agent.js";
+import { createMessage } from "../prompt/prompt-builder.js";
+import { checkArguments, isNil } from "../utils/type-utils.js";
+import { MessageQueue, } from "./message-queue.js";
+export class ExecutionContext extends EventEmitter {
+    engine;
+    constructor(engine) {
+        super();
+        this.engine = engine;
+        this.limits = engine?.limits;
+        this.messageQueue = engine?.messageQueue ?? new MessageQueue();
+    }
+    messageQueue;
+    get model() {
+        return this.engine?.model;
+    }
+    get tools() {
+        return this.engine?.tools;
+    }
+    usage = {
+        promptTokens: 0,
+        completionTokens: 0,
+        agentCalls: 0,
+    };
+    limits;
+    abortController = new AbortController();
+    timer;
+    initTimeout() {
+        if (this.timer)
+            return;
+        const timeout = this.limits?.timeout;
+        if (timeout) {
+            this.timer = setTimeout(() => {
+                this.abortController.abort();
+            }, timeout);
+        }
+    }
+    get status() {
+        return this.abortController.signal.aborted ? "timeout" : "normal";
+    }
+    call(agent, message, options) {
+        checkArguments("ExecutionEngine._call", executionEngineCallArgsSchema, {
+            agent,
+            message,
+            options,
+        });
+        if (isNil(message)) {
+            let activeAgent = agent;
+            return UserAgent.from({
+                context: this,
+                process: async (input, context) => {
+                    const [output, agent] = await context.call(activeAgent, input, {
+                        returnActiveAgent: true,
+                    });
+                    activeAgent = agent;
+                    return output;
+                },
+            });
+        }
+        return withAbortSignal(this.abortController.signal, new Error("ExecutionEngine is timeout"), async () => {
+            this.initTimeout();
+            const msg = createMessageFromInput(message);
+            return this.callAgent(agent, msg).then(async ({ output, agent: activeAgent }) => {
+                if (activeAgent instanceof Agent) {
+                    const publishTopics = typeof activeAgent.publishTopic === "function"
+                        ? await activeAgent.publishTopic(output)
+                        : activeAgent.publishTopic;
+                    if (publishTopics?.length) {
+                        this.publish(publishTopics, output, activeAgent);
+                    }
+                }
+                if (options?.returnActiveAgent) {
+                    return [output, activeAgent];
+                }
+                return output;
+            });
+        });
+    }
+    publish(topic, message, from) {
+        checkArguments("ExecutionEngineContext.publish", executionEnginePublishArgsSchema, {
+            topic,
+            message,
+            from,
+        });
+        const request = {
+            role: !from || from instanceof UserAgent ? "user" : "agent",
+            source: from?.name,
+            message: createMessageFromInput(message),
+            context: this,
+        };
+        this.messageQueue.publish(topic, request);
+    }
+    subscribe(topic, listener) {
+        return this.messageQueue.subscribe(topic, listener);
+    }
+    unsubscribe(topic, listener) {
+        this.messageQueue.unsubscribe(topic, listener);
+    }
+    emit(eventName, ...args) {
+        if (this.engine?.emit)
+            return this.engine.emit(eventName, ...args);
+        return super.emit(eventName, ...args);
+    }
+    async callAgent(agent, input) {
+        let activeAgent = agent;
+        let output;
+        for (;;) {
+            let result;
+            if (typeof activeAgent === "function") {
+                result = await activeAgent(input, this);
+            }
+            else {
+                result = await activeAgent.call(input, this);
+            }
+            if (!(result instanceof Agent))
+                output = result;
+            const transferToAgent = result instanceof Agent
+                ? result
+                : isTransferAgentOutput(result)
+                    ? result[transferAgentOutputKey].agent
+                    : undefined;
+            if (transferToAgent) {
+                activeAgent = transferToAgent;
+            }
+            else {
+                break;
+            }
+        }
+        if (!output)
+            throw new Error("Unexpected empty output");
+        return {
+            agent: activeAgent,
+            output,
+        };
+    }
+}
+function withAbortSignal(signal, error, fn) {
+    return new Promise((resolve, reject) => {
+        const listener = () => reject(error);
+        signal.addEventListener("abort", listener);
+        fn()
+            .then(resolve, reject)
+            .finally(() => {
+            signal.removeEventListener("abort", listener);
+        });
+    });
+}
+function createMessageFromInput(message) {
+    return typeof message === "string" ? createMessage(message) : message;
+}
+const executionEnginePublishArgsSchema = z.object({
+    topic: z.union([z.string(), z.array(z.string())]),
+    message: z.union([z.string(), z.record(z.unknown())]),
+    from: z.instanceof(Agent).optional(),
+});
+const executionEngineCallArgsSchema = z.object({
+    agent: z.union([z.function(), z.instanceof(Agent)]),
+    message: z.union([z.record(z.unknown()), z.string()]).optional(),
+    options: z.object({ returnActiveAgent: z.boolean().optional() }).optional(),
+});

package/lib/esm/execution-engine/execution-engine.d.ts

@@ -1,22 +1,24 @@
 import EventEmitter from "node:events";
 import { Agent, type Message } from "../agents/agent.js";
-import { UserAgent } from "../agents/user-agent.js";
+import type { UserAgent } from "../agents/user-agent.js";
 import { ChatModel } from "../models/chat-model.js";
-import type { Context, Runnable } from "./context.js";
+import { type ContextLimits, ExecutionContext, type Runnable } from "./context.js";
 import { type MessagePayload, MessageQueue, type MessageQueueListener, type Unsubscribe } from "./message-queue.js";
 export interface ExecutionEngineOptions {
     model?: ChatModel;
     tools?: Agent[];
     agents?: Agent[];
+    limits?: ContextLimits;
 }
-export declare class ExecutionEngine extends EventEmitter implements Context {
+export declare class ExecutionEngine extends EventEmitter {
     constructor(options?: ExecutionEngineOptions);
     readonly messageQueue: MessageQueue;
     model?: ChatModel;
     tools: Agent[];
     private agents;
+    limits?: ContextLimits;
     addAgent(...agents: Agent[]): void;
-    private createMessageFromInput;
+    newContext(): ExecutionContext;
     /**
      * Publish a message to a topic, the engine will call the listeners of the topic
      * @param topic topic name, or an array of topic names
@@ -45,20 +47,15 @@ export declare class ExecutionEngine extends EventEmitter implements Context {
      * @returns the output of the agent and the final active agent
      */
     call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options: {
-        returnActiveAgent?: true;
+        returnActiveAgent: true;
     }): Promise<[O, Runnable]>;
-    /**
-     * Call an agent with a message and return the output and the active agent
-     * @param agent Agent to call
-     * @param message Message to pass to the agent
-     * @returns the output of the agent and the final active agent
-     */
-    private _call;
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message?: I | string, options?: {
+        returnActiveAgent?: boolean;
+    }): UserAgent<I, O> | Promise<O | [O, Runnable]>;
     subscribe(topic: string, listener?: undefined): Promise<MessagePayload>;
     subscribe(topic: string, listener: MessageQueueListener): Unsubscribe;
     subscribe(topic: string, listener?: MessageQueueListener): Unsubscribe | Promise<MessagePayload>;
     unsubscribe(topic: string, listener: MessageQueueListener): void;
-    private callAgent;
     shutdown(): Promise<void>;
     private initProcessExitHandler;
 }

package/lib/esm/execution-engine/execution-engine.js

@@ -1,12 +1,9 @@
 import EventEmitter from "node:events";
-import { isNil } from "lodash-es";
 import { z } from "zod";
 import { Agent } from "../agents/agent.js";
-import { isTransferAgentOutput, transferAgentOutputKey } from "../agents/types.js";
-import { UserAgent } from "../agents/user-agent.js";
 import { ChatModel } from "../models/chat-model.js";
-import { createMessage } from "../prompt/prompt-builder.js";
 import { checkArguments } from "../utils/type-utils.js";
+import { ExecutionContext } from "./context.js";
 import { MessageQueue, } from "./message-queue.js";
 export class ExecutionEngine extends EventEmitter {
     constructor(options) {
@@ -15,6 +12,7 @@ export class ExecutionEngine extends EventEmitter {
         super();
         this.model = options?.model;
         this.tools = options?.tools ?? [];
+        this.limits = options?.limits;
         if (options?.agents?.length)
             this.addAgent(...options.agents);
         this.initProcessExitHandler();
@@ -23,6 +21,7 @@ export class ExecutionEngine extends EventEmitter {
     model;
     tools;
     agents = [];
+    limits;
     addAgent(...agents) {
         checkArguments("ExecutionEngine.addAgent", executionEngineAddAgentArgsSchema, agents);
         for (const agent of agents) {
@@ -30,8 +29,8 @@ export class ExecutionEngine extends EventEmitter {
             agent.attach(this);
         }
     }
-    createMessageFromInput(message) {
-        return typeof message === "string" ? createMessage(message) : message;
+    newContext() {
+        return new ExecutionContext(this);
     }
     /**
      * Publish a message to a topic, the engine will call the listeners of the topic
@@ -40,56 +39,10 @@ export class ExecutionEngine extends EventEmitter {
      * @param from the agent who publish the message, if not provided, it will be treated as a user message
      */
     publish(topic, message, from) {
-        checkArguments("ExecutionEngine.publish", executionEnginePublishArgsSchema, {
-            topic,
-            message,
-            from,
-        });
-        const request = {
-            role: !from || from instanceof UserAgent ? "user" : "agent",
-            source: from?.name,
-            message: this.createMessageFromInput(message),
-        };
-        this.messageQueue.publish(topic, request);
+        return new ExecutionContext(this).publish(topic, message, from);
     }
     call(agent, message, options) {
-        checkArguments("ExecutionEngine.call", executionEngineCallArgsSchema, {
-            agent,
-            message,
-            options,
-        });
-        if (isNil(message)) {
-            let activeAgent = agent;
-            return UserAgent.from({
-                context: this,
-                process: async (input, context) => {
-                    const [output, agent] = await context.call(activeAgent, input, {
-                        returnActiveAgent: true,
-                    });
-                    activeAgent = agent;
-                    return output;
-                },
-            });
-        }
-        return this._call(agent, this.createMessageFromInput(message)).then((result) => options?.returnActiveAgent ? [result.output, result.agent] : result.output);
-    }
-    /**
-     * Call an agent with a message and return the output and the active agent
-     * @param agent Agent to call
-     * @param message Message to pass to the agent
-     * @returns the output of the agent and the final active agent
-     */
-    async _call(agent, message) {
-        const { output, agent: activeAgent } = await this.callAgent(agent, message);
-        if (activeAgent instanceof Agent) {
-            const publishTopics = typeof activeAgent.publishTopic === "function"
-                ? await activeAgent.publishTopic(output)
-                : activeAgent.publishTopic;
-            if (publishTopics?.length) {
-                this.publish(publishTopics, output, activeAgent);
-            }
-        }
-        return { output, agent: activeAgent };
+        return new ExecutionContext(this).call(agent, message, options);
     }
     subscribe(topic, listener) {
         checkArguments("ExecutionEngine.subscribe", executionEngineSubscribeArgsSchema, {
@@ -105,38 +58,6 @@ export class ExecutionEngine extends EventEmitter {
         });
         this.messageQueue.unsubscribe(topic, listener);
     }
-    async callAgent(agent, input) {
-        let activeAgent = agent;
-        let output;
-        for (;;) {
-            let result;
-            if (typeof activeAgent === "function") {
-                result = await activeAgent(input, this);
-            }
-            else {
-                result = await activeAgent.call(input, this);
-            }
-            if (!(result instanceof Agent))
-                output = result;
-            const transferToAgent = result instanceof Agent
-                ? result
-                : isTransferAgentOutput(result)
-                    ? result[transferAgentOutputKey].agent
-                    : undefined;
-            if (transferToAgent) {
-                activeAgent = transferToAgent;
-            }
-            else {
-                break;
-            }
-        }
-        if (!output)
-            throw new Error("Unexpected empty output");
-        return {
-            agent: activeAgent,
-            output,
-        };
-    }
     async shutdown() {
         for (const tool of this.tools) {
             await tool.shutdown();
@@ -157,16 +78,6 @@ const executionEngineOptionsSchema = z.object({
     agents: z.array(z.instanceof(Agent)).optional(),
 });
 const executionEngineAddAgentArgsSchema = z.array(z.instanceof(Agent));
-const executionEnginePublishArgsSchema = z.object({
-    topic: z.union([z.string(), z.array(z.string())]),
-    message: z.union([z.string(), z.record(z.unknown())]),
-    from: z.instanceof(Agent).optional(),
-});
-const executionEngineCallArgsSchema = z.object({
-    agent: z.union([z.function(), z.instanceof(Agent)]),
-    message: z.union([z.record(z.unknown()), z.string()]).optional(),
-    options: z.object({ returnActiveAgent: z.boolean().optional() }).optional(),
-});
 const executionEngineSubscribeArgsSchema = z.object({
     topic: z.string(),
     listener: z.function(z.tuple([z.any()]), z.any()).optional(),

package/lib/esm/execution-engine/message-queue.d.ts

@@ -1,10 +1,12 @@
 import type { Message } from "../agents/agent.js";
+import type { Context } from "./context.js";
 export declare const UserInputTopic = "UserInputTopic";
 export declare const UserOutputTopic = "UserOutputTopic";
 export interface MessagePayload {
     role: "user" | "agent";
     source?: string;
     message: Message;
+    context: Context;
 }
 export type MessageQueueListener = (message: MessagePayload) => void;
 export type MessageRequest = MessagePayload;

package/lib/esm/execution-engine/utils.d.ts

@@ -1,4 +1,4 @@
 import { type FunctionAgentFn } from "../agents/agent.js";
 import type { Runnable } from "./context.js";
-export declare function sequential(..._agents: [Runnable, ...Runnable[]]): FunctionAgentFn;
-export declare function parallel(..._agents: [Runnable, ...Runnable[]]): FunctionAgentFn;
+export declare function sequential(...agents: [Runnable, ...Runnable[]]): FunctionAgentFn;
+export declare function parallel(...agents: [Runnable, ...Runnable[]]): FunctionAgentFn;

package/lib/esm/execution-engine/utils.js

@@ -1,32 +1,28 @@
 import { z } from "zod";
 import { Agent } from "../agents/agent.js";
 import { checkArguments } from "../utils/type-utils.js";
-export function sequential(..._agents) {
-    checkArguments("sequential", agentArraySchema, _agents);
-    let agents = [..._agents];
+export function sequential(...agents) {
+    checkArguments("sequential", agentArraySchema, agents);
+    let _agents = [...agents];
     return async (input, context) => {
-        if (!context)
-            throw new Error("Context is required for executing sequential agents. Please provide a valid context.");
         const output = {};
         // Clone the agents to run, so that we can update the agents list during the loop
-        const agentsToRun = [...agents];
-        agents = [];
+        const agentsToRun = [..._agents];
+        _agents = [];
         for (const agent of agentsToRun) {
             const [o, transferToAgent] = await context.call(agent, { ...input, ...output }, { returnActiveAgent: true });
             Object.assign(output, o);
-            agents.push(transferToAgent);
+            _agents.push(transferToAgent);
         }
         return output;
     };
 }
-export function parallel(..._agents) {
-    checkArguments("parallel", agentArraySchema, _agents);
-    let agents = [..._agents];
+export function parallel(...agents) {
+    checkArguments("parallel", agentArraySchema, agents);
+    let _agents = [...agents];
     return async (input, context) => {
-        if (!context)
-            throw new Error("Context is required for executing parallel agents. Please provide a valid context.");
-        const result = await Promise.all(agents.map((agent) => context.call(agent, input, { returnActiveAgent: true })));
-        agents = result.map((i) => i[1]);
+        const result = await Promise.all(_agents.map((agent) => context.call(agent, input, { returnActiveAgent: true })));
+        _agents = result.map((i) => i[1]);
         const outputs = result.map((i) => i[0]);
         return Object.assign({}, ...outputs);
     };

package/lib/esm/index.d.ts

@@ -6,12 +6,5 @@ export * from "./agents/types.js";
 export * from "./agents/user-agent.js";
 export * from "./execution-engine/index.js";
 export * from "./models/chat-model.js";
-export * from "./models/claude-chat-model.js";
-export * from "./models/openai-chat-model.js";
-export * from "./models/xai-chat-model.js";
 export * from "./prompt/prompt-builder.js";
 export * from "./prompt/template.js";
-export * from "./utils/json-schema.js";
-export * from "./utils/logger.js";
-export * from "./utils/run-chat-loop.js";
-export * from "./utils/type-utils.js";

package/lib/esm/index.js

@@ -6,12 +6,5 @@ export * from "./agents/types.js";
 export * from "./agents/user-agent.js";
 export * from "./execution-engine/index.js";
 export * from "./models/chat-model.js";
-export * from "./models/claude-chat-model.js";
-export * from "./models/openai-chat-model.js";
-export * from "./models/xai-chat-model.js";
 export * from "./prompt/prompt-builder.js";
 export * from "./prompt/template.js";
-export * from "./utils/json-schema.js";
-export * from "./utils/logger.js";
-export * from "./utils/run-chat-loop.js";
-export * from "./utils/type-utils.js";

package/lib/esm/models/chat-model.d.ts

@@ -1,6 +1,9 @@
 import { Agent, type Message } from "../agents/agent.js";
+import type { Context } from "../execution-engine/context.js";
 export declare abstract class ChatModel extends Agent<ChatModelInput, ChatModelOutput> {
     constructor();
+    protected preprocess(input: ChatModelInput, context: Context): void;
+    protected postprocess(input: ChatModelInput, output: ChatModelOutput, context: Context): void;
 }
 export interface ChatModelInput extends Message {
     messages: ChatModelInputMessage[];
@@ -71,6 +74,7 @@ export interface ChatModelOutput extends Message {
     text?: string;
     json?: object;
     toolCalls?: ChatModelOutputToolCall[];
+    usage?: ChatModelOutputUsage;
 }
 export interface ChatModelOutputToolCall {
     id: string;
@@ -80,3 +84,7 @@ export interface ChatModelOutputToolCall {
         arguments: Message;
     };
 }
+export interface ChatModelOutputUsage {
+    promptTokens: number;
+    completionTokens: number;
+}

package/lib/esm/models/chat-model.js

@@ -7,6 +7,22 @@ export class ChatModel extends Agent {
             outputSchema: chatModelOutputSchema,
         });
     }
+    preprocess(input, context) {
+        super.preprocess(input, context);
+        const { limits, usage } = context;
+        const usedTokens = usage.completionTokens + usage.promptTokens;
+        if (limits?.maxTokens && usedTokens >= limits.maxTokens) {
+            throw new Error(`Exceeded max tokens ${usedTokens}/${limits.maxTokens}`);
+        }
+    }
+    postprocess(input, output, context) {
+        super.postprocess(input, output, context);
+        const { usage } = output;
+        if (usage) {
+            context.usage.completionTokens += usage.completionTokens;
+            context.usage.promptTokens += usage.promptTokens;
+        }
+    }
 }
 const chatModelInputMessageSchema = z.object({
     role: z.union([z.literal("system"), z.literal("user"), z.literal("agent"), z.literal("tool")]),
@@ -81,8 +97,13 @@ const chatModelOutputToolCallSchema = z.object({
         arguments: z.record(z.unknown()),
     }),
 });
+const chatModelOutputUsageSchema = z.object({
+    promptTokens: z.number(),
+    completionTokens: z.number(),
+});
 const chatModelOutputSchema = z.object({
     text: z.string().optional(),
     json: z.record(z.unknown()).optional(),
     toolCalls: z.array(chatModelOutputToolCallSchema).optional(),
+    usage: chatModelOutputUsageSchema.optional(),
 });