@nomyx/assistant/src/ai/providers/OpenRouterProvider/streaming.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import axios from 'axios';
import { ChatMessage, ChatOptions, ProviderResponse } from '../../../types/chat';
import { Tool } from '../../../types/tool';
import { OpenRouterConfig } from './config';
import { OpenRouterChatRequestBody, OpenRouterChatResponse, OpenRouterMessage, OpenRouterDelta, OpenRouterMessageContent, OpenRouterContentPart } from './types';
import { convertToolSchema, convertMessages, extractToolCalls } from './tools';
import { getMaxTokens, supportsTools } from './modelConfig';

function convertContentToOpenRouterFormat(content: string): string {
  return content;
}

export async function* streamChat(
  config: OpenRouterConfig,
  messages: ChatMessage[],
  options: ChatOptions,
  tools?: Tool[]
): AsyncIterableIterator<ProviderResponse> {
  const url = config.getApiUrl('chat/completions');
  const model = config.model;
  const logger = config.getLogger();

  const requestBody: OpenRouterChatRequestBody = {
    messages: convertMessages(messages).map(msg => ({
      ...msg,
      content: convertContentToOpenRouterFormat(Array.isArray(msg.content) ? msg.content.map((part: any) => part.text).join(' ') : msg.content)
    })),
    model: model,
    max_tokens: options.maxTokens || getMaxTokens(model),
    temperature: options.temperature,
    stream: true
  };

  const modelSupportsTools = supportsTools(model);

  if (tools && tools.length > 0) {
    if (modelSupportsTools) {
      requestBody.tools = tools.map(tool => convertToolSchema(tool.getSchema()));
      requestBody.tool_choice = 'auto';
    } else {
      // Append instructions for using tools
      const toolInstructions = `
You have access to the following tools:
${tools.map(tool => `- ${tool.name}: ${tool.description}`).join('\n')}

To use a tool, output a function call in a code block like this:
\`\`\`function_call
{
  "name": "tool_name",
  "arguments": {
    "arg1": "value1",
    "arg2": "value2"
  }
}
\`\`\`
Replace "tool_name" with the actual tool name and provide the necessary arguments.

For example, to use the "tool_name" tool with arguments "arg1" and "arg2":

\`\`\`function_call
{
  "name": "tool_name",
  "arguments": {
    "arg1": "value1",
    "arg2": "value2"
  }
}
\`\`\`
`;
      requestBody.messages[requestBody.messages.length - 1].content += '\n\n' + toolInstructions;
    }
  }

  const headers = {
    ...config.getHeaders(),
    'HTTP-Referer': config.httpReferer,
    'X-Title': config.xTitle
  };

  logger.debug('OpenRouter API streaming request', { url, model, toolsCount: tools?.length });

  try {
    const response = await axios.post<NodeJS.ReadableStream>(url, requestBody, {
      headers,
      responseType: 'stream',
    });
    console.clear();
    logger.debug('OpenRouter API streaming response received');

    let buffer = '';
    let accumulatedMessage: OpenRouterMessage = { role: 'assistant', content: '', tool_calls: [] };
    for await (const chunk of response.data) {
      buffer += chunk.toString();
      let newlineIndex;
      while ((newlineIndex = buffer.indexOf('\n')) !== -1) {
        const line = buffer.slice(0, newlineIndex).trim();
        buffer = buffer.slice(newlineIndex + 1);
        if (line.startsWith('data: ')) {
          const jsonData = line.slice(6);
          //logger.debug('Received JSON data', { jsonData });
          
          if (jsonData === '[DONE]') {
            //logger.debug('End of stream detected');
            return;
          }

          try {
            const data = JSON.parse(jsonData) as OpenRouterChatResponse;
            yield* createStreamingProviderResponse(data, modelSupportsTools, logger, accumulatedMessage);
          } catch (error) {
            logger.error('Error parsing JSON', { line, error });
          }
        } 
      }
    }
  } catch (error) {
    if (axios.isAxiosError(error)) {
      logger.error('OpenRouter API streaming error', {
        status: error.response?.status,
        statusText: error.response?.statusText,
        data: error.response?.data
      });
    } else {
      logger.error('OpenRouter API streaming error', { error });
    }
    throw error;
  }
}

function* createStreamingProviderResponse(
  data: OpenRouterChatResponse, 
  modelSupportsTools: boolean,
  logger: ReturnType<OpenRouterConfig['getLogger']>,
  accumulatedMessage: OpenRouterMessage
): Generator<ProviderResponse> {
  const choice = data.choices[0];
  if (!choice) {
    logger.error('Unexpected response format from OpenRouter API: No choices', { data });
    throw new Error('Unexpected response format from OpenRouter API: No choices');
  }

  if (choice.delta) {
    // Handle streaming format
    updateAccumulatedMessage(accumulatedMessage, choice.delta);
  } else if (choice.message) {
    // Handle non-streaming format
    accumulatedMessage = ensureValidOpenRouterMessage(choice.message);
  } else {
    logger.error('Unexpected response format from OpenRouter API: No delta or message in choice', { choice });
    throw new Error('Unexpected response format from OpenRouter API: No delta or message in choice');
  }

  let content = getContentAsString(accumulatedMessage.content);
  let toolCalls = modelSupportsTools ? extractToolCalls(accumulatedMessage) : [];

  if (!modelSupportsTools) {
    const functionCallRegex = /```function_call\n([\s\S]*?)```/g;
    let match;
    while ((match = functionCallRegex.exec(content)) !== null) {
      try {
        const functionCall = JSON.parse(match[1]);
        toolCalls.push({
          name: functionCall.name,
          arguments: functionCall.arguments
        });
      } catch (error) {
        logger.error('Error parsing function call', { error, match: match[1] });
      }
    }
    content = content.replace(functionCallRegex, '');
  }

  yield {
    content: content.trim(),
    toolCalls: toolCalls,
    usage: {
      promptTokens: data.usage?.prompt_tokens || 0,
      completionTokens: data.usage?.completion_tokens || 0,
      totalTokens: data.usage?.total_tokens || 0
    },
  };
}

function updateAccumulatedMessage(accumulatedMessage: OpenRouterMessage, delta: OpenRouterDelta) {
  if (delta.role && (delta.role === 'assistant' || delta.role === 'user' || delta.role === 'system' || delta.role === 'tool')) {
    accumulatedMessage.role = delta.role;
  }
  if (delta.content) {
    if (typeof accumulatedMessage.content === 'string') {
      accumulatedMessage.content += delta.content;
      process.stdout.write(delta.content);
    } else if (Array.isArray(accumulatedMessage.content)) {
      if (accumulatedMessage.content.length === 0 || typeof accumulatedMessage.content[accumulatedMessage.content.length - 1] !== 'string') {
        accumulatedMessage.content.push({ type: 'text', text: delta.content });
      } else {
        const lastPart = accumulatedMessage.content[accumulatedMessage.content.length - 1];
        if (lastPart.type === 'text') {
          lastPart.text += delta.content;
        } else {
          accumulatedMessage.content.push({ type: 'text', text: delta.content });
        }
      }
    }
  }
  if (delta.tool_calls) {
    if (!accumulatedMessage.tool_calls) {
      accumulatedMessage.tool_calls = [];
    }
    delta.tool_calls.forEach((toolCall, index) => {
      if (!accumulatedMessage.tool_calls![index]) {
        accumulatedMessage.tool_calls![index] = { id: '', type: 'function', function: { name: '', arguments: '' } };
      }
      if (toolCall.id) {
        accumulatedMessage.tool_calls![index].id = toolCall.id;
      }
      if (toolCall.function) {
        if (toolCall.function.name) {
          accumulatedMessage.tool_calls![index].function.name = toolCall.function.name;
        }
        if (toolCall.function.arguments) {
          accumulatedMessage.tool_calls![index].function.arguments += toolCall.function.arguments;
        }
      }
    });
  }
}

function getContentAsString(content: OpenRouterMessageContent): string {
  if (typeof content === 'string') {
    return content;
  } else if (Array.isArray(content)) {
    return content.map(part => {
      if (part.type === 'text') {
        return part.text;
      } else if (part.type === 'image_url') {
        return `[Image: ${part.image_url.url}]`;
      }
      return '';
    }).join(' ');
  }
  return '';
}

function ensureValidOpenRouterMessage(message: any): OpenRouterMessage {
  const validRoles: OpenRouterMessage['role'][] = ['assistant', 'user', 'system', 'tool'];
  const role = validRoles.includes(message.role) ? message.role : 'assistant';
  return {
    role,
    content: message.content || '',
    tool_calls: message.tool_calls || []
  };
}