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 new file mode 100644 index 000000000..08265fed2 --- /dev/null +++ b/apps/open-swe/src/utils/llms/chat-openrouter.ts @@ -0,0 +1,124 @@ +import { + BaseChatModel, + type BaseChatModelParams, +} from "@langchain/core/language_models/chat_models"; +import { + AIMessage, + AIMessageChunk, + BaseMessage, +} from "@langchain/core/messages"; +import { ChatResult } from "@langchain/core/outputs"; +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; + 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; + + constructor(fields: ChatOpenRouterParams) { + super(fields); + this.modelName = fields.modelName; + this.temperature = fields.temperature; + this.maxTokens = fields.maxTokens; + this.graphConfig = fields.graphConfig; + 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 { + 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({ + content: getMessageContentString(response.content), + }), + text: getMessageContentString(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()) as { + choices: { message: { content: string } }[]; + }; + return new AIMessageChunk({ + content: json.choices[0].message.content, + }); + } + + public _combineLLMOutput(): 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/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/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..b3422f28f --- /dev/null +++ b/packages/shared/src/open-swe/openrouter.ts @@ -0,0 +1,82 @@ +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()) as { data: OpenRouterProvider[] }; + 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; + } + + public getKeys(): string[] { + return this.keys; + } +} 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. */