illuma-agents 1.0.2

package/dist/esm/graphs/Graph.mjs

@@ -0,0 +1,578 @@
+import { nanoid } from 'nanoid';
+import { concat } from '@langchain/core/utils/stream';
+import { ChatVertexAI } from '@langchain/google-vertexai';
+import { StateGraph, START, END } from '@langchain/langgraph';
+import { dispatchCustomEvent } from '@langchain/core/callbacks/dispatch';
+import { SystemMessage, AIMessageChunk, ToolMessage } from '@langchain/core/messages';
+import { GraphNodeKeys, ContentTypes, Providers, GraphEvents, StepTypes } from '../common/enum.mjs';
+import { getChatModelClass, manualToolStreamProviders } from '../llm/providers.mjs';
+import { ToolNode, toolsCondition } from '../tools/ToolNode.mjs';
+import { convertMessagesToContent, formatAnthropicArtifactContent, formatArtifactPayload, modifyDeltaProperties } from '../messages/core.mjs';
+import { createPruneMessages } from '../messages/prune.mjs';
+import { resetIfNotEmpty, joinKeys } from '../utils/graph.mjs';
+import { isOpenAILike, isGoogleLike } from '../utils/llm.mjs';
+import { sleep } from '../utils/run.mjs';
+import 'js-tiktoken/lite';
+import { ChatOpenAI, AzureChatOpenAI } from '../llm/openai/index.mjs';
+import { createFakeStreamingLLM } from '../llm/fake.mjs';
+
+/* eslint-disable no-console */
+// src/graphs/Graph.ts
+const { AGENT, TOOLS } = GraphNodeKeys;
+class Graph {
+    lastToken;
+    tokenTypeSwitch;
+    reasoningKey = 'reasoning_content';
+    currentTokenType = ContentTypes.TEXT;
+    messageStepHasToolCalls = new Map();
+    messageIdsByStepKey = new Map();
+    prelimMessageIdsByStepKey = new Map();
+    config;
+    contentData = [];
+    stepKeyIds = new Map();
+    contentIndexMap = new Map();
+    toolCallStepIds = new Map();
+    currentUsage;
+    indexTokenCountMap = {};
+    maxContextTokens;
+    pruneMessages;
+    /** The amount of time that should pass before another consecutive API call */
+    streamBuffer;
+    tokenCounter;
+    signal;
+    /** Set of invoked tool call IDs from non-message run steps completed mid-run, if any */
+    invokedToolIds;
+    handlerRegistry;
+}
+class StandardGraph extends Graph {
+    graphState;
+    clientOptions;
+    boundModel;
+    /** The last recorded timestamp that a stream API call was invoked */
+    lastStreamCall;
+    systemMessage;
+    messages = [];
+    runId;
+    tools;
+    toolMap;
+    startIndex = 0;
+    provider;
+    toolEnd;
+    signal;
+    constructor({ runId, tools, signal, toolMap, provider, streamBuffer, instructions, reasoningKey, clientOptions, toolEnd = false, additional_instructions = '', }) {
+        super();
+        this.runId = runId;
+        this.tools = tools;
+        this.signal = signal;
+        this.toolEnd = toolEnd;
+        this.toolMap = toolMap;
+        this.provider = provider;
+        this.streamBuffer = streamBuffer;
+        this.clientOptions = clientOptions;
+        this.graphState = this.createGraphState();
+        this.boundModel = this.initializeModel();
+        if (reasoningKey) {
+            this.reasoningKey = reasoningKey;
+        }
+        let finalInstructions = instructions;
+        if (additional_instructions) {
+            finalInstructions =
+                finalInstructions != null && finalInstructions
+                    ? `${finalInstructions}\n\n${additional_instructions}`
+                    : additional_instructions;
+        }
+        if (finalInstructions != null &&
+            finalInstructions &&
+            provider === Providers.ANTHROPIC &&
+            (clientOptions.clientOptions
+                ?.defaultHeaders?.['anthropic-beta']?.includes('prompt-caching') ??
+                false)) {
+            finalInstructions = {
+                content: [
+                    {
+                        type: 'text',
+                        text: instructions,
+                        cache_control: { type: 'ephemeral' },
+                    },
+                ],
+            };
+        }
+        if (finalInstructions != null && finalInstructions !== '') {
+            this.systemMessage = new SystemMessage(finalInstructions);
+        }
+    }
+    /* Init */
+    resetValues(keepContent) {
+        this.messages = [];
+        this.config = resetIfNotEmpty(this.config, undefined);
+        if (keepContent !== true) {
+            this.contentData = resetIfNotEmpty(this.contentData, []);
+            this.contentIndexMap = resetIfNotEmpty(this.contentIndexMap, new Map());
+        }
+        this.stepKeyIds = resetIfNotEmpty(this.stepKeyIds, new Map());
+        this.toolCallStepIds = resetIfNotEmpty(this.toolCallStepIds, new Map());
+        this.messageIdsByStepKey = resetIfNotEmpty(this.messageIdsByStepKey, new Map());
+        this.messageStepHasToolCalls = resetIfNotEmpty(this.prelimMessageIdsByStepKey, new Map());
+        this.prelimMessageIdsByStepKey = resetIfNotEmpty(this.prelimMessageIdsByStepKey, new Map());
+        this.currentTokenType = resetIfNotEmpty(this.currentTokenType, ContentTypes.TEXT);
+        this.lastToken = resetIfNotEmpty(this.lastToken, undefined);
+        this.tokenTypeSwitch = resetIfNotEmpty(this.tokenTypeSwitch, undefined);
+        this.indexTokenCountMap = resetIfNotEmpty(this.indexTokenCountMap, {});
+        this.currentUsage = resetIfNotEmpty(this.currentUsage, undefined);
+        this.tokenCounter = resetIfNotEmpty(this.tokenCounter, undefined);
+        this.maxContextTokens = resetIfNotEmpty(this.maxContextTokens, undefined);
+        this.invokedToolIds = resetIfNotEmpty(this.invokedToolIds, undefined);
+    }
+    /* Run Step Processing */
+    getRunStep(stepId) {
+        const index = this.contentIndexMap.get(stepId);
+        if (index !== undefined) {
+            return this.contentData[index];
+        }
+        return undefined;
+    }
+    getStepKey(metadata) {
+        if (!metadata)
+            return '';
+        const keyList = this.getKeyList(metadata);
+        if (this.checkKeyList(keyList)) {
+            throw new Error('Missing metadata');
+        }
+        return joinKeys(keyList);
+    }
+    getStepIdByKey(stepKey, index) {
+        const stepIds = this.stepKeyIds.get(stepKey);
+        if (!stepIds) {
+            throw new Error(`No step IDs found for stepKey ${stepKey}`);
+        }
+        if (index === undefined) {
+            return stepIds[stepIds.length - 1];
+        }
+        return stepIds[index];
+    }
+    generateStepId(stepKey) {
+        const stepIds = this.stepKeyIds.get(stepKey);
+        let newStepId;
+        let stepIndex = 0;
+        if (stepIds) {
+            stepIndex = stepIds.length;
+            newStepId = `step_${nanoid()}`;
+            stepIds.push(newStepId);
+            this.stepKeyIds.set(stepKey, stepIds);
+        }
+        else {
+            newStepId = `step_${nanoid()}`;
+            this.stepKeyIds.set(stepKey, [newStepId]);
+        }
+        return [newStepId, stepIndex];
+    }
+    getKeyList(metadata) {
+        if (!metadata)
+            return [];
+        const keyList = [
+            metadata.run_id,
+            metadata.thread_id,
+            metadata.langgraph_node,
+            metadata.langgraph_step,
+            metadata.checkpoint_ns,
+        ];
+        if (this.currentTokenType === ContentTypes.THINK ||
+            this.currentTokenType === 'think_and_text') {
+            keyList.push('reasoning');
+        }
+        else if (this.tokenTypeSwitch === 'content') {
+            keyList.push('post-reasoning');
+        }
+        if (this.invokedToolIds != null && this.invokedToolIds.size > 0) {
+            keyList.push(this.invokedToolIds.size + '');
+        }
+        return keyList;
+    }
+    checkKeyList(keyList) {
+        return keyList.some((key) => key === undefined);
+    }
+    /* Misc.*/
+    getRunMessages() {
+        return this.messages.slice(this.startIndex);
+    }
+    getContentParts() {
+        return convertMessagesToContent(this.messages.slice(this.startIndex));
+    }
+    /* Graph */
+    createGraphState() {
+        return {
+            messages: {
+                value: (x, y) => {
+                    if (!x.length) {
+                        if (this.systemMessage) {
+                            x.push(this.systemMessage);
+                        }
+                        this.startIndex = x.length + y.length;
+                    }
+                    const current = x.concat(y);
+                    this.messages = current;
+                    return current;
+                },
+                default: () => [],
+            },
+        };
+    }
+    initializeTools() {
+        // return new ToolNode<t.BaseGraphState>(this.tools);
+        return new ToolNode({
+            tools: this.tools || [],
+            toolMap: this.toolMap,
+            toolCallStepIds: this.toolCallStepIds,
+            errorHandler: (data, metadata) => StandardGraph.handleToolCallErrorStatic(this, data, metadata),
+        });
+    }
+    initializeModel() {
+        const ChatModelClass = getChatModelClass(this.provider);
+        const model = new ChatModelClass(this.clientOptions);
+        if (isOpenAILike(this.provider) &&
+            (model instanceof ChatOpenAI || model instanceof AzureChatOpenAI)) {
+            model.temperature = this.clientOptions
+                .temperature;
+            model.topP = this.clientOptions.topP;
+            model.frequencyPenalty = this.clientOptions
+                .frequencyPenalty;
+            model.presencePenalty = this.clientOptions
+                .presencePenalty;
+            model.n = this.clientOptions.n;
+        }
+        else if (this.provider === Providers.VERTEXAI &&
+            model instanceof ChatVertexAI) {
+            model.temperature = this.clientOptions
+                .temperature;
+            model.topP = this.clientOptions
+                .topP;
+            model.topK = this.clientOptions
+                .topK;
+            model.topLogprobs = this.clientOptions
+                .topLogprobs;
+            model.frequencyPenalty = this.clientOptions
+                .frequencyPenalty;
+            model.presencePenalty = this.clientOptions
+                .presencePenalty;
+            model.maxOutputTokens = this.clientOptions
+                .maxOutputTokens;
+        }
+        if (!this.tools || this.tools.length === 0) {
+            return model;
+        }
+        return model.bindTools(this.tools);
+    }
+    overrideTestModel(responses, sleep, toolCalls) {
+        this.boundModel = createFakeStreamingLLM({
+            responses,
+            sleep,
+            toolCalls,
+        });
+    }
+    getNewModel({ provider, clientOptions, omitOptions, }) {
+        const ChatModelClass = getChatModelClass(provider);
+        const options = omitOptions && clientOptions == null
+            ? Object.assign(Object.fromEntries(Object.entries(this.clientOptions).filter(([key]) => !omitOptions.has(key))), clientOptions)
+            : (clientOptions ?? this.clientOptions);
+        return new ChatModelClass(options);
+    }
+    storeUsageMetadata(finalMessage) {
+        if (finalMessage &&
+            'usage_metadata' in finalMessage &&
+            finalMessage.usage_metadata != null) {
+            this.currentUsage = finalMessage.usage_metadata;
+        }
+    }
+    cleanupSignalListener() {
+        if (!this.signal) {
+            return;
+        }
+        if (!this.boundModel) {
+            return;
+        }
+        const client = this.boundModel?.exposedClient;
+        if (!client?.abortHandler) {
+            return;
+        }
+        this.signal.removeEventListener('abort', client.abortHandler);
+        client.abortHandler = undefined;
+    }
+    createCallModel() {
+        return async (state, config) => {
+            const { provider = '' } = config?.configurable ?? {};
+            if (this.boundModel == null) {
+                throw new Error('No Graph model found');
+            }
+            if (!config || !provider) {
+                throw new Error(`No ${config ? 'provider' : 'config'} provided`);
+            }
+            if (!config.signal) {
+                config.signal = this.signal;
+            }
+            this.config = config;
+            const { messages } = state;
+            let messagesToUse = messages;
+            if (!this.pruneMessages &&
+                this.tokenCounter &&
+                this.maxContextTokens != null &&
+                this.indexTokenCountMap[0] != null) {
+                const isAnthropicWithThinking = (this.provider === Providers.ANTHROPIC &&
+                    this.clientOptions.thinking !=
+                        null) ||
+                    (this.provider === Providers.BEDROCK &&
+                        this.clientOptions
+                            .additionalModelRequestFields?.['thinking'] != null) ||
+                    (this.provider === Providers.OPENAI &&
+                        this.clientOptions.modelKwargs
+                            ?.thinking?.type === 'enabled');
+                this.pruneMessages = createPruneMessages({
+                    provider: this.provider,
+                    indexTokenCountMap: this.indexTokenCountMap,
+                    maxTokens: this.maxContextTokens,
+                    tokenCounter: this.tokenCounter,
+                    startIndex: this.startIndex,
+                    thinkingEnabled: isAnthropicWithThinking,
+                });
+            }
+            if (this.pruneMessages) {
+                const { context, indexTokenCountMap } = this.pruneMessages({
+                    messages,
+                    usageMetadata: this.currentUsage,
+                    // startOnMessageType: 'human',
+                });
+                this.indexTokenCountMap = indexTokenCountMap;
+                messagesToUse = context;
+            }
+            const finalMessages = messagesToUse;
+            const lastMessageX = finalMessages.length >= 2
+                ? finalMessages[finalMessages.length - 2]
+                : null;
+            const lastMessageY = finalMessages.length >= 1
+                ? finalMessages[finalMessages.length - 1]
+                : null;
+            if (provider === Providers.BEDROCK &&
+                lastMessageX instanceof AIMessageChunk &&
+                lastMessageY instanceof ToolMessage &&
+                typeof lastMessageX.content === 'string') {
+                finalMessages[finalMessages.length - 2].content = '';
+            }
+            const isLatestToolMessage = lastMessageY instanceof ToolMessage;
+            if (isLatestToolMessage && provider === Providers.ANTHROPIC) {
+                formatAnthropicArtifactContent(finalMessages);
+            }
+            else if (isLatestToolMessage &&
+                (isOpenAILike(provider) || isGoogleLike(provider))) {
+                formatArtifactPayload(finalMessages);
+            }
+            if (this.lastStreamCall != null && this.streamBuffer != null) {
+                const timeSinceLastCall = Date.now() - this.lastStreamCall;
+                if (timeSinceLastCall < this.streamBuffer) {
+                    const timeToWait = Math.ceil((this.streamBuffer - timeSinceLastCall) / 1000) * 1000;
+                    await sleep(timeToWait);
+                }
+            }
+            this.lastStreamCall = Date.now();
+            let result;
+            if ((this.tools?.length ?? 0) > 0 &&
+                manualToolStreamProviders.has(provider)) {
+                const stream = await this.boundModel.stream(finalMessages, config);
+                let finalChunk;
+                for await (const chunk of stream) {
+                    dispatchCustomEvent(GraphEvents.CHAT_MODEL_STREAM, { chunk }, config);
+                    if (!finalChunk) {
+                        finalChunk = chunk;
+                    }
+                    else {
+                        finalChunk = concat(finalChunk, chunk);
+                    }
+                }
+                finalChunk = modifyDeltaProperties(this.provider, finalChunk);
+                result = { messages: [finalChunk] };
+            }
+            else {
+                const finalMessage = (await this.boundModel.invoke(finalMessages, config));
+                if ((finalMessage.tool_calls?.length ?? 0) > 0) {
+                    finalMessage.tool_calls = finalMessage.tool_calls?.filter((tool_call) => {
+                        if (!tool_call.name) {
+                            return false;
+                        }
+                        return true;
+                    });
+                }
+                result = { messages: [finalMessage] };
+            }
+            this.storeUsageMetadata(result.messages?.[0]);
+            this.cleanupSignalListener();
+            return result;
+        };
+    }
+    createWorkflow() {
+        const routeMessage = (state, config) => {
+            this.config = config;
+            return toolsCondition(state, this.invokedToolIds);
+        };
+        const workflow = new StateGraph({
+            channels: this.graphState,
+        })
+            .addNode(AGENT, this.createCallModel())
+            .addNode(TOOLS, this.initializeTools())
+            .addEdge(START, AGENT)
+            .addConditionalEdges(AGENT, routeMessage)
+            .addEdge(TOOLS, this.toolEnd ? END : AGENT);
+        return workflow.compile();
+    }
+    /* Dispatchers */
+    /**
+     * Dispatches a run step to the client, returns the step ID
+     */
+    dispatchRunStep(stepKey, stepDetails) {
+        if (!this.config) {
+            throw new Error('No config provided');
+        }
+        const [stepId, stepIndex] = this.generateStepId(stepKey);
+        if (stepDetails.type === StepTypes.TOOL_CALLS && stepDetails.tool_calls) {
+            for (const tool_call of stepDetails.tool_calls) {
+                const toolCallId = tool_call.id ?? '';
+                if (!toolCallId || this.toolCallStepIds.has(toolCallId)) {
+                    continue;
+                }
+                this.toolCallStepIds.set(toolCallId, stepId);
+            }
+        }
+        const runStep = {
+            stepIndex,
+            id: stepId,
+            type: stepDetails.type,
+            index: this.contentData.length,
+            stepDetails,
+            usage: null,
+        };
+        const runId = this.runId ?? '';
+        if (runId) {
+            runStep.runId = runId;
+        }
+        this.contentData.push(runStep);
+        this.contentIndexMap.set(stepId, runStep.index);
+        dispatchCustomEvent(GraphEvents.ON_RUN_STEP, runStep, this.config);
+        return stepId;
+    }
+    handleToolCallCompleted(data, metadata, omitOutput) {
+        if (!this.config) {
+            throw new Error('No config provided');
+        }
+        if (!data.output) {
+            return;
+        }
+        const { input, output } = data;
+        const { tool_call_id } = output;
+        const stepId = this.toolCallStepIds.get(tool_call_id) ?? '';
+        if (!stepId) {
+            throw new Error(`No stepId found for tool_call_id ${tool_call_id}`);
+        }
+        const runStep = this.getRunStep(stepId);
+        if (!runStep) {
+            throw new Error(`No run step found for stepId ${stepId}`);
+        }
+        const dispatchedOutput = typeof output.content === 'string'
+            ? output.content
+            : JSON.stringify(output.content);
+        const args = typeof input === 'string' ? input : input.input;
+        const tool_call = {
+            args: typeof args === 'string' ? args : JSON.stringify(args),
+            name: output.name ?? '',
+            id: output.tool_call_id,
+            output: omitOutput === true ? '' : dispatchedOutput,
+            progress: 1,
+        };
+        this.handlerRegistry?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)?.handle(GraphEvents.ON_RUN_STEP_COMPLETED, {
+            result: {
+                id: stepId,
+                index: runStep.index,
+                type: 'tool_call',
+                tool_call,
+            },
+        }, metadata, this);
+    }
+    /**
+     * Static version of handleToolCallError to avoid creating strong references
+     * that prevent garbage collection
+     */
+    static handleToolCallErrorStatic(graph, data, metadata) {
+        if (!graph.config) {
+            throw new Error('No config provided');
+        }
+        if (!data.id) {
+            console.warn('No Tool ID provided for Tool Error');
+            return;
+        }
+        const stepId = graph.toolCallStepIds.get(data.id) ?? '';
+        if (!stepId) {
+            throw new Error(`No stepId found for tool_call_id ${data.id}`);
+        }
+        const { name, input: args, error } = data;
+        const runStep = graph.getRunStep(stepId);
+        if (!runStep) {
+            throw new Error(`No run step found for stepId ${stepId}`);
+        }
+        const tool_call = {
+            id: data.id,
+            name: name || '',
+            args: typeof args === 'string' ? args : JSON.stringify(args),
+            output: `Error processing tool${error?.message != null ? `: ${error.message}` : ''}`,
+            progress: 1,
+        };
+        graph.handlerRegistry
+            ?.getHandler(GraphEvents.ON_RUN_STEP_COMPLETED)
+            ?.handle(GraphEvents.ON_RUN_STEP_COMPLETED, {
+            result: {
+                id: stepId,
+                index: runStep.index,
+                type: 'tool_call',
+                tool_call,
+            },
+        }, metadata, graph);
+    }
+    /**
+     * Instance method that delegates to the static method
+     * Kept for backward compatibility
+     */
+    handleToolCallError(data, metadata) {
+        StandardGraph.handleToolCallErrorStatic(this, data, metadata);
+    }
+    dispatchRunStepDelta(id, delta) {
+        if (!this.config) {
+            throw new Error('No config provided');
+        }
+        else if (!id) {
+            throw new Error('No step ID found');
+        }
+        const runStepDelta = {
+            id,
+            delta,
+        };
+        dispatchCustomEvent(GraphEvents.ON_RUN_STEP_DELTA, runStepDelta, this.config);
+    }
+    dispatchMessageDelta(id, delta) {
+        if (!this.config) {
+            throw new Error('No config provided');
+        }
+        const messageDelta = {
+            id,
+            delta,
+        };
+        dispatchCustomEvent(GraphEvents.ON_MESSAGE_DELTA, messageDelta, this.config);
+    }
+    dispatchReasoningDelta = (stepId, delta) => {
+        if (!this.config) {
+            throw new Error('No config provided');
+        }
+        const reasoningDelta = {
+            id: stepId,
+            delta,
+        };
+        dispatchCustomEvent(GraphEvents.ON_REASONING_DELTA, reasoningDelta, this.config);
+    };
+}
+
+export { Graph, StandardGraph };
+//# sourceMappingURL=Graph.mjs.map