@pikku/core 0.11.1 → 0.12.0

package/dist/wirings/ai-agent/ai-agent-stream.js

@@ -0,0 +1,503 @@
+import { pikkuState } from '../../pikku-state.js';
+import { combineChannelMiddleware, wrapChannelWithMiddleware, } from '../channel/channel-middleware-runner.js';
+import { randomUUID } from 'crypto';
+import { parseWorkingMemory, deepMergeWorkingMemory, resolveMemoryServices, loadContextMessages, trimMessages, } from './ai-agent-memory.js';
+import { prepareAgentRun, resolveAgent, buildInstructions, buildToolDefs, createScopedChannel, ToolApprovalRequired, } from './ai-agent-prepare.js';
+import { resolveModelConfig } from './ai-agent-model-config.js';
+function createPersistingChannel(parent, storage, threadId) {
+    let fullText = '';
+    let stepText = '';
+    let stepToolCalls = [];
+    let stepToolResults = [];
+    const flushStep = async () => {
+        if (!storage)
+            return;
+        const text = stepText;
+        const calls = stepToolCalls;
+        const results = stepToolResults;
+        stepText = '';
+        stepToolCalls = [];
+        stepToolResults = [];
+        if (text || calls.length > 0) {
+            const assistantMsg = {
+                id: randomUUID(),
+                role: 'assistant',
+                content: text || undefined,
+                toolCalls: calls.length > 0 ? calls : undefined,
+                createdAt: new Date(),
+            };
+            const messages = [assistantMsg];
+            if (results.length > 0) {
+                messages.push({
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: results,
+                    createdAt: new Date(),
+                });
+            }
+            await storage.saveMessages(threadId, messages);
+        }
+    };
+    const channel = {
+        channelId: parent.channelId,
+        openingData: parent.openingData,
+        get state() {
+            return parent.state;
+        },
+        get fullText() {
+            return fullText;
+        },
+        flush: flushStep,
+        close: () => parent.close(),
+        send: (event) => {
+            if (storage) {
+                switch (event.type) {
+                    case 'text-delta':
+                        stepText += event.text;
+                        fullText += event.text;
+                        break;
+                    case 'tool-call':
+                        stepToolCalls.push({
+                            id: event.toolCallId,
+                            name: event.toolName,
+                            args: event.args,
+                        });
+                        break;
+                    case 'tool-result':
+                        stepToolResults.push({
+                            id: event.toolCallId,
+                            name: event.toolName,
+                            result: typeof event.result === 'string'
+                                ? event.result
+                                : JSON.stringify(event.result),
+                        });
+                        break;
+                    case 'usage':
+                        flushStep();
+                        break;
+                    case 'done':
+                        flushStep();
+                        break;
+                }
+            }
+            parent.send(event);
+        },
+    };
+    return channel;
+}
+async function postStreamCleanup(persistingChannel, aiMiddlewares, singletonServices, storage, memoryConfig, threadId, workingMemorySchemaName, messages, aiRunState, runId) {
+    let outputText = persistingChannel.fullText;
+    let outputMessages = messages;
+    for (let i = aiMiddlewares.length - 1; i >= 0; i--) {
+        const mw = aiMiddlewares[i];
+        if (mw.modifyOutput) {
+            const result = await mw.modifyOutput(singletonServices, {
+                text: outputText,
+                messages: outputMessages,
+                usage: { inputTokens: 0, outputTokens: 0 },
+            });
+            outputText = result.text;
+            outputMessages = result.messages;
+        }
+    }
+    if (storage && memoryConfig?.workingMemory && outputText) {
+        const parsed = parseWorkingMemory(outputText);
+        if (parsed) {
+            const existing = (await storage.getWorkingMemory(threadId, 'thread')) ?? {};
+            const merged = deepMergeWorkingMemory(existing, parsed);
+            if (singletonServices.schema && workingMemorySchemaName) {
+                try {
+                    await singletonServices.schema.validateSchema(workingMemorySchemaName, merged);
+                }
+                catch (err) {
+                    singletonServices.logger.warn(`Working memory validation failed: ${err instanceof Error ? err.message : String(err)}`);
+                }
+            }
+            await storage.saveWorkingMemory(threadId, 'thread', merged);
+        }
+    }
+    await aiRunState.updateRun(runId, { status: 'completed' });
+}
+export async function streamAIAgent(agentName, input, channel, params, agentSessionMap, options) {
+    const sessionMap = agentSessionMap ?? new Map();
+    const streamContext = { channel, options };
+    const { agent, packageName, resolvedName, agentRunner, storage, memoryConfig, threadId, userMessage, runnerParams, missingRpcs, workingMemorySchemaName, } = await prepareAgentRun(agentName, input, params, sessionMap, streamContext);
+    const { singletonServices } = params;
+    const { aiRunState } = singletonServices;
+    if (!aiRunState) {
+        throw new Error('AIRunStateService not available in singletonServices');
+    }
+    if (missingRpcs.length > 0) {
+        await aiRunState.createRun({
+            agentName,
+            threadId,
+            resourceId: input.resourceId,
+            status: 'suspended',
+            suspendReason: 'rpc-missing',
+            missingRpcs,
+            usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
+            createdAt: new Date(),
+            updatedAt: new Date(),
+        });
+        channel.send({ type: 'suspended', reason: 'rpc-missing', missingRpcs });
+        channel.send({ type: 'done' });
+        return;
+    }
+    const aiMiddlewares = agent.aiMiddleware ?? [];
+    let modifiedMessages = runnerParams.messages;
+    let modifiedInstructions = runnerParams.instructions;
+    for (const mw of aiMiddlewares) {
+        if (mw.modifyInput) {
+            const result = await mw.modifyInput(singletonServices, {
+                messages: modifiedMessages,
+                instructions: modifiedInstructions,
+            });
+            modifiedMessages = result.messages;
+            modifiedInstructions = result.instructions;
+        }
+    }
+    runnerParams.messages = modifiedMessages;
+    runnerParams.instructions = modifiedInstructions;
+    const runId = await aiRunState.createRun({
+        agentName,
+        threadId,
+        resourceId: input.resourceId,
+        status: 'running',
+        usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
+        createdAt: new Date(),
+        updatedAt: new Date(),
+    });
+    if (storage) {
+        await storage.saveMessages(threadId, [userMessage]);
+    }
+    const streamMiddleware = aiMiddlewares
+        .filter((mw) => mw.modifyOutputStream)
+        .map((mw) => {
+        const state = {};
+        const allEvents = [];
+        return async (services, event, next) => {
+            allEvents.push(event);
+            const result = await mw.modifyOutputStream(services, {
+                event,
+                allEvents,
+                state,
+            });
+            if (result != null)
+                await next(result);
+        };
+    });
+    const agentsMeta = pikkuState(packageName, 'agent', 'agentsMeta');
+    const meta = agentsMeta[resolvedName];
+    const allChannelMiddleware = combineChannelMiddleware('agent', `stream:${agentName}`, {
+        wireInheritedChannelMiddleware: meta?.channelMiddleware,
+        wireChannelMiddleware: [
+            ...(agent.channelMiddleware ?? []),
+            ...streamMiddleware,
+        ],
+    });
+    const wrappedChannel = allChannelMiddleware.length > 0
+        ? wrapChannelWithMiddleware({ channel }, singletonServices, allChannelMiddleware).channel
+        : channel;
+    const persistingChannel = createPersistingChannel(wrappedChannel, storage, threadId);
+    try {
+        await agentRunner.stream(runnerParams, persistingChannel);
+        await postStreamCleanup(persistingChannel, aiMiddlewares, singletonServices, storage, memoryConfig, threadId, workingMemorySchemaName, runnerParams.messages, aiRunState, runId);
+    }
+    catch (err) {
+        if (err instanceof ToolApprovalRequired) {
+            await persistingChannel.flush();
+            const pendingApproval = err.agentRunId
+                ? {
+                    type: 'agent-call',
+                    toolCallId: err.toolCallId,
+                    agentName: err.toolName,
+                    agentRunId: err.agentRunId,
+                    displayToolName: err.displayToolName ?? err.toolName,
+                    displayArgs: err.displayArgs ?? err.args,
+                }
+                : {
+                    type: 'tool-call',
+                    toolCallId: err.toolCallId,
+                    toolName: err.toolName,
+                    args: err.args,
+                };
+            await aiRunState.updateRun(runId, {
+                status: 'suspended',
+                suspendReason: 'approval',
+                pendingApprovals: [pendingApproval],
+            });
+            const approvalEvent = {
+                type: 'approval-request',
+                toolCallId: err.toolCallId,
+                toolName: err.displayToolName ?? err.toolName,
+                args: err.displayArgs ?? err.args,
+                reason: err.reason,
+                runId,
+            };
+            channel.send(approvalEvent);
+            channel.send({ type: 'done' });
+            channel.close();
+            return;
+        }
+        await aiRunState.updateRun(runId, { status: 'failed' });
+        channel.send({
+            type: 'error',
+            message: err instanceof Error ? err.message : String(err),
+        });
+        channel.send({ type: 'done' });
+        channel.close();
+    }
+}
+export async function resumeAIAgent(input, channel, params, options) {
+    const { singletonServices } = params;
+    const { aiRunState } = singletonServices;
+    if (!aiRunState) {
+        throw new Error('AIRunStateService not available in singletonServices');
+    }
+    const found = await aiRunState.findRunByToolCallId(input.toolCallId);
+    if (!found) {
+        throw new Error(`No suspended run found for toolCallId ${input.toolCallId}`);
+    }
+    const { run, approval: pending } = found;
+    await aiRunState.resolveApproval(input.toolCallId, input.approved ? 'approved' : 'denied');
+    const { agent, packageName, resolvedName } = resolveAgent(run.agentName);
+    const { storage } = resolveMemoryServices(agent, singletonServices);
+    const memoryConfig = agent.memory;
+    const agentRunner = singletonServices.aiAgentRunner;
+    if (!agentRunner) {
+        throw new Error('AIAgentRunnerService not available in singletonServices');
+    }
+    if (!input.approved) {
+        if (pending.type === 'agent-call') {
+            await aiRunState.updateRun(pending.agentRunId, { status: 'failed' });
+        }
+        const denialResult = 'Tool call was not approved by the user';
+        if (storage) {
+            await storage.saveMessages(run.threadId, [
+                {
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: [
+                        {
+                            id: input.toolCallId,
+                            name: pending.type === 'tool-call'
+                                ? pending.toolName
+                                : pending.agentName,
+                            result: denialResult,
+                        },
+                    ],
+                    createdAt: new Date(),
+                },
+            ]);
+        }
+        await aiRunState.updateRun(run.runId, { status: 'running' });
+        await continueAfterToolResult(run, agent, packageName, resolvedName, storage, memoryConfig, agentRunner, channel, params, aiRunState, options);
+        return;
+    }
+    if (pending.type === 'agent-call') {
+        const subRun = await aiRunState.getRun(pending.agentRunId);
+        if (!subRun) {
+            throw new Error(`Sub-agent run not found: ${pending.agentRunId}`);
+        }
+        const subPending = subRun.pendingApprovals?.[0];
+        if (!subPending) {
+            throw new Error(`No pending approval on sub-agent run ${pending.agentRunId}`);
+        }
+        const subChannel = createScopedChannel(channel, subRun.agentName, 'resume');
+        channel.send({
+            type: 'agent-call',
+            agentName: subRun.agentName,
+            session: 'resume',
+            input: null,
+        });
+        await resumeAIAgent({ toolCallId: subPending.toolCallId, approved: true }, subChannel, params, options);
+        channel.send({
+            type: 'agent-result',
+            agentName: subRun.agentName,
+            session: 'resume',
+            result: null,
+        });
+        if (storage) {
+            await storage.saveMessages(run.threadId, [
+                {
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: [
+                        {
+                            id: input.toolCallId,
+                            name: pending.agentName,
+                            result: 'Sub-agent completed successfully',
+                        },
+                    ],
+                    createdAt: new Date(),
+                },
+            ]);
+        }
+    }
+    else {
+        const streamContext = {
+            channel,
+            options: { ...options, requiresToolApproval: false },
+        };
+        const { tools } = buildToolDefs(singletonServices, params, new Map(), run.resourceId, resolvedName, packageName, streamContext);
+        const matchingTool = tools.find((t) => t.name === pending.toolName);
+        if (!matchingTool) {
+            throw new Error(`Tool "${pending.toolName}" not found in agent definition`);
+        }
+        const toolArgs = typeof pending.args === 'string' ? JSON.parse(pending.args) : pending.args;
+        const { toolApprovalReason: _ignored, ...cleanArgs } = toolArgs;
+        let toolResult;
+        try {
+            toolResult = await matchingTool.execute(cleanArgs);
+        }
+        catch (execErr) {
+            toolResult = `Error: ${execErr instanceof Error ? execErr.message : String(execErr)}`;
+        }
+        const resultStr = typeof toolResult === 'string' ? toolResult : JSON.stringify(toolResult);
+        if (storage) {
+            await storage.saveMessages(run.threadId, [
+                {
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: [
+                        {
+                            id: input.toolCallId,
+                            name: pending.toolName,
+                            result: resultStr,
+                        },
+                    ],
+                    createdAt: new Date(),
+                },
+            ]);
+        }
+        channel.send({
+            type: 'tool-result',
+            toolCallId: input.toolCallId,
+            toolName: pending.toolName,
+            result: toolResult,
+        });
+    }
+    await aiRunState.updateRun(run.runId, { status: 'running' });
+    await continueAfterToolResult(run, agent, packageName, resolvedName, storage, memoryConfig, agentRunner, channel, params, aiRunState, options);
+}
+async function continueAfterToolResult(run, agent, packageName, resolvedName, storage, memoryConfig, agentRunner, channel, params, aiRunState, options) {
+    const { singletonServices } = params;
+    const agentsMeta = pikkuState(packageName, 'agent', 'agentsMeta');
+    const meta = agentsMeta[resolvedName];
+    const workingMemorySchemaName = meta?.workingMemorySchema ?? null;
+    const messages = storage
+        ? await storage.getMessages(run.threadId, {
+            lastN: memoryConfig?.lastMessages ?? 20,
+        })
+        : [];
+    const workingMemoryJsonSchema = workingMemorySchemaName
+        ? pikkuState(packageName, 'misc', 'schemas').get(workingMemorySchemaName)
+        : undefined;
+    const contextMessages = await loadContextMessages(memoryConfig, storage, { message: '', threadId: run.threadId, resourceId: run.resourceId }, workingMemoryJsonSchema);
+    const allMessages = [...contextMessages, ...messages];
+    const trimmedMessages = trimMessages(allMessages);
+    const instructions = buildInstructions(resolvedName, packageName);
+    const aiMiddlewares = agent.aiMiddleware ?? [];
+    let modifiedMessages = trimmedMessages;
+    let modifiedInstructions = instructions;
+    for (const mw of aiMiddlewares) {
+        if (mw.modifyInput) {
+            const result = await mw.modifyInput(singletonServices, {
+                messages: modifiedMessages,
+                instructions: modifiedInstructions,
+            });
+            modifiedMessages = result.messages;
+            modifiedInstructions = result.instructions;
+        }
+    }
+    const streamMiddleware = aiMiddlewares
+        .filter((mw) => mw.modifyOutputStream)
+        .map((mw) => {
+        const state = {};
+        const allEvents = [];
+        return async (services, event, next) => {
+            allEvents.push(event);
+            const result = await mw.modifyOutputStream(services, {
+                event,
+                allEvents,
+                state,
+            });
+            if (result != null)
+                await next(result);
+        };
+    });
+    const allChannelMiddleware = combineChannelMiddleware('agent', `stream:${run.agentName}`, {
+        wireInheritedChannelMiddleware: meta?.channelMiddleware,
+        wireChannelMiddleware: [
+            ...(agent.channelMiddleware ?? []),
+            ...streamMiddleware,
+        ],
+    });
+    const wrappedChannel = allChannelMiddleware.length > 0
+        ? wrapChannelWithMiddleware({ channel }, singletonServices, allChannelMiddleware).channel
+        : channel;
+    const persistingChannel = createPersistingChannel(wrappedChannel, storage, run.threadId);
+    const streamContext = { channel, options };
+    const resumeTools = buildToolDefs(singletonServices, params, new Map(), run.resourceId, resolvedName, packageName, streamContext).tools;
+    const resolved = resolveModelConfig(resolvedName, agent);
+    const runnerParams = {
+        model: resolved.model,
+        temperature: resolved.temperature,
+        instructions: modifiedInstructions,
+        messages: modifiedMessages,
+        tools: resumeTools,
+        maxSteps: resolved.maxSteps ?? 10,
+        toolChoice: (agent.toolChoice ?? 'auto'),
+        outputSchema: meta?.outputSchema
+            ? pikkuState(packageName, 'misc', 'schemas').get(meta.outputSchema)
+            : undefined,
+    };
+    try {
+        await agentRunner.stream(runnerParams, persistingChannel);
+        await postStreamCleanup(persistingChannel, aiMiddlewares, singletonServices, storage, memoryConfig, run.threadId, workingMemorySchemaName, runnerParams.messages, aiRunState, run.runId);
+    }
+    catch (err) {
+        if (err instanceof ToolApprovalRequired) {
+            await persistingChannel.flush();
+            const pendingApproval = err.agentRunId
+                ? {
+                    type: 'agent-call',
+                    toolCallId: err.toolCallId,
+                    agentName: err.toolName,
+                    agentRunId: err.agentRunId,
+                    displayToolName: err.displayToolName ?? err.toolName,
+                    displayArgs: err.displayArgs ?? err.args,
+                }
+                : {
+                    type: 'tool-call',
+                    toolCallId: err.toolCallId,
+                    toolName: err.toolName,
+                    args: err.args,
+                };
+            await aiRunState.updateRun(run.runId, {
+                status: 'suspended',
+                suspendReason: 'approval',
+                pendingApprovals: [pendingApproval],
+            });
+            const approvalEvent = {
+                type: 'approval-request',
+                toolCallId: err.toolCallId,
+                toolName: err.displayToolName ?? err.toolName,
+                args: err.displayArgs ?? err.args,
+                reason: err.reason,
+                runId: run.runId,
+            };
+            channel.send(approvalEvent);
+            channel.send({ type: 'done' });
+            channel.close();
+            return;
+        }
+        await aiRunState.updateRun(run.runId, { status: 'failed' });
+        channel.send({
+            type: 'error',
+            message: err instanceof Error ? err.message : String(err),
+        });
+        channel.send({ type: 'done' });
+        channel.close();
+    }
+}