illuma-agents 1.0.2

package/src/llm/google/utils/tools.ts

@@ -0,0 +1,160 @@
+import {
+  Tool as GenerativeAITool,
+  ToolConfig,
+  FunctionCallingMode,
+  FunctionDeclaration,
+  FunctionDeclarationsTool,
+  FunctionDeclarationSchema,
+} from '@google/generative-ai';
+import { ToolChoice } from '@langchain/core/language_models/chat_models';
+import { StructuredToolInterface } from '@langchain/core/tools';
+import { isLangChainTool } from '@langchain/core/utils/function_calling';
+import {
+  isOpenAITool,
+  ToolDefinition,
+} from '@langchain/core/language_models/base';
+import { convertToGenerativeAITools } from './common';
+import { GoogleGenerativeAIToolType } from '../types';
+import { removeAdditionalProperties } from './zod_to_genai_parameters';
+
+export function convertToolsToGenAI(
+  tools: GoogleGenerativeAIToolType[],
+  extra?: {
+    toolChoice?: ToolChoice;
+    allowedFunctionNames?: string[];
+  }
+): {
+  tools: GenerativeAITool[];
+  toolConfig?: ToolConfig;
+} {
+  // Extract function declaration processing to a separate function
+  const genAITools = processTools(tools);
+
+  // Simplify tool config creation
+  const toolConfig = createToolConfig(genAITools, extra);
+
+  return { tools: genAITools, toolConfig };
+}
+
+function processTools(tools: GoogleGenerativeAIToolType[]): GenerativeAITool[] {
+  let functionDeclarationTools: FunctionDeclaration[] = [];
+  const genAITools: GenerativeAITool[] = [];
+
+  tools.forEach((tool) => {
+    if (isLangChainTool(tool)) {
+      const [convertedTool] = convertToGenerativeAITools([
+        tool as StructuredToolInterface,
+      ]);
+      if (convertedTool.functionDeclarations) {
+        functionDeclarationTools.push(...convertedTool.functionDeclarations);
+      }
+    } else if (isOpenAITool(tool)) {
+      const { functionDeclarations } = convertOpenAIToolToGenAI(tool);
+      if (functionDeclarations) {
+        functionDeclarationTools.push(...functionDeclarations);
+      } else {
+        throw new Error(
+          'Failed to convert OpenAI structured tool to GenerativeAI tool'
+        );
+      }
+    } else {
+      genAITools.push(tool as GenerativeAITool);
+    }
+  });
+
+  const genAIFunctionDeclaration = genAITools.find(
+    (t) => 'functionDeclarations' in t
+  );
+  if (genAIFunctionDeclaration) {
+    return genAITools.map((tool) => {
+      if (
+        functionDeclarationTools.length > 0 &&
+        'functionDeclarations' in tool
+      ) {
+        const newTool = {
+          functionDeclarations: [
+            ...(tool.functionDeclarations || []),
+            ...functionDeclarationTools,
+          ],
+        };
+        // Clear the functionDeclarationTools array so it is not passed again
+        functionDeclarationTools = [];
+        return newTool;
+      }
+      return tool;
+    });
+  }
+
+  return [
+    ...genAITools,
+    ...(functionDeclarationTools.length > 0
+      ? [
+        {
+          functionDeclarations: functionDeclarationTools,
+        },
+      ]
+      : []),
+  ];
+}
+
+function convertOpenAIToolToGenAI(
+  tool: ToolDefinition
+): FunctionDeclarationsTool {
+  return {
+    functionDeclarations: [
+      {
+        name: tool.function.name,
+        description: tool.function.description,
+        parameters: removeAdditionalProperties(
+          tool.function.parameters
+        ) as FunctionDeclarationSchema,
+      },
+    ],
+  };
+}
+
+function createToolConfig(
+  genAITools: GenerativeAITool[],
+  extra?: {
+    toolChoice?: ToolChoice;
+    allowedFunctionNames?: string[];
+  }
+): ToolConfig | undefined {
+  if (!genAITools.length || !extra) return undefined;
+
+  const { toolChoice, allowedFunctionNames } = extra;
+
+  const modeMap: Record<string, FunctionCallingMode> = {
+    any: FunctionCallingMode.ANY,
+    auto: FunctionCallingMode.AUTO,
+    none: FunctionCallingMode.NONE,
+  };
+
+  if (
+    toolChoice != null &&
+    ['any', 'auto', 'none'].includes(toolChoice as string)
+  ) {
+    return {
+      functionCallingConfig: {
+        mode: modeMap[toolChoice as keyof typeof modeMap] ?? 'MODE_UNSPECIFIED',
+        allowedFunctionNames,
+      },
+    };
+  }
+
+  if (typeof toolChoice === 'string' || allowedFunctionNames) {
+    return {
+      functionCallingConfig: {
+        mode: FunctionCallingMode.ANY,
+        allowedFunctionNames: [
+          ...(allowedFunctionNames ?? []),
+          ...(toolChoice != null && typeof toolChoice === 'string'
+            ? [toolChoice]
+            : []),
+        ],
+      },
+    };
+  }
+
+  return undefined;
+}

package/src/llm/google/utils/zod_to_genai_parameters.ts

@@ -0,0 +1,88 @@
+/* eslint-disable @typescript-eslint/no-unused-vars */
+
+import {
+  type FunctionDeclarationSchema as GenerativeAIFunctionDeclarationSchema,
+  type SchemaType as FunctionDeclarationSchemaType,
+} from '@google/generative-ai';
+import {
+  InteropZodType,
+  isInteropZodSchema,
+} from '@langchain/core/utils/types';
+import {
+  type JsonSchema7Type,
+  toJsonSchema,
+} from '@langchain/core/utils/json_schema';
+
+export interface GenerativeAIJsonSchema extends Record<string, unknown> {
+  properties?: Record<string, GenerativeAIJsonSchema>;
+  type: FunctionDeclarationSchemaType;
+}
+
+export interface GenerativeAIJsonSchemaDirty extends GenerativeAIJsonSchema {
+  properties?: Record<string, GenerativeAIJsonSchemaDirty>;
+  additionalProperties?: boolean;
+}
+
+export function removeAdditionalProperties(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  obj: Record<string, any>
+): GenerativeAIJsonSchema {
+  if (typeof obj === 'object' && obj !== null) {
+    const newObj = { ...obj };
+
+    if ('additionalProperties' in newObj) {
+      delete newObj.additionalProperties;
+    }
+    if ('$schema' in newObj) {
+      delete newObj.$schema;
+    }
+    if ('strict' in newObj) {
+      delete newObj.strict;
+    }
+
+    for (const key in newObj) {
+      if (key in newObj) {
+        if (Array.isArray(newObj[key])) {
+          newObj[key] = newObj[key].map(removeAdditionalProperties);
+        } else if (typeof newObj[key] === 'object' && newObj[key] !== null) {
+          newObj[key] = removeAdditionalProperties(newObj[key]);
+        }
+      }
+    }
+
+    return newObj as GenerativeAIJsonSchema;
+  }
+
+  return obj as GenerativeAIJsonSchema;
+}
+
+export function schemaToGenerativeAIParameters<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  RunOutput extends Record<string, any> = Record<string, any>,
+>(
+  schema: InteropZodType<RunOutput> | JsonSchema7Type
+): GenerativeAIFunctionDeclarationSchema {
+  // GenerativeAI doesn't accept either the $schema or additionalProperties
+  // attributes, so we need to explicitly remove them.
+  const jsonSchema = removeAdditionalProperties(
+    isInteropZodSchema(schema) ? toJsonSchema(schema) : schema
+  );
+  const { $schema, ...rest } = jsonSchema;
+
+  return rest as GenerativeAIFunctionDeclarationSchema;
+}
+
+export function jsonSchemaToGeminiParameters(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  schema: Record<string, any>
+): GenerativeAIFunctionDeclarationSchema {
+  // Gemini doesn't accept either the $schema or additionalProperties
+  // attributes, so we need to explicitly remove them.
+
+  const jsonSchema = removeAdditionalProperties(
+    schema as GenerativeAIJsonSchemaDirty
+  );
+  const { $schema, ...rest } = jsonSchema;
+
+  return rest as GenerativeAIFunctionDeclarationSchema;
+}

package/src/llm/ollama/index.ts

@@ -0,0 +1,92 @@
+import { AIMessageChunk } from '@langchain/core/messages';
+import { ChatGenerationChunk } from '@langchain/core/outputs';
+import { ChatOllama as BaseChatOllama } from '@langchain/ollama';
+import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
+import type {
+  ChatResponse as OllamaChatResponse,
+  Message as OllamaMessage,
+} from 'ollama';
+import type { UsageMetadata, BaseMessage } from '@langchain/core/messages';
+import {
+  convertOllamaMessagesToLangChain,
+  convertToOllamaMessages,
+} from './utils';
+
+export class ChatOllama extends BaseChatOllama {
+  static lc_name(): 'IllumaOllama' {
+    return 'IllumaOllama';
+  }
+  async *_streamResponseChunks(
+    messages: BaseMessage[],
+    options: this['ParsedCallOptions'],
+    runManager?: CallbackManagerForLLMRun
+  ): AsyncGenerator<ChatGenerationChunk> {
+    if (this.checkOrPullModel) {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      if (!((await this.checkModelExistsOnMachine(this.model)) as boolean)) {
+        await this.pull(this.model, {
+          logProgress: true,
+        });
+      }
+    }
+
+    const params = this.invocationParams(options);
+    // TODO: remove cast after SDK adds support for tool calls
+    const ollamaMessages = convertToOllamaMessages(messages) as OllamaMessage[];
+
+    const usageMetadata: UsageMetadata = {
+      input_tokens: 0,
+      output_tokens: 0,
+      total_tokens: 0,
+    };
+
+    const stream = await this.client.chat({
+      ...params,
+      messages: ollamaMessages,
+      stream: true,
+    });
+
+    let lastMetadata: Omit<OllamaChatResponse, 'message'> | undefined;
+
+    for await (const chunk of stream) {
+      if (options.signal?.aborted === true) {
+        this.client.abort();
+      }
+      const { message: responseMessage, ...rest } =
+        chunk as Partial<OllamaChatResponse>;
+      usageMetadata.input_tokens += rest.prompt_eval_count ?? 0;
+      usageMetadata.output_tokens += rest.eval_count ?? 0;
+      usageMetadata.total_tokens =
+        usageMetadata.input_tokens + usageMetadata.output_tokens;
+      lastMetadata = rest as Omit<OllamaChatResponse, 'message'>;
+      if (!responseMessage) {
+        continue;
+      }
+      const message = convertOllamaMessagesToLangChain(responseMessage);
+      const generationChunk = new ChatGenerationChunk({
+        text: responseMessage.content || '',
+        message,
+      });
+      yield generationChunk;
+      await runManager?.handleLLMNewToken(
+        responseMessage.content || '',
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        { chunk: generationChunk }
+      );
+    }
+
+    // Yield the `response_metadata` as the final chunk.
+    yield new ChatGenerationChunk({
+      text: '',
+      message: new AIMessageChunk({
+        content: '',
+        response_metadata: lastMetadata,
+        usage_metadata: usageMetadata,
+      }),
+    });
+  }
+}

package/src/llm/ollama/utils.ts

@@ -0,0 +1,193 @@
+import {
+  AIMessage,
+  AIMessageChunk,
+  BaseMessage,
+  HumanMessage,
+  MessageContentText,
+  SystemMessage,
+  ToolMessage,
+  UsageMetadata,
+} from '@langchain/core/messages';
+import type {
+  Message as OllamaMessage,
+  ToolCall as OllamaToolCall,
+} from 'ollama';
+import { v4 as uuidv4 } from 'uuid';
+
+export function convertOllamaMessagesToLangChain(
+  messages: OllamaMessage,
+  extra?: {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    responseMetadata?: Record<string, any>;
+    usageMetadata?: UsageMetadata;
+  }
+): AIMessageChunk {
+  const additional_kwargs: BaseMessage['additional_kwargs'] = {};
+  if ('thinking' in messages) {
+    additional_kwargs.reasoning_content = messages.thinking as string;
+  }
+  return new AIMessageChunk({
+    content: messages.content || '',
+    tool_call_chunks: messages.tool_calls?.map((tc) => ({
+      name: tc.function.name,
+      args: JSON.stringify(tc.function.arguments),
+      type: 'tool_call_chunk',
+      index: 0,
+      id: uuidv4(),
+    })),
+    response_metadata: extra?.responseMetadata,
+    usage_metadata: extra?.usageMetadata,
+    additional_kwargs,
+  });
+}
+
+function extractBase64FromDataUrl(dataUrl: string): string {
+  const match = dataUrl.match(/^data:.*?;base64,(.*)$/);
+  return match ? match[1] : '';
+}
+
+function convertAMessagesToOllama(messages: AIMessage): OllamaMessage[] {
+  if (typeof messages.content === 'string') {
+    return [
+      {
+        role: 'assistant',
+        content: messages.content,
+      },
+    ];
+  }
+
+  const textFields = messages.content.filter(
+    (c) => c.type === 'text' && typeof c.text === 'string'
+  );
+  const textMessages = (textFields as MessageContentText[]).map((c) => ({
+    role: 'assistant',
+    content: c.text,
+  }));
+  let toolCallMsgs: OllamaMessage | undefined;
+
+  if (
+    messages.content.find((c) => c.type === 'tool_use') &&
+    messages.tool_calls?.length
+  ) {
+    // `tool_use` content types are accepted if the message has tool calls
+    const toolCalls: OllamaToolCall[] | undefined = messages.tool_calls.map(
+      (tc) => ({
+        id: tc.id,
+        type: 'function',
+        function: {
+          name: tc.name,
+          arguments: tc.args,
+        },
+      })
+    );
+
+    if (toolCalls) {
+      toolCallMsgs = {
+        role: 'assistant',
+        tool_calls: toolCalls,
+        content: '',
+      };
+    }
+  } else if (
+    messages.content.find((c) => c.type === 'tool_use') &&
+    !messages.tool_calls?.length
+  ) {
+    throw new Error(
+      '\'tool_use\' content type is not supported without tool calls.'
+    );
+  }
+
+  return [...textMessages, ...(toolCallMsgs ? [toolCallMsgs] : [])];
+}
+
+function convertHumanGenericMessagesToOllama(
+  message: HumanMessage
+): OllamaMessage[] {
+  if (typeof message.content === 'string') {
+    return [
+      {
+        role: 'user',
+        content: message.content,
+      },
+    ];
+  }
+  return message.content.map((c) => {
+    if (c.type === 'text') {
+      return {
+        role: 'user',
+        content: c.text,
+      };
+    } else if (c.type === 'image_url') {
+      if (typeof c.image_url === 'string') {
+        return {
+          role: 'user',
+          content: '',
+          images: [extractBase64FromDataUrl(c.image_url)],
+        };
+      } else if (c.image_url.url && typeof c.image_url.url === 'string') {
+        return {
+          role: 'user',
+          content: '',
+          images: [extractBase64FromDataUrl(c.image_url.url)],
+        };
+      }
+    }
+    throw new Error(`Unsupported content type: ${c.type}`);
+  });
+}
+
+function convertSystemMessageToOllama(message: SystemMessage): OllamaMessage[] {
+  if (typeof message.content === 'string') {
+    return [
+      {
+        role: 'system',
+        content: message.content,
+      },
+    ];
+  } else if (
+    message.content.every(
+      (c) => c.type === 'text' && typeof c.text === 'string'
+    )
+  ) {
+    return (message.content as MessageContentText[]).map((c) => ({
+      role: 'system',
+      content: c.text,
+    }));
+  } else {
+    throw new Error(
+      `Unsupported content type(s): ${message.content
+        .map((c) => c.type)
+        .join(', ')}`
+    );
+  }
+}
+
+function convertToolMessageToOllama(message: ToolMessage): OllamaMessage[] {
+  if (typeof message.content !== 'string') {
+    throw new Error('Non string tool message content is not supported');
+  }
+  return [
+    {
+      role: 'tool',
+      content: message.content,
+    },
+  ];
+}
+
+export function convertToOllamaMessages(
+  messages: BaseMessage[]
+): OllamaMessage[] {
+  return messages.flatMap((msg) => {
+    if (['human', 'generic'].includes(msg._getType())) {
+      return convertHumanGenericMessagesToOllama(msg);
+    } else if (msg._getType() === 'ai') {
+      return convertAMessagesToOllama(msg);
+    } else if (msg._getType() === 'system') {
+      return convertSystemMessageToOllama(msg);
+    } else if (msg._getType() === 'tool') {
+      return convertToolMessageToOllama(msg as ToolMessage);
+    } else {
+      throw new Error(`Unsupported message type: ${msg._getType()}`);
+    }
+  });
+}