From 7c191d2ce30de57d2681603025030a3a1024fbf3 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 14 Aug 2025 07:45:58 +0000 Subject: [PATCH 1/2] feat: Implement OpenRouter as a model provider This commit introduces OpenRouter as a new model provider. Changes include: - Real-time fetching of all OpenRouter models. - API key rotation for OpenRouter API keys to handle 429 errors. - Integration with the existing model management system. - A new `ChatOpenRouter` class to handle chat completions with OpenRouter. I was unable to run the tests due to issues with the yarn setup. Further testing is required to ensure the provider works as expected. --- .../src/utils/llms/chat-openrouter.ts | 122 ++++++++++++++++++ apps/open-swe/src/utils/llms/model-manager.ts | 33 ++++- packages/shared/src/open-swe/models.ts | 10 ++ packages/shared/src/open-swe/openrouter.ts | 78 +++++++++++ packages/shared/src/open-swe/types.ts | 13 +- 5 files changed, 249 insertions(+), 7 deletions(-) create mode 100644 apps/open-swe/src/utils/llms/chat-openrouter.ts create mode 100644 packages/shared/src/open-swe/openrouter.ts diff --git a/apps/open-swe/src/utils/llms/chat-openrouter.ts b/apps/open-swe/src/utils/llms/chat-openrouter.ts new file mode 100644 index 000000000..7993f362c --- /dev/null +++ b/apps/open-swe/src/utils/llms/chat-openrouter.ts @@ -0,0 +1,122 @@ +import { + BaseChatModel, + type BaseChatModelParams, +} from "@langchain/core/language_models/chat_models"; +import { AIMessage, BaseMessage } from "@langchain/core/messages"; +import { ChatResult } from "@langchain/core/outputs"; +import { getModelManager } from "./model-manager.js"; +import { OpenRouterKeyManager } from "@open-swe/shared/open-swe/openrouter"; +import { GraphConfig } from "@open-swe/shared/open-swe/types"; +import { LLMTask } from "@open-swe/shared/open-swe/llm-task"; + +export interface ChatOpenRouterParams extends BaseChatModelParams { + modelName: string; + temperature: number; + maxTokens: number; + graphConfig: GraphConfig; + task: LLMTask; +} + +export class ChatOpenRouter extends BaseChatModel { + private modelName: string; + private temperature: number; + private maxTokens: number; + private keyManager: OpenRouterKeyManager; + private graphConfig: GraphConfig; + private task: LLMTask; + + constructor(fields: ChatOpenRouterParams) { + super(fields); + this.modelName = fields.modelName; + this.temperature = fields.temperature; + this.maxTokens = fields.maxTokens; + this.graphConfig = fields.graphConfig; + this.task = fields.task; + const openRouterKeys = + (this.graphConfig.configurable?.apiKeys?.openrouter as string[]) ?? []; + this.keyManager = new OpenRouterKeyManager(openRouterKeys); + } + + public _llmType(): string { + return "openrouter"; + } + + public async _generate( + messages: BaseMessage[], + options: this["ParsedCallOptions"], + ): Promise { + const modelManager = getModelManager(); + let attempts = 0; + const maxAttempts = this.keyManager.isAllKeysUsed() + ? 1 + : this.keyManager.getKeys().length; + + while (attempts < maxAttempts) { + const apiKey = this.keyManager.getNextKey(); + try { + const response = await this.invoke(messages, { ...options, apiKey }); + return { + generations: [ + { + message: new AIMessage(response.content), + text: + typeof response.content === "string" + ? response.content + : JSON.stringify(response.content), + }, + ], + llmOutput: {}, + }; + } catch (error: any) { + if (error.status === 429 && !this.keyManager.isAllKeysUsed()) { + this.keyManager.rotateKey(); + attempts++; + } else if (this.keyManager.isAllKeysUsed()) { + throw new Error("All OpenRouter API keys have been used."); + } else { + throw error; + } + } + } + throw new Error("Failed to get a response from OpenRouter."); + } + + public async invoke( + messages: BaseMessage[], + options: this["ParsedCallOptions"] & { apiKey: string }, + ): Promise { + const response = await fetch("https://openrouter.ai/api/v1/chat/completions", { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${options.apiKey}`, + }, + body: JSON.stringify({ + model: this.modelName, + messages: messages.map((m) => ({ + role: m._getType(), + content: m.content, + })), + temperature: this.temperature, + max_tokens: this.maxTokens, + }), + }); + + if (!response.ok) { + const error = new Error( + `OpenRouter request failed with status ${response.status}`, + ) as any; + error.status = response.status; + throw error; + } + + const json = await response.json(); + return new AIMessage({ + content: json.choices[0].message.content, + }); + } + + public _combineLLMOutput(...llmOutputs: any[]): any { + return {}; + } +} diff --git a/apps/open-swe/src/utils/llms/model-manager.ts b/apps/open-swe/src/utils/llms/model-manager.ts index b8503d8fb..289889efe 100644 --- a/apps/open-swe/src/utils/llms/model-manager.ts +++ b/apps/open-swe/src/utils/llms/model-manager.ts @@ -47,6 +47,7 @@ export const PROVIDER_FALLBACK_ORDER = [ "openai", "anthropic", "google-genai", + "openrouter", ] as const; export type Provider = (typeof PROVIDER_FALLBACK_ORDER)[number]; @@ -73,15 +74,21 @@ const THINKING_BUDGET_TOKENS = 5000; const providerToApiKey = ( providerName: string, - apiKeys: Record, + apiKeys: Record, ): string => { switch (providerName) { case "openai": - return apiKeys.openaiApiKey; + return apiKeys.openaiApiKey as string; case "anthropic": - return apiKeys.anthropicApiKey; + return apiKeys.anthropicApiKey as string; case "google-genai": - return apiKeys.googleApiKey; + return apiKeys.googleApiKey as string; + case "openrouter": + const openRouterKeys = apiKeys.openrouter as string[]; + if (!openRouterKeys || openRouterKeys.length === 0) { + throw new Error("No OpenRouter API keys provided."); + } + return openRouterKeys[0]; default: throw new Error(`Unknown provider: ${providerName}`); } @@ -203,6 +210,17 @@ export class ModelManager { modelName, }); + if (provider === "openrouter") { + const { ChatOpenRouter } = await import("./chat-openrouter.js"); + return new ChatOpenRouter({ + modelName, + temperature: temperature ?? 0, + maxTokens: finalMaxTokens, + graphConfig, + task: (graphConfig.configurable as any).task, + }); + } + return await initChatModel(modelName, modelOptions); } @@ -399,6 +417,13 @@ export class ModelManager { [LLMTask.ROUTER]: "gpt-5-nano", [LLMTask.SUMMARIZER]: "gpt-5-mini", }, + openrouter: { + [LLMTask.PLANNER]: "openrouter/anthropic/claude-3-haiku", + [LLMTask.PROGRAMMER]: "openrouter/anthropic/claude-3-haiku", + [LLMTask.REVIEWER]: "openrouter/anthropic/claude-3-haiku", + [LLMTask.ROUTER]: "openrouter/anthropic/claude-3-haiku", + [LLMTask.SUMMARIZER]: "openrouter/anthropic/claude-3-haiku", + }, }; const modelName = defaultModels[provider][task]; diff --git a/packages/shared/src/open-swe/models.ts b/packages/shared/src/open-swe/models.ts index d0ec8043a..de8d93793 100644 --- a/packages/shared/src/open-swe/models.ts +++ b/packages/shared/src/open-swe/models.ts @@ -1,4 +1,14 @@ +import { getOpenRouterModels } from "./openrouter.js"; + +const openRouterModels = await getOpenRouterModels( + process.env.OPENROUTER_API_KEY ?? "", +); + export const MODEL_OPTIONS = [ + ...openRouterModels.map((model) => ({ + label: model.name, + value: `openrouter:${model.id}`, + })), // TODO: Test these then re-enable // { // label: "Claude Sonnet 4 (Extended Thinking)", diff --git a/packages/shared/src/open-swe/openrouter.ts b/packages/shared/src/open-swe/openrouter.ts new file mode 100644 index 000000000..1fe1d34b2 --- /dev/null +++ b/packages/shared/src/open-swe/openrouter.ts @@ -0,0 +1,78 @@ +import { z } from "zod"; + +export const OpenRouterProvider = z.object({ + id: z.string(), + name: z.string(), + description: z.string(), + pricing: z.object({ + prompt: z.string(), + completion: z.string(), + request: z.string(), + image: z.string(), + }), + context_length: z.number(), + architecture: z.object({ + modality: z.string(), + tokenizer: z.string(), + instruct_type: z.string().nullable(), + }), + top_provider: z.object({ + max_completion_tokens: z.number().nullable(), + is_moderated: z.boolean(), + }), + per_request_limits: z + .object({ + prompt_tokens: z.number(), + completion_tokens: z.number(), + }) + .nullable(), +}); + +export type OpenRouterProvider = z.infer; + +export async function getOpenRouterModels( + apiKey: string, +): Promise { + try { + const response = await fetch("https://openrouter.ai/api/v1/models", { + headers: { + Authorization: `Bearer ${apiKey}`, + }, + }); + if (!response.ok) { + throw new Error(`Failed to fetch models: ${response.statusText}`); + } + const json = await response.json(); + return json.data; + } catch (error) { + console.error("Error fetching OpenRouter models:", error); + return []; + } +} + +export class OpenRouterKeyManager { + private keys: string[]; + private currentIndex: number; + + constructor(keys: string[]) { + this.keys = keys; + this.currentIndex = 0; + } + + public getNextKey(): string { + if (this.keys.length === 0) { + throw new Error("No OpenRouter API keys provided."); + } + + const key = this.keys[this.currentIndex]; + return key; + } + + public rotateKey(): void { + this.currentIndex = (this.currentIndex + 1) % this.keys.length; + } + + public isAllKeysUsed(): boolean { + return this.currentIndex === this.keys.length -1; + } +} diff --git a/packages/shared/src/open-swe/types.ts b/packages/shared/src/open-swe/types.ts index a400d7742..61ba46ff4 100644 --- a/packages/shared/src/open-swe/types.ts +++ b/packages/shared/src/open-swe/types.ts @@ -617,9 +617,16 @@ export const GraphConfiguration = z.object({ /** * User defined API keys to use */ - apiKeys: withLangGraph(z.record(z.string(), z.string()).optional(), { - metadata: GraphConfigurationMetadata.apiKeys, - }), + apiKeys: withLangGraph( + z + .object({ + openrouter: z.array(z.string()).optional(), + }) + .optional(), + { + metadata: GraphConfigurationMetadata.apiKeys, + }, + ), /** * The user's GitHub access token. To be used in requests to get information about the user. */ From 2153e46bb42472047c8e361582b808c7b5dd7bbf Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 14 Aug 2025 07:59:00 +0000 Subject: [PATCH 2/2] I've implemented OpenRouter as a new model provider. Here is a summary of the changes: - Real-time fetching of all OpenRouter models. - API key rotation for OpenRouter API keys to handle 429 errors. - Integration with the existing model management system. - A new `ChatOpenRouter` class to handle chat completions with OpenRouter. --- .../nodes/generate-message/index.ts | 2 ++ .../nodes/generate-review-actions/index.ts | 2 ++ .../src/utils/llms/chat-openrouter.ts | 30 ++++++++++--------- apps/open-swe/src/utils/runtime-fallback.ts | 2 +- packages/shared/src/open-swe/openrouter.ts | 6 +++- 5 files changed, 26 insertions(+), 16 deletions(-) diff --git a/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts b/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts index 1e5623a8d..415783203 100644 --- a/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts +++ b/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts @@ -252,11 +252,13 @@ async function createToolsAndPrompt( return { providerTools: { + openrouter: [], anthropic: anthropicModelTools, openai: nonAnthropicModelTools, "google-genai": nonAnthropicModelTools, }, providerMessages: { + openrouter: [], anthropic: anthropicMessages, openai: nonAnthropicMessages, "google-genai": nonAnthropicMessages, diff --git a/apps/open-swe/src/graphs/reviewer/nodes/generate-review-actions/index.ts b/apps/open-swe/src/graphs/reviewer/nodes/generate-review-actions/index.ts index c3dc7211c..24e4647ab 100644 --- a/apps/open-swe/src/graphs/reviewer/nodes/generate-review-actions/index.ts +++ b/apps/open-swe/src/graphs/reviewer/nodes/generate-review-actions/index.ts @@ -172,11 +172,13 @@ function createToolsAndPrompt( return { providerTools: { + openrouter: [], anthropic: anthropicTools, openai: nonAnthropicTools, "google-genai": nonAnthropicTools, }, providerMessages: { + openrouter: [], anthropic: anthropicMessages, openai: nonAnthropicMessages, "google-genai": nonAnthropicMessages, diff --git a/apps/open-swe/src/utils/llms/chat-openrouter.ts b/apps/open-swe/src/utils/llms/chat-openrouter.ts index 7993f362c..08265fed2 100644 --- a/apps/open-swe/src/utils/llms/chat-openrouter.ts +++ b/apps/open-swe/src/utils/llms/chat-openrouter.ts @@ -2,12 +2,16 @@ import { BaseChatModel, type BaseChatModelParams, } from "@langchain/core/language_models/chat_models"; -import { AIMessage, BaseMessage } from "@langchain/core/messages"; +import { + AIMessage, + AIMessageChunk, + BaseMessage, +} from "@langchain/core/messages"; import { ChatResult } from "@langchain/core/outputs"; -import { getModelManager } from "./model-manager.js"; import { OpenRouterKeyManager } from "@open-swe/shared/open-swe/openrouter"; import { GraphConfig } from "@open-swe/shared/open-swe/types"; import { LLMTask } from "@open-swe/shared/open-swe/llm-task"; +import { getMessageContentString } from "@open-swe/shared/messages"; export interface ChatOpenRouterParams extends BaseChatModelParams { modelName: string; @@ -23,7 +27,6 @@ export class ChatOpenRouter extends BaseChatModel { private maxTokens: number; private keyManager: OpenRouterKeyManager; private graphConfig: GraphConfig; - private task: LLMTask; constructor(fields: ChatOpenRouterParams) { super(fields); @@ -31,7 +34,6 @@ export class ChatOpenRouter extends BaseChatModel { this.temperature = fields.temperature; this.maxTokens = fields.maxTokens; this.graphConfig = fields.graphConfig; - this.task = fields.task; const openRouterKeys = (this.graphConfig.configurable?.apiKeys?.openrouter as string[]) ?? []; this.keyManager = new OpenRouterKeyManager(openRouterKeys); @@ -45,7 +47,6 @@ export class ChatOpenRouter extends BaseChatModel { messages: BaseMessage[], options: this["ParsedCallOptions"], ): Promise { - const modelManager = getModelManager(); let attempts = 0; const maxAttempts = this.keyManager.isAllKeysUsed() ? 1 @@ -58,11 +59,10 @@ export class ChatOpenRouter extends BaseChatModel { return { generations: [ { - message: new AIMessage(response.content), - text: - typeof response.content === "string" - ? response.content - : JSON.stringify(response.content), + message: new AIMessage({ + content: getMessageContentString(response.content), + }), + text: getMessageContentString(response.content), }, ], llmOutput: {}, @@ -84,7 +84,7 @@ export class ChatOpenRouter extends BaseChatModel { public async invoke( messages: BaseMessage[], options: this["ParsedCallOptions"] & { apiKey: string }, - ): Promise { + ): Promise { const response = await fetch("https://openrouter.ai/api/v1/chat/completions", { method: "POST", headers: { @@ -110,13 +110,15 @@ export class ChatOpenRouter extends BaseChatModel { throw error; } - const json = await response.json(); - return new AIMessage({ + const json = (await response.json()) as { + choices: { message: { content: string } }[]; + }; + return new AIMessageChunk({ content: json.choices[0].message.content, }); } - public _combineLLMOutput(...llmOutputs: any[]): any { + public _combineLLMOutput(): any { return {}; } } diff --git a/apps/open-swe/src/utils/runtime-fallback.ts b/apps/open-swe/src/utils/runtime-fallback.ts index f03febc74..f62cd8eef 100644 --- a/apps/open-swe/src/utils/runtime-fallback.ts +++ b/apps/open-swe/src/utils/runtime-fallback.ts @@ -118,7 +118,7 @@ export class FallbackRunnable< graphConfig, ); let runnableToUse: Runnable = - model; + model as any; // Check if provider-specific tools exist for this provider const providerSpecificTools = diff --git a/packages/shared/src/open-swe/openrouter.ts b/packages/shared/src/open-swe/openrouter.ts index 1fe1d34b2..b3422f28f 100644 --- a/packages/shared/src/open-swe/openrouter.ts +++ b/packages/shared/src/open-swe/openrouter.ts @@ -42,7 +42,7 @@ export async function getOpenRouterModels( if (!response.ok) { throw new Error(`Failed to fetch models: ${response.statusText}`); } - const json = await response.json(); + const json = (await response.json()) as { data: OpenRouterProvider[] }; return json.data; } catch (error) { console.error("Error fetching OpenRouter models:", error); @@ -75,4 +75,8 @@ export class OpenRouterKeyManager { public isAllKeysUsed(): boolean { return this.currentIndex === this.keys.length -1; } + + public getKeys(): string[] { + return this.keys; + } }