illuma-agents 1.0.2

package/src/run.ts

@@ -0,0 +1,381 @@
+// src/run.ts
+import './instrumentation';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+import { CallbackHandler } from '@langfuse/langchain';
+import { PromptTemplate } from '@langchain/core/prompts';
+import { SystemMessage } from '@langchain/core/messages';
+import { RunnableLambda } from '@langchain/core/runnables';
+import { AzureChatOpenAI, ChatOpenAI } from '@langchain/openai';
+import type {
+  MessageContentComplex,
+  BaseMessage,
+} from '@langchain/core/messages';
+import type { StringPromptValue } from '@langchain/core/prompt_values';
+import type { RunnableConfig } from '@langchain/core/runnables';
+import type { ClientCallbacks, SystemCallbacks } from '@/graphs/Graph';
+import type * as t from '@/types';
+import { GraphEvents, Providers, Callback, TitleMethod } from '@/common';
+import { manualToolStreamProviders } from '@/llm/providers';
+import { shiftIndexTokenCountMap } from '@/messages/format';
+import {
+  createTitleRunnable,
+  createCompletionTitleRunnable,
+} from '@/utils/title';
+import { createTokenCounter } from '@/utils/tokens';
+import { StandardGraph } from '@/graphs/Graph';
+import { HandlerRegistry } from '@/events';
+import { isOpenAILike } from '@/utils/llm';
+import { isPresent } from '@/utils/misc';
+
+export const defaultOmitOptions = new Set([
+  'stream',
+  'thinking',
+  'streaming',
+  'maxTokens',
+  'clientOptions',
+  'thinkingConfig',
+  'thinkingBudget',
+  'includeThoughts',
+  'maxOutputTokens',
+  'additionalModelRequestFields',
+]);
+
+export class Run<T extends t.BaseGraphState> {
+  graphRunnable?: t.CompiledWorkflow<T, Partial<T>, string>;
+  // private collab!: CollabGraph;
+  // private taskManager!: TaskManager;
+  private handlerRegistry: HandlerRegistry;
+  id: string;
+  Graph: StandardGraph | undefined;
+  provider: Providers | undefined;
+  returnContent: boolean = false;
+
+  private constructor(config: Partial<t.RunConfig>) {
+    const runId = config.runId ?? '';
+    if (!runId) {
+      throw new Error('Run ID not provided');
+    }
+
+    this.id = runId;
+
+    const handlerRegistry = new HandlerRegistry();
+
+    if (config.customHandlers) {
+      for (const [eventType, handler] of Object.entries(
+        config.customHandlers
+      )) {
+        handlerRegistry.register(eventType, handler);
+      }
+    }
+
+    this.handlerRegistry = handlerRegistry;
+
+    if (!config.graphConfig) {
+      throw new Error('Graph config not provided');
+    }
+
+    if (config.graphConfig.type === 'standard' || !config.graphConfig.type) {
+      this.provider = config.graphConfig.llmConfig.provider;
+      this.graphRunnable = this.createStandardGraph(
+        config.graphConfig
+      ) as unknown as t.CompiledWorkflow<T, Partial<T>, string>;
+      if (this.Graph) {
+        this.Graph.handlerRegistry = handlerRegistry;
+      }
+    }
+
+    this.returnContent = config.returnContent ?? false;
+  }
+
+  private createStandardGraph(
+    config: t.StandardGraphConfig
+  ): t.CompiledWorkflow<t.IState, Partial<t.IState>, string> {
+    const { llmConfig, tools = [], ...graphInput } = config;
+    const { provider, ...clientOptions } = llmConfig;
+
+    const standardGraph = new StandardGraph({
+      tools,
+      provider,
+      clientOptions,
+      ...graphInput,
+      runId: this.id,
+    });
+    this.Graph = standardGraph;
+    return standardGraph.createWorkflow();
+  }
+
+  static async create<T extends t.BaseGraphState>(
+    config: t.RunConfig
+  ): Promise<Run<T>> {
+    return new Run<T>(config);
+  }
+
+  getRunMessages(): BaseMessage[] | undefined {
+    if (!this.Graph) {
+      throw new Error(
+        'Graph not initialized. Make sure to use Run.create() to instantiate the Run.'
+      );
+    }
+    return this.Graph.getRunMessages();
+  }
+
+  async processStream(
+    inputs: t.IState,
+    config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; run_id?: string },
+    streamOptions?: t.EventStreamOptions
+  ): Promise<MessageContentComplex[] | undefined> {
+    if (!this.graphRunnable) {
+      throw new Error(
+        'Run not initialized. Make sure to use Run.create() to instantiate the Run.'
+      );
+    }
+    if (!this.Graph) {
+      throw new Error(
+        'Graph not initialized. Make sure to use Run.create() to instantiate the Run.'
+      );
+    }
+
+    this.Graph.resetValues(streamOptions?.keepContent);
+    const provider = this.Graph.provider;
+    const hasTools = this.Graph.tools ? this.Graph.tools.length > 0 : false;
+    if (streamOptions?.callbacks) {
+      /* TODO: conflicts with callback manager */
+      const callbacks = (config.callbacks as t.ProvidedCallbacks) ?? [];
+      config.callbacks = callbacks.concat(
+        this.getCallbacks(streamOptions.callbacks)
+      );
+    }
+
+    if (
+      isPresent(process.env.LANGFUSE_SECRET_KEY) &&
+      isPresent(process.env.LANGFUSE_PUBLIC_KEY) &&
+      isPresent(process.env.LANGFUSE_BASE_URL)
+    ) {
+      const userId = config.configurable?.user_id;
+      const sessionId = config.configurable?.thread_id;
+      const traceMetadata = {
+        messageId: this.id,
+        parentMessageId: config.configurable?.requestBody?.parentMessageId,
+      };
+      const handler = new CallbackHandler({
+        userId,
+        sessionId,
+        traceMetadata,
+      });
+      config.callbacks = (
+        (config.callbacks as t.ProvidedCallbacks) ?? []
+      ).concat([handler]);
+    }
+
+    if (!this.id) {
+      throw new Error('Run ID not provided');
+    }
+
+    const tokenCounter =
+      streamOptions?.tokenCounter ??
+      (streamOptions?.indexTokenCountMap
+        ? await createTokenCounter()
+        : undefined);
+    const tools = this.Graph.tools as
+      | Array<t.GenericTool | undefined>
+      | undefined;
+    const toolTokens = tokenCounter
+      ? (tools?.reduce((acc, tool) => {
+        if (!(tool as Partial<t.GenericTool>).schema) {
+          return acc;
+        }
+
+        const jsonSchema = zodToJsonSchema(
+          (tool?.schema as t.ZodObjectAny).describe(tool?.description ?? ''),
+          tool?.name ?? ''
+        );
+        return (
+          acc + tokenCounter(new SystemMessage(JSON.stringify(jsonSchema)))
+        );
+      }, 0) ?? 0)
+      : 0;
+    let instructionTokens = toolTokens;
+    if (this.Graph.systemMessage && tokenCounter) {
+      instructionTokens += tokenCounter(this.Graph.systemMessage);
+    }
+    const tokenMap = streamOptions?.indexTokenCountMap ?? {};
+    if (this.Graph.systemMessage && instructionTokens > 0) {
+      this.Graph.indexTokenCountMap = shiftIndexTokenCountMap(
+        tokenMap,
+        instructionTokens
+      );
+    } else if (instructionTokens > 0) {
+      tokenMap[0] = tokenMap[0] + instructionTokens;
+      this.Graph.indexTokenCountMap = tokenMap;
+    } else {
+      this.Graph.indexTokenCountMap = tokenMap;
+    }
+
+    this.Graph.maxContextTokens = streamOptions?.maxContextTokens;
+    this.Graph.tokenCounter = tokenCounter;
+
+    config.run_id = this.id;
+    config.configurable = Object.assign(config.configurable ?? {}, {
+      run_id: this.id,
+      provider: this.provider,
+    });
+
+    const stream = this.graphRunnable.streamEvents(inputs, config, {
+      raiseError: true,
+    });
+
+    for await (const event of stream) {
+      const { data, name, metadata, ...info } = event;
+
+      let eventName: t.EventName = info.event;
+      if (
+        hasTools &&
+        manualToolStreamProviders.has(provider) &&
+        eventName === GraphEvents.CHAT_MODEL_STREAM
+      ) {
+        /* Skipping CHAT_MODEL_STREAM event due to double-call edge case */
+        continue;
+      }
+
+      if (eventName && eventName === GraphEvents.ON_CUSTOM_EVENT) {
+        eventName = name;
+      }
+
+      const handler = this.handlerRegistry.getHandler(eventName);
+      if (handler) {
+        handler.handle(eventName, data, metadata, this.Graph);
+      }
+    }
+
+    if (this.returnContent) {
+      return this.Graph.getContentParts();
+    }
+  }
+
+  private createSystemCallback<K extends keyof ClientCallbacks>(
+    clientCallbacks: ClientCallbacks,
+    key: K
+  ): SystemCallbacks[K] {
+    return ((...args: unknown[]) => {
+      const clientCallback = clientCallbacks[key];
+      if (clientCallback && this.Graph) {
+        (clientCallback as (...args: unknown[]) => void)(this.Graph, ...args);
+      }
+    }) as SystemCallbacks[K];
+  }
+
+  getCallbacks(clientCallbacks: ClientCallbacks): SystemCallbacks {
+    return {
+      [Callback.TOOL_ERROR]: this.createSystemCallback(
+        clientCallbacks,
+        Callback.TOOL_ERROR
+      ),
+      [Callback.TOOL_START]: this.createSystemCallback(
+        clientCallbacks,
+        Callback.TOOL_START
+      ),
+      [Callback.TOOL_END]: this.createSystemCallback(
+        clientCallbacks,
+        Callback.TOOL_END
+      ),
+    };
+  }
+
+  async generateTitle({
+    provider,
+    inputText,
+    contentParts,
+    titlePrompt,
+    clientOptions,
+    chainOptions,
+    skipLanguage,
+    omitOptions = defaultOmitOptions,
+    titleMethod = TitleMethod.COMPLETION,
+    titlePromptTemplate,
+  }: t.RunTitleOptions): Promise<{ language?: string; title?: string }> {
+    if (
+      chainOptions != null &&
+      isPresent(process.env.LANGFUSE_SECRET_KEY) &&
+      isPresent(process.env.LANGFUSE_PUBLIC_KEY) &&
+      isPresent(process.env.LANGFUSE_BASE_URL)
+    ) {
+      const userId = chainOptions.configurable?.user_id;
+      const sessionId = chainOptions.configurable?.thread_id;
+      const traceMetadata = {
+        messageId: 'title-' + this.id,
+      };
+      const handler = new CallbackHandler({
+        userId,
+        sessionId,
+        traceMetadata,
+      });
+      chainOptions.callbacks = (
+        (chainOptions.callbacks as t.ProvidedCallbacks) ?? []
+      ).concat([handler]);
+    }
+
+    const convoTemplate = PromptTemplate.fromTemplate(
+      titlePromptTemplate ?? 'User: {input}\nAI: {output}'
+    );
+
+    const response = contentParts
+      .map((part) => {
+        if (part?.type === 'text') return part.text;
+        return '';
+      })
+      .join('\n');
+
+    const model = this.Graph?.getNewModel({
+      provider,
+      omitOptions,
+      clientOptions,
+    });
+    if (!model) {
+      return { language: '', title: '' };
+    }
+    if (
+      isOpenAILike(provider) &&
+      (model instanceof ChatOpenAI || model instanceof AzureChatOpenAI)
+    ) {
+      model.temperature = (clientOptions as t.OpenAIClientOptions | undefined)
+        ?.temperature as number;
+      model.topP = (clientOptions as t.OpenAIClientOptions | undefined)
+        ?.topP as number;
+      model.frequencyPenalty = (
+        clientOptions as t.OpenAIClientOptions | undefined
+      )?.frequencyPenalty as number;
+      model.presencePenalty = (
+        clientOptions as t.OpenAIClientOptions | undefined
+      )?.presencePenalty as number;
+      model.n = (clientOptions as t.OpenAIClientOptions | undefined)
+        ?.n as number;
+    }
+
+    const convoToTitleInput = new RunnableLambda({
+      func: (
+        promptValue: StringPromptValue
+      ): { convo: string; inputText: string; skipLanguage?: boolean } => ({
+        convo: promptValue.value,
+        inputText,
+        skipLanguage,
+      }),
+    }).withConfig({ runName: 'ConvoTransform' });
+
+    const titleChain =
+      titleMethod === TitleMethod.COMPLETION
+        ? await createCompletionTitleRunnable(model, titlePrompt)
+        : await createTitleRunnable(model, titlePrompt);
+
+    /** Pipes `convoTemplate` -> `transformer` -> `titleChain` */
+    const fullChain = convoTemplate
+      .withConfig({ runName: 'ConvoTemplate' })
+      .pipe(convoToTitleInput)
+      .pipe(titleChain)
+      .withConfig({ runName: 'TitleChain' });
+
+    return await fullChain.invoke(
+      { input: inputText, output: response },
+      chainOptions
+    );
+  }
+}

package/src/scripts/abort.ts

@@ -0,0 +1,138 @@
+/* eslint-disable no-console */
+// src/scripts/cli.ts
+import { config } from 'dotenv';
+config();
+import { HumanMessage, BaseMessage } from '@langchain/core/messages';
+import { TavilySearchResults } from '@langchain/community/tools/tavily_search';
+import type * as t from '@/types';
+import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
+import { ToolEndHandler } from '@/events';
+
+import { getArgs } from '@/scripts/args';
+import { Run } from '@/run';
+import { GraphEvents, Callback } from '@/common';
+import { getLLMConfig } from '@/utils/llmConfig';
+
+const conversationHistory: BaseMessage[] = [];
+
+async function testStandardStreaming(): Promise<void> {
+  const { userName, location, provider, currentDate } = await getArgs();
+  const { contentParts, aggregateContent } = createContentAggregator();
+  const controller = new AbortController();
+
+  const customHandlers = {
+    [GraphEvents.TOOL_END]: new ToolEndHandler(),
+    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
+    [GraphEvents.ON_RUN_STEP_COMPLETED]: {
+      handle: (event: GraphEvents.ON_RUN_STEP_COMPLETED, data: t.StreamEventData): void => {
+        console.log('====== ON_RUN_STEP_COMPLETED ======');
+        aggregateContent({ event, data: data as unknown as { result: t.ToolEndEvent } });
+      }
+    },
+    [GraphEvents.ON_RUN_STEP]: {
+      handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData): void => {
+        console.log('====== ON_RUN_STEP ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.RunStep });
+      }
+    },
+    [GraphEvents.ON_RUN_STEP_DELTA]: {
+      handle: (event: GraphEvents.ON_RUN_STEP_DELTA, data: t.StreamEventData): void => {
+        console.log('====== ON_RUN_STEP_DELTA ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.RunStepDeltaEvent });
+      }
+    },
+    [GraphEvents.ON_MESSAGE_DELTA]: {
+      handle: (event: GraphEvents.ON_MESSAGE_DELTA, data: t.StreamEventData): void => {
+        console.log('====== ON_MESSAGE_DELTA ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.MessageDeltaEvent });
+      }
+    },
+    [GraphEvents.TOOL_START]: {
+      handle: (_event: string, data: t.StreamEventData, metadata?: Record<string, unknown>): void => {
+        console.log('====== TOOL_START ======');
+      }
+    },
+  };
+
+  const llmConfig = getLLMConfig(provider);
+  const signal = controller.signal;
+  const run = await Run.create<t.IState>({
+    runId: 'test-run-id',
+    graphConfig: {
+      type: 'standard',
+      signal,
+      llmConfig,
+      tools: [new TavilySearchResults()],
+      instructions: 'You are a friendly AI assistant. Always address the user by their name.',
+      additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
+    },
+    customHandlers,
+  });
+
+  const config = {
+    configurable: {
+      provider,
+      thread_id: 'conversation-num-1',
+    },
+    signal,
+    streamMode: 'values',
+    version: 'v2' as const,
+  };
+
+  console.log('Test 1: Weather query (content parts test)');
+
+  const userMessage = `
+  Make a search for the weather in ${location} today, which is ${currentDate}.
+  Make sure to always refer to me by name, which is ${userName}.
+  After giving me a thorough summary, tell me a joke about the weather forecast we went over.
+  `;
+
+  conversationHistory.push(new HumanMessage(userMessage));
+
+  const inputs = {
+    messages: conversationHistory,
+  };
+
+  // Set a timeout to abort the operation after 5 seconds
+  setTimeout(() => {
+    controller.abort();
+    console.log('Operation aborted');
+    console.log('Current content parts:');
+    console.dir(contentParts, { depth: null });
+  }, 4000);
+
+  try {
+    const finalContentParts = await run.processStream(inputs, config);
+    const finalMessages = run.getRunMessages();
+    if (finalMessages) {
+      conversationHistory.push(...finalMessages);
+      console.dir(conversationHistory, { depth: null });
+    }
+    console.dir(finalContentParts, { depth: null });
+  } catch (error) {
+    if ((error as Error)?.name === 'AbortError') {
+      console.log('Operation was aborted');
+    } else {
+      console.error('An error occurred:', error);
+    }
+  }
+
+  console.log('\n\n====================\n\n');
+  console.dir(contentParts, { depth: null });
+}
+
+process.on('unhandledRejection', (reason, promise) => {
+  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
+  console.log('Conversation history:');
+  process.exit(1);
+});
+
+testStandardStreaming().catch((err) => {
+  console.error(err);
+  console.log('Conversation history:');
+  console.dir(conversationHistory, { depth: null });
+  process.exit(1);
+});

package/src/scripts/ant_web_search.ts

@@ -0,0 +1,158 @@
+/* eslint-disable no-console */
+// src/scripts/cli.ts
+import { config } from 'dotenv';
+config();
+import { HumanMessage, BaseMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+
+import { getArgs } from '@/scripts/args';
+import { Run } from '@/run';
+import { GraphEvents, Callback, Providers } from '@/common';
+import { getLLMConfig } from '@/utils/llmConfig';
+
+const conversationHistory: BaseMessage[] = [];
+let _contentParts: (t.MessageContentComplex | undefined)[] = [];
+async function testStandardStreaming(): Promise<void> {
+  const { userName, location, currentDate } = await getArgs();
+  const { contentParts, aggregateContent } = createContentAggregator();
+  _contentParts = contentParts;
+  const customHandlers = {
+    [GraphEvents.TOOL_END]: new ToolEndHandler(),
+    [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
+    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
+    [GraphEvents.ON_RUN_STEP_COMPLETED]: {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_COMPLETED,
+        data: t.StreamEventData
+      ): void => {
+        console.log('====== ON_RUN_STEP_COMPLETED ======');
+        // console.dir(data, { depth: null });
+        aggregateContent({
+          event,
+          data: data as unknown as { result: t.ToolEndEvent },
+        });
+      },
+    },
+    [GraphEvents.ON_RUN_STEP]: {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP,
+        data: t.StreamEventData
+      ): void => {
+        console.log('====== ON_RUN_STEP ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.RunStep });
+      },
+    },
+    [GraphEvents.ON_RUN_STEP_DELTA]: {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_DELTA,
+        data: t.StreamEventData
+      ): void => {
+        console.log('====== ON_RUN_STEP_DELTA ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.RunStepDeltaEvent });
+      },
+    },
+    [GraphEvents.ON_MESSAGE_DELTA]: {
+      handle: (
+        event: GraphEvents.ON_MESSAGE_DELTA,
+        data: t.StreamEventData
+      ): void => {
+        console.log('====== ON_MESSAGE_DELTA ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.MessageDeltaEvent });
+      },
+    },
+    [GraphEvents.TOOL_START]: {
+      handle: (
+        _event: string,
+        data: t.StreamEventData,
+        metadata?: Record<string, unknown>
+      ): void => {
+        console.log('====== TOOL_START ======');
+        // console.dir(data, { depth: null });
+      },
+    },
+  };
+
+  const llmConfig = getLLMConfig(
+    Providers.ANTHROPIC
+  ) as t.AnthropicClientOptions & t.SharedLLMConfig;
+  llmConfig.model = 'claude-3-5-sonnet-latest';
+
+  const run = await Run.create<t.IState>({
+    runId: 'test-run-id',
+    graphConfig: {
+      type: 'standard',
+      llmConfig,
+      tools: [
+        {
+          type: 'web_search_20250305',
+          name: 'web_search',
+          max_uses: 5,
+        },
+      ],
+      instructions: 'You are a friendly AI assistant.',
+      // additional_instructions: `Always address the user by their name. The user's name is ${userName} and they are located in ${location}.`,
+    },
+    returnContent: true,
+    customHandlers,
+  });
+
+  const config = {
+    configurable: {
+      provider: Providers.ANTHROPIC,
+      thread_id: 'conversation-num-1',
+    },
+    streamMode: 'values',
+    version: 'v2' as const,
+  };
+
+  console.log('Test 1: Search query (search tool test)');
+
+  // const userMessage = `
+  // Make a search for the weather in ${location} today, which is ${currentDate}.
+  // Before making the search, please let me know what you're about to do, then immediately start searching without hesitation.
+  // Make sure to always refer to me by name, which is ${userName}.
+  // After giving me a thorough summary, tell me a joke about the weather forecast we went over.
+  // `;
+  // const userMessage = 'Are massage guns good?';
+  // const userMessage = 'What is functional programming?';
+  const userMessage = "Get me today's trending news.";
+  // const userMessage = "search recent italy earthquake volcano activity";
+  // const userMessage =
+  //   "use 'Trump' as the exact search query and tell me what you find.";
+
+  conversationHistory.push(new HumanMessage(userMessage));
+
+  const inputs = {
+    messages: conversationHistory,
+  };
+  const finalContentParts = await run.processStream(inputs, config);
+  const finalMessages = run.getRunMessages();
+  if (finalMessages) {
+    conversationHistory.push(...finalMessages);
+    console.dir(conversationHistory, { depth: null });
+  }
+  // console.dir(finalContentParts, { depth: null });
+  console.log('\n\n====================\n\n');
+  // console.dir(contentParts, { depth: null });
+}
+
+process.on('unhandledRejection', (reason, promise) => {
+  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
+  console.log('Content Parts:');
+  console.dir(_contentParts, { depth: null });
+  process.exit(1);
+});
+
+testStandardStreaming().catch((err) => {
+  console.error(err);
+  console.log('Conversation history:');
+  console.dir(conversationHistory, { depth: null });
+  console.log('Content Parts:');
+  console.dir(_contentParts, { depth: null });
+  process.exit(1);
+});