@aigne/core 1.18.6 → 1.20.1

package/lib/esm/aigne/context.js

@@ -2,13 +2,13 @@ import equal from "fast-deep-equal";
 import { Emitter } from "strict-event-emitter";
 import { v7 } from "uuid";
 import { z } from "zod";
-import { Agent, isEmptyChunk, } from "../agents/agent.js";
+import { Agent, isAgentResponseDelta, isEmptyChunk, } 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 { AgentResponseProgressStream } from "../utils/event-stream.js";
 import { promiseWithResolvers } from "../utils/promise.js";
-import { agentResponseStreamToObject, asyncGeneratorToReadableStream, onAgentResponseStreamEnd, } from "../utils/stream-utils.js";
-import { checkArguments, isEmpty, isNil, omitBy, } from "../utils/type-utils.js";
+import { agentResponseStreamToObject, asyncGeneratorToReadableStream, mergeReadableStreams, onAgentResponseStreamEnd, } from "../utils/stream-utils.js";
+import { checkArguments, isEmpty, isNil, omit, } from "../utils/type-utils.js";
 import { MessageQueue, toMessagePayload, } from "./message-queue.js";
 import { newEmptyContextUsage } from "./usage.js";
 /**
@@ -80,38 +80,54 @@ export class AIGNEContext {
             });
         }
         const newContext = this.newContext();
-        const msg = createMessage(message);
-        return Promise.resolve(newContext.internal.invoke(agent, msg, newContext, options)).then(async (response) => {
+        return Promise.resolve(newContext.internal.invoke(agent, message, newContext, options)).then(async (response) => {
             if (!options?.streaming) {
-                const { __activeAgent__: activeAgent, ...output } = await agentResponseStreamToObject(response);
+                let { __activeAgent__: activeAgent, ...output } = await agentResponseStreamToObject(response);
+                output = await this.onInvocationResult(output, options);
                 if (options?.returnActiveAgent) {
                     return [output, activeAgent];
                 }
                 return output;
             }
             const activeAgentPromise = promiseWithResolvers();
-            const stream = onAgentResponseStreamEnd(asyncGeneratorToReadableStream(response), async ({ __activeAgent__: activeAgent }) => {
-                activeAgentPromise.resolve(activeAgent);
-            }, {
-                processChunk(chunk) {
-                    if (chunk.delta.json) {
+            const stream = onAgentResponseStreamEnd(asyncGeneratorToReadableStream(response), {
+                onChunk(chunk) {
+                    if (isAgentResponseDelta(chunk) && chunk.delta.json) {
                         return {
                             ...chunk,
                             delta: {
                                 ...chunk.delta,
-                                json: omitBy(chunk.delta.json, (_, k) => k === "__activeAgent__"),
+                                json: omit(chunk.delta.json, "__activeAgent__"),
                             },
                         };
                     }
-                    return chunk;
+                },
+                onResult: async (output) => {
+                    activeAgentPromise.resolve(output.__activeAgent__);
+                    return await this.onInvocationResult(output, options);
                 },
             });
+            const finalStream = !options.returnProgressChunks
+                ? stream
+                : mergeReadableStreams(stream, new AgentResponseProgressStream(newContext));
             if (options.returnActiveAgent) {
-                return [stream, activeAgentPromise.promise];
+                return [finalStream, activeAgentPromise.promise];
             }
-            return stream;
+            return finalStream;
         });
     });
+    async onInvocationResult(output, options) {
+        if (!options?.returnMetadata) {
+            return output;
+        }
+        return {
+            ...output,
+            $meta: {
+                ...output.$meta,
+                usage: this.usage,
+            },
+        };
+    }
     publish = ((topic, payload, options) => {
         if (options?.userContext) {
             Object.assign(this.userContext, options.userContext);
@@ -164,7 +180,6 @@ class AIGNEContextShared {
         this.memories = overrides?.memories ?? [];
     }
     messageQueue;
-    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
     events = new Emitter();
     get model() {
         return this.parent?.model;
@@ -198,44 +213,54 @@ class AIGNEContextShared {
         return withAbortSignal(this.abortController.signal, new Error("AIGNEContext is timeout"), () => this.invokeAgent(agent, input, context, options));
     }
     async *invokeAgent(agent, input, context, options) {
-        let activeAgent = agent;
-        for (;;) {
-            const result = {};
-            if (options?.sourceAgent && activeAgent !== options.sourceAgent) {
-                options.sourceAgent.hooks.onHandoff?.({
-                    context,
-                    source: options.sourceAgent,
-                    target: activeAgent,
-                    input,
-                });
-            }
-            const stream = await activeAgent.invoke(input, { ...options, context, streaming: true });
-            for await (const value of stream) {
-                if (value.delta.json) {
-                    value.delta.json = omitExistsProperties(result, value.delta.json);
-                    Object.assign(result, value.delta.json);
+        const startedAt = Date.now();
+        try {
+            let activeAgent = agent;
+            for (;;) {
+                const result = {};
+                if (options?.sourceAgent && activeAgent !== options.sourceAgent) {
+                    options.sourceAgent.hooks.onHandoff?.({
+                        context,
+                        source: options.sourceAgent,
+                        target: activeAgent,
+                        input,
+                    });
                 }
-                delete value.delta.json?.[transferAgentOutputKey];
-                if (isEmptyChunk(value))
-                    continue;
-                yield value;
-            }
-            if (!options?.disableTransfer) {
-                const transferToAgent = isTransferAgentOutput(result)
-                    ? result[transferAgentOutputKey].agent
-                    : undefined;
-                if (transferToAgent) {
-                    activeAgent = transferToAgent;
-                    continue;
+                const stream = await activeAgent.invoke(input, { ...options, context, streaming: true });
+                for await (const value of stream) {
+                    if (isAgentResponseDelta(value)) {
+                        if (value.delta.json) {
+                            value.delta.json = omitExistsProperties(result, value.delta.json);
+                            Object.assign(result, value.delta.json);
+                        }
+                        delete value.delta.json?.[transferAgentOutputKey];
+                    }
+                    if (isEmptyChunk(value))
+                        continue;
+                    yield value;
+                }
+                if (!options?.disableTransfer) {
+                    const transferToAgent = isTransferAgentOutput(result)
+                        ? result[transferAgentOutputKey].agent
+                        : undefined;
+                    if (transferToAgent) {
+                        activeAgent = transferToAgent;
+                        continue;
+                    }
                 }
+                break;
             }
-            break;
+            yield {
+                delta: {
+                    json: { __activeAgent__: activeAgent },
+                },
+            };
+        }
+        finally {
+            const endedAt = Date.now();
+            const duration = endedAt - startedAt;
+            this.usage.duration += duration;
         }
-        yield {
-            delta: {
-                json: { __activeAgent__: activeAgent },
-            },
-        };
     }
 }
 function omitExistsProperties(result, { ...delta }) {

package/lib/esm/aigne/message-queue.d.ts

@@ -21,7 +21,7 @@ export interface MessagePayload {
 /**
  * @hidden
  */
-export declare function toMessagePayload(payload: Omit<MessagePayload, "context"> | string | Message, options?: Partial<Pick<MessagePayload, "role" | "source">>): Omit<MessagePayload, "context">;
+export declare function toMessagePayload(payload: Omit<MessagePayload, "context"> | Message, options?: Partial<Pick<MessagePayload, "role" | "source">>): Omit<MessagePayload, "context">;
 /**
  * @hidden
  */

package/lib/esm/aigne/message-queue.js

@@ -1,6 +1,5 @@
 import { Emitter } from "strict-event-emitter";
 import { z } from "zod";
-import { createMessage } from "../prompt/prompt-builder.js";
 import { checkArguments, isNil, orArrayToArray } from "../utils/type-utils.js";
 /**
  * @hidden
@@ -24,12 +23,12 @@ function isMessagePayload(payload) {
  */
 export function toMessagePayload(payload, options) {
     if (isMessagePayload(payload)) {
-        return { ...payload, message: createMessage(payload.message), ...options };
+        return { ...payload, ...options };
     }
     return {
         role: options?.role || "user",
         source: options?.source,
-        message: createMessage(payload),
+        message: payload,
     };
 }
 /**

package/lib/esm/aigne/usage.d.ts

@@ -5,6 +5,7 @@ export interface ContextUsage {
     inputTokens: number;
     outputTokens: number;
     agentCalls: number;
+    duration: number;
 }
 /**
  * @hidden

package/lib/esm/aigne/usage.js

@@ -6,5 +6,6 @@ export function newEmptyContextUsage() {
         inputTokens: 0,
         outputTokens: 0,
         agentCalls: 0,
+        duration: 0,
     };
 }

package/lib/esm/loader/agent-yaml.d.ts

@@ -5,11 +5,11 @@ export declare function loadAgentFromYamlFile(path: string): Promise<{
     name: string;
     description?: string | undefined;
     skills?: string[] | undefined;
-    instructions?: string | undefined;
     memory?: true | {
         provider: string;
         subscribeTopic?: string[] | undefined;
     } | undefined;
+    instructions?: string | undefined;
     inputSchema?: ZodObject<Record<string, ZodType<any, z.ZodTypeDef, any>>, z.UnknownKeysParam, z.ZodTypeAny, {
         [x: string]: any;
     }, {
@@ -20,6 +20,7 @@ export declare function loadAgentFromYamlFile(path: string): Promise<{
     }, {
         [x: string]: any;
     }> | undefined;
+    inputKey?: string | undefined;
     outputKey?: string | undefined;
     toolChoice?: AIAgentToolChoice | undefined;
 } | {

package/lib/esm/loader/agent-yaml.js

@@ -18,6 +18,10 @@ const agentFileSchema = z.discriminatedUnion("type", [
             .string()
             .nullish()
             .transform((v) => v ?? undefined),
+        input_key: z
+            .string()
+            .nullish()
+            .transform((v) => v ?? undefined),
         input_schema: inputOutputSchema
             .nullish()
             .transform((v) => (v ? jsonSchemaToZod(v) : undefined)),

package/lib/esm/prompt/prompt-builder.d.ts

@@ -3,33 +3,19 @@ import { Agent, type AgentInvokeOptions, type Message } from "../agents/agent.js
 import type { AIAgent } from "../agents/ai-agent.js";
 import type { ChatModel, ChatModelInput } from "../agents/chat-model.js";
 import { ChatMessagesTemplate } from "./template.js";
-export declare const MESSAGE_KEY = "$message";
-export declare function createMessage<V extends Message>(message: string, variables?: V): {
-    [MESSAGE_KEY]: string;
-} & typeof variables;
-export declare function createMessage<I extends Message, V extends Message>(message: I, variables?: V): I & typeof variables;
-export declare function createMessage<I extends Message, V extends Message>(message: string | I, variables?: V): ({
-    [MESSAGE_KEY]: string;
-} | I) & typeof variables;
-export declare function getMessage(input: Message): string | undefined;
 export interface PromptBuilderOptions {
     instructions?: string | ChatMessagesTemplate;
 }
 export interface PromptBuildOptions extends Pick<AgentInvokeOptions, "context"> {
-    agent?: AIAgent;
+    agent?: AIAgent<any, any, any>;
     input?: Message;
     model?: ChatModel;
     outputSchema?: Agent["outputSchema"];
 }
 export declare class PromptBuilder {
-    static from(instructions: string): PromptBuilder;
-    static from(instructions: GetPromptResult): PromptBuilder;
-    static from(instructions: {
-        path: string;
-    }): Promise<PromptBuilder>;
     static from(instructions: string | {
         path: string;
-    } | GetPromptResult): PromptBuilder | Promise<PromptBuilder>;
+    } | GetPromptResult): PromptBuilder;
     private static fromFile;
     private static fromMCPPromptResult;
     constructor(options?: PromptBuilderOptions);

package/lib/esm/prompt/prompt-builder.js

@@ -3,23 +3,9 @@ import { stringify } from "yaml";
 import { ZodObject } from "zod";
 import { Agent } from "../agents/agent.js";
 import { outputSchemaToResponseFormatSchema } from "../utils/json-schema.js";
-import { isNil, unique } from "../utils/type-utils.js";
+import { unique } from "../utils/type-utils.js";
 import { MEMORY_MESSAGE_TEMPLATE } from "./prompts/memory-message-template.js";
 import { AgentMessageTemplate, ChatMessagesTemplate, PromptTemplate, SystemMessageTemplate, UserMessageTemplate, } from "./template.js";
-export const MESSAGE_KEY = "$message";
-export function createMessage(message, variables) {
-    return (typeof message === "string"
-        ? { [MESSAGE_KEY]: message, ...variables }
-        : { ...message, ...variables });
-}
-export function getMessage(input) {
-    const userInputMessage = input[MESSAGE_KEY];
-    if (typeof userInputMessage === "string")
-        return userInputMessage;
-    if (!isNil(userInputMessage))
-        return JSON.stringify(userInputMessage);
-    return undefined;
-}
 export class PromptBuilder {
     static from(instructions) {
         if (typeof instructions === "string")
@@ -30,8 +16,8 @@ export class PromptBuilder {
             return PromptBuilder.fromFile(instructions.path);
         throw new Error(`Invalid instructions ${instructions}`);
     }
-    static async fromFile(path) {
-        const text = await nodejs.fs.readFile(path, "utf-8");
+    static fromFile(path) {
+        const text = nodejs.fsSync.readFileSync(path, "utf-8");
         return PromptBuilder.from(text);
     }
     static fromMCPPromptResult(result) {
@@ -75,22 +61,25 @@ export class PromptBuilder {
     }
     async buildMessages(options) {
         const { input } = options;
+        const inputKey = options.agent?.inputKey;
+        const message = inputKey && typeof input?.[inputKey] === "string" ? input[inputKey] : undefined;
         const messages = (typeof this.instructions === "string"
             ? ChatMessagesTemplate.from([SystemMessageTemplate.from(this.instructions)])
             : this.instructions)?.format(options.input) ?? [];
         const memories = [];
-        if (options.agent) {
-            memories.push(...(await options.agent.retrieveMemories({ search: options.input }, options)));
+        if (options.agent?.inputKey) {
+            memories.push(...(await options.agent.retrieveMemories({ search: message }, options)));
         }
         if (options.context.memories?.length) {
             memories.push(...options.context.memories);
         }
         if (memories.length)
             messages.push(...this.convertMemoriesToMessages(memories, options));
-        const content = input && getMessage(input);
-        // add user input if it's not the same as the last message
-        if (content && messages.at(-1)?.content !== content) {
-            messages.push({ role: "user", content });
+        if (message) {
+            messages.push({
+                role: "user",
+                content: message,
+            });
         }
         return messages;
     }

package/lib/esm/utils/event-stream.d.ts

@@ -1,4 +1,5 @@
-import type { AgentResponseChunk, AgentResponseStream, Message } from "../agents/agent.js";
+import { type AgentResponseChunk, type AgentResponseProgress, type AgentResponseStream, type Message } from "../agents/agent.js";
+import type { Context } from "../aigne/context.js";
 export declare class EventStreamParser<T> extends TransformStream<string, T | Error> {
     constructor();
 }
@@ -9,3 +10,6 @@ export declare class AgentResponseStreamParser<O extends Message> extends Transf
 export declare class AgentResponseStreamSSE<O extends Message> extends ReadableStream<string> {
     constructor(stream: AgentResponseStream<O>);
 }
+export declare class AgentResponseProgressStream extends ReadableStream<AgentResponseProgress> {
+    constructor(context: Context);
+}

package/lib/esm/utils/event-stream.js

@@ -1,5 +1,6 @@
 import { createParser } from "eventsource-parser";
 import { produce } from "immer";
+import { isAgentResponseDelta, isAgentResponseProgress, } from "../agents/agent.js";
 import { tryOrThrow } from "./type-utils.js";
 export class EventStreamParser extends TransformStream {
     constructor() {
@@ -12,11 +13,16 @@ export class EventStreamParser extends TransformStream {
                             controller.enqueue(new Error(`Parse response chunk json error: ${e.message} ${event.data}`));
                         });
                         if (json) {
-                            if (event.event === "error") {
-                                controller.enqueue(new Error(json.message));
-                            }
-                            else {
-                                controller.enqueue(json);
+                            switch (event.event) {
+                                case "error":
+                                    controller.enqueue(new Error(json.message));
+                                    break;
+                                default: {
+                                    if (!event.event)
+                                        controller.enqueue(json);
+                                    else
+                                        console.warn(`Unknown event type: ${event.event}`, event.data);
+                                }
                             }
                         }
                     },
@@ -38,25 +44,35 @@ export class AgentResponseStreamParser extends TransformStream {
                     controller.terminate();
                     return;
                 }
-                this.json = produce(this.json, (draft) => {
-                    if (chunk.delta.json)
-                        Object.assign(draft, chunk.delta.json);
-                    if (chunk.delta.text) {
-                        for (const [key, text] of Object.entries(chunk.delta.text)) {
-                            const original = draft[key];
-                            const t = (original || "") + (text || "");
-                            if (t)
-                                Object.assign(draft, { [key]: t });
+                if (isAgentResponseDelta(chunk)) {
+                    this.json = produce(this.json, (draft) => {
+                        if (chunk.delta.json)
+                            Object.assign(draft, chunk.delta.json);
+                        if (chunk.delta.text) {
+                            for (const [key, text] of Object.entries(chunk.delta.text)) {
+                                const original = draft[key];
+                                const t = (original || "") + (text || "");
+                                if (t)
+                                    Object.assign(draft, { [key]: t });
+                            }
                         }
+                    });
+                    controller.enqueue({
+                        ...chunk,
+                        delta: {
+                            ...chunk.delta,
+                            json: this.json,
+                        },
+                    });
+                }
+                else if (isAgentResponseProgress(chunk)) {
+                    if (chunk.progress.event === "agentFailed") {
+                        const { name, message } = chunk.progress.error;
+                        chunk.progress.error = new Error(message);
+                        chunk.progress.error.name = name;
                     }
-                });
-                controller.enqueue({
-                    ...chunk,
-                    delta: {
-                        ...chunk.delta,
-                        json: this.json,
-                    },
-                });
+                    controller.enqueue(chunk);
+                }
             },
         });
     }
@@ -73,6 +89,14 @@ export class AgentResponseStreamSSE extends ReadableStream {
                         controller.close();
                         return;
                     }
+                    if (isAgentResponseProgress(value)) {
+                        if (value.progress.event === "agentFailed") {
+                            value.progress.error = {
+                                name: value.progress.error.name,
+                                message: value.progress.error.message,
+                            };
+                        }
+                    }
                     controller.enqueue(`data: ${JSON.stringify(value)}\n\n`);
                 }
                 catch (error) {
@@ -83,3 +107,43 @@ export class AgentResponseStreamSSE extends ReadableStream {
         });
     }
 }
+export class AgentResponseProgressStream extends ReadableStream {
+    constructor(context) {
+        super({
+            async start(controller) {
+                const writeEvent = (eventName, event) => {
+                    const progress = {
+                        ...event,
+                        event: eventName,
+                        agent: { name: event.agent.name },
+                    };
+                    controller.enqueue({ progress });
+                };
+                const close = () => {
+                    context.off("agentStarted", onAgentStarted);
+                    context.off("agentSucceed", onAgentSucceed);
+                    context.off("agentFailed", onAgentFailed);
+                    controller.close();
+                };
+                const onAgentStarted = (event) => {
+                    writeEvent("agentStarted", event);
+                };
+                const onAgentSucceed = (event) => {
+                    writeEvent("agentSucceed", event);
+                    if (event.contextId === context.id) {
+                        close();
+                    }
+                };
+                const onAgentFailed = (event) => {
+                    writeEvent("agentFailed", event);
+                    if (event.contextId === context.id) {
+                        close();
+                    }
+                };
+                context.on("agentStarted", onAgentStarted);
+                context.on("agentSucceed", onAgentSucceed);
+                context.on("agentFailed", onAgentFailed);
+            },
+        });
+    }
+}

package/lib/esm/utils/stream-utils.d.ts

@@ -1,15 +1,15 @@
 import { type AgentProcessAsyncGenerator, type AgentResponseChunk, type AgentResponseStream, type Message } from "../agents/agent.js";
-import type { MESSAGE_KEY } from "../prompt/prompt-builder.js";
 import { type PromiseOrValue } from "./type-utils.js";
 import "./stream-polyfill.js";
 export declare function objectToAgentResponseStream<T extends Message>(json: T): AgentResponseStream<T>;
 export declare function mergeAgentResponseChunk<T extends Message>(output: T, chunk: AgentResponseChunk<T>): T;
 export declare function agentResponseStreamToObject<T extends Message>(stream: AgentResponseStream<T> | AgentProcessAsyncGenerator<T>): Promise<T>;
 export declare function asyncGeneratorToReadableStream<T extends Message>(generator: AgentProcessAsyncGenerator<T>): AgentResponseStream<T>;
-export declare function onAgentResponseStreamEnd<T extends Message>(stream: AgentResponseStream<T>, callback: (result: T) => PromiseOrValue<Partial<T> | void>, options?: {
-    errorCallback?: (error: Error) => PromiseOrValue<Error>;
-    processChunk?: (chunk: AgentResponseChunk<T>) => AgentResponseChunk<T>;
-}): ReadableStream<any>;
+export declare function onAgentResponseStreamEnd<T extends Message>(stream: AgentResponseStream<T>, options?: {
+    onChunk?: (chunk: AgentResponseChunk<T>) => PromiseOrValue<AgentResponseChunk<T> | undefined | void>;
+    onResult?: (result: T) => PromiseOrValue<Partial<T> | undefined | void>;
+    onError?: (error: Error) => PromiseOrValue<Error>;
+}): AgentResponseStream<T>;
 export declare function isAsyncGenerator<T extends AsyncGenerator>(value: AsyncGenerator | unknown): value is T;
 export declare function arrayToAgentProcessAsyncGenerator<T extends Message>(chunks: (AgentResponseChunk<T> | Error)[], result?: Partial<T>): AgentProcessAsyncGenerator<T>;
 export declare function arrayToReadableStream<T>(chunks: (T | Error)[]): ReadableStream<T>;
@@ -17,8 +17,10 @@ export declare function readableStreamToArray<T>(stream: ReadableStream<T>, opti
     catchError: true;
 }): Promise<(T | Error)[]>;
 export declare function readableStreamToArray<T>(stream: ReadableStream<T>, options?: {
-    catchError?: false;
+    catchError?: boolean;
 }): Promise<T[]>;
-export declare function stringToAgentResponseStream(str: string, key?: "text" | typeof MESSAGE_KEY | string): AgentResponseStream<Message>;
+export declare function stringToAgentResponseStream(str: string, key?: "text" | string): AgentResponseStream<Message>;
 export declare function toReadableStream(stream: NodeJS.ReadStream): ReadableStream<Uint8Array<ArrayBufferLike>>;
 export declare function readAllString(stream: NodeJS.ReadStream | ReadableStream): Promise<string>;
+export declare function mergeReadableStreams<T1, T2>(s1: ReadableStream<T1>, s2: ReadableStream<T2>): ReadableStream<T1 | T2>;
+export declare function mergeReadableStreams(...streams: ReadableStream<any>[]): ReadableStream<any>;