mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
Extract router models fetching logic into dedicated service
Refactor to improve separation of concerns: - Create src/services/router-models/index.ts to handle provider model fetching - Extract buildProviderFetchList() function for fetch options construction - Extract fetchRouterModels() function for coordinated model fetching - Move 150+ lines of provider-specific logic out of webviewMessageHandler - Add comprehensive tests in router-models-service.spec.ts (11 test cases) Benefits: - Cleaner webviewMessageHandler with less business logic - Reusable service for router model operations - Better testability with isolated unit tests - Clear separation between UI message handling and data fetching Files changed: - New: src/services/router-models/index.ts - New: src/services/router-models/__tests__/router-models-service.spec.ts - Modified: src/core/webview/webviewMessageHandler.ts (simplified)
This commit is contained in:
daniel-lxs
parent
97beb1366c
commit
f5cc227c09
3 changed files with 490 additions and 263 deletions
|
|
@ -54,9 +54,9 @@ import { openMention } from "../mentions"
|
|||
import { getWorkspacePath } from "../../utils/path"
|
||||
import { Mode, defaultModeSlug } from "../../shared/modes"
|
||||
import { getModels, flushModels } from "../../api/providers/fetchers/modelCache"
|
||||
import { GetModelsOptions } from "../../shared/api"
|
||||
import { generateSystemPrompt } from "./generateSystemPrompt"
|
||||
import { getCommand } from "../../utils/commands"
|
||||
import { fetchRouterModels } from "../../services/router-models"
|
||||
|
||||
const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"])
|
||||
|
||||
|
|
@ -768,152 +768,27 @@ export const webviewMessageHandler = async (
|
|||
case "requestRouterModels": {
|
||||
// Phase 2: Scope to active provider during chat/task flows
|
||||
const { apiConfiguration } = await provider.getState()
|
||||
const providerStr = apiConfiguration.apiProvider
|
||||
const activeProvider: RouterName | undefined =
|
||||
providerStr && isRouterName(providerStr) ? providerStr : undefined
|
||||
|
||||
const routerModels: Partial<Record<RouterName, ModelRecord>> = {
|
||||
openrouter: {},
|
||||
"vercel-ai-gateway": {},
|
||||
huggingface: {},
|
||||
litellm: {},
|
||||
deepinfra: {},
|
||||
"io-intelligence": {},
|
||||
requesty: {},
|
||||
unbound: {},
|
||||
glama: {},
|
||||
ollama: {},
|
||||
lmstudio: {},
|
||||
roo: {},
|
||||
}
|
||||
const { routerModels, errors } = await fetchRouterModels({
|
||||
apiConfiguration,
|
||||
activeProviderOnly: true,
|
||||
litellmOverrides: message?.values
|
||||
? {
|
||||
apiKey: message.values.litellmApiKey,
|
||||
baseUrl: message.values.litellmBaseUrl,
|
||||
}
|
||||
: undefined,
|
||||
})
|
||||
|
||||
const safeGetModels = async (options: GetModelsOptions): Promise<ModelRecord> => {
|
||||
try {
|
||||
return await getModels(options)
|
||||
} catch (error) {
|
||||
provider.log(
|
||||
`Failed to fetch models in webviewMessageHandler requestRouterModels for ${options.provider}: ${error instanceof Error ? error.message : String(error)}`,
|
||||
)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
// Build full list then filter to active provider
|
||||
const allFetches: { key: RouterName; options: GetModelsOptions }[] = [
|
||||
{ key: "openrouter", options: { provider: "openrouter" } },
|
||||
{
|
||||
key: "requesty",
|
||||
options: {
|
||||
provider: "requesty",
|
||||
apiKey: apiConfiguration.requestyApiKey,
|
||||
baseUrl: apiConfiguration.requestyBaseUrl,
|
||||
},
|
||||
},
|
||||
{ key: "glama", options: { provider: "glama" } },
|
||||
{ key: "unbound", options: { provider: "unbound", apiKey: apiConfiguration.unboundApiKey } },
|
||||
{ key: "vercel-ai-gateway", options: { provider: "vercel-ai-gateway" } },
|
||||
{
|
||||
key: "deepinfra",
|
||||
options: {
|
||||
provider: "deepinfra",
|
||||
apiKey: apiConfiguration.deepInfraApiKey,
|
||||
baseUrl: apiConfiguration.deepInfraBaseUrl,
|
||||
},
|
||||
},
|
||||
{
|
||||
key: "roo",
|
||||
options: {
|
||||
provider: "roo",
|
||||
baseUrl: process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy",
|
||||
apiKey: CloudService.hasInstance()
|
||||
? CloudService.instance.authService?.getSessionToken()
|
||||
: undefined,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
// Include local providers (ollama, lmstudio, huggingface) when they are the active provider
|
||||
if (activeProvider === "ollama") {
|
||||
allFetches.push({
|
||||
key: "ollama",
|
||||
options: {
|
||||
provider: "ollama",
|
||||
baseUrl: apiConfiguration.ollamaBaseUrl,
|
||||
apiKey: apiConfiguration.ollamaApiKey,
|
||||
},
|
||||
// Send error notifications for failed providers
|
||||
errors.forEach((err) => {
|
||||
provider.log(`Error fetching models for ${err.provider}: ${err.error}`)
|
||||
provider.postMessageToWebview({
|
||||
type: "singleRouterModelFetchResponse",
|
||||
success: false,
|
||||
error: err.error,
|
||||
values: { provider: err.provider },
|
||||
})
|
||||
}
|
||||
if (activeProvider === "lmstudio") {
|
||||
allFetches.push({
|
||||
key: "lmstudio",
|
||||
options: {
|
||||
provider: "lmstudio",
|
||||
baseUrl: apiConfiguration.lmStudioBaseUrl,
|
||||
},
|
||||
})
|
||||
}
|
||||
if (activeProvider === "huggingface") {
|
||||
allFetches.push({
|
||||
key: "huggingface",
|
||||
options: {
|
||||
provider: "huggingface",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// IO Intelligence (optional)
|
||||
if (apiConfiguration.ioIntelligenceApiKey) {
|
||||
allFetches.push({
|
||||
key: "io-intelligence",
|
||||
options: { provider: "io-intelligence", apiKey: apiConfiguration.ioIntelligenceApiKey },
|
||||
})
|
||||
}
|
||||
|
||||
// LiteLLM (optional)
|
||||
const litellmApiKey = apiConfiguration.litellmApiKey || message?.values?.litellmApiKey
|
||||
const litellmBaseUrl = apiConfiguration.litellmBaseUrl || message?.values?.litellmBaseUrl
|
||||
if (litellmApiKey && litellmBaseUrl) {
|
||||
allFetches.push({
|
||||
key: "litellm",
|
||||
options: { provider: "litellm", apiKey: litellmApiKey, baseUrl: litellmBaseUrl },
|
||||
})
|
||||
}
|
||||
|
||||
const modelFetchPromises = activeProvider
|
||||
? allFetches.filter(({ key }) => key === activeProvider)
|
||||
: allFetches
|
||||
|
||||
// If nothing matched (edge case), still post empty structure for stability
|
||||
if (modelFetchPromises.length === 0) {
|
||||
await provider.postMessageToWebview({
|
||||
type: "routerModels",
|
||||
routerModels: routerModels as RouterModels,
|
||||
})
|
||||
break
|
||||
}
|
||||
|
||||
const results = await Promise.allSettled(
|
||||
modelFetchPromises.map(async ({ key, options }) => {
|
||||
const models = await safeGetModels(options)
|
||||
return { key, models }
|
||||
}),
|
||||
)
|
||||
|
||||
results.forEach((result, index) => {
|
||||
const routerName = modelFetchPromises[index].key
|
||||
if (result.status === "fulfilled") {
|
||||
routerModels[routerName] = result.value.models
|
||||
} else {
|
||||
const errorMessage = result.reason instanceof Error ? result.reason.message : String(result.reason)
|
||||
provider.log(`Error fetching models for ${routerName}: ${errorMessage}`)
|
||||
routerModels[routerName] = {}
|
||||
provider.postMessageToWebview({
|
||||
type: "singleRouterModelFetchResponse",
|
||||
success: false,
|
||||
error: errorMessage,
|
||||
values: { provider: routerName },
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
provider.postMessageToWebview({ type: "routerModels", routerModels: routerModels as RouterModels })
|
||||
|
|
@ -923,127 +798,42 @@ export const webviewMessageHandler = async (
|
|||
// Settings and activation: fetch all providers (legacy behavior)
|
||||
const { apiConfiguration } = await provider.getState()
|
||||
|
||||
const routerModels: Partial<Record<RouterName, ModelRecord>> = {
|
||||
openrouter: {},
|
||||
"vercel-ai-gateway": {},
|
||||
huggingface: {},
|
||||
litellm: {},
|
||||
deepinfra: {},
|
||||
"io-intelligence": {},
|
||||
requesty: {},
|
||||
unbound: {},
|
||||
glama: {},
|
||||
ollama: {},
|
||||
lmstudio: {},
|
||||
roo: {},
|
||||
}
|
||||
|
||||
const safeGetModels = async (options: GetModelsOptions): Promise<ModelRecord> => {
|
||||
try {
|
||||
return await getModels(options)
|
||||
} catch (error) {
|
||||
provider.log(
|
||||
`Failed to fetch models in webviewMessageHandler requestRouterModelsAll for ${options.provider}: ${error instanceof Error ? error.message : String(error)}`,
|
||||
)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
const modelFetchPromises: { key: RouterName; options: GetModelsOptions }[] = [
|
||||
{ key: "openrouter", options: { provider: "openrouter" } },
|
||||
{
|
||||
key: "requesty",
|
||||
options: {
|
||||
provider: "requesty",
|
||||
apiKey: apiConfiguration.requestyApiKey,
|
||||
baseUrl: apiConfiguration.requestyBaseUrl,
|
||||
},
|
||||
},
|
||||
{ key: "glama", options: { provider: "glama" } },
|
||||
{ key: "unbound", options: { provider: "unbound", apiKey: apiConfiguration.unboundApiKey } },
|
||||
{ key: "vercel-ai-gateway", options: { provider: "vercel-ai-gateway" } },
|
||||
{
|
||||
key: "deepinfra",
|
||||
options: {
|
||||
provider: "deepinfra",
|
||||
apiKey: apiConfiguration.deepInfraApiKey,
|
||||
baseUrl: apiConfiguration.deepInfraBaseUrl,
|
||||
},
|
||||
},
|
||||
{
|
||||
key: "roo",
|
||||
options: {
|
||||
provider: "roo",
|
||||
baseUrl: process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy",
|
||||
apiKey: CloudService.hasInstance()
|
||||
? CloudService.instance.authService?.getSessionToken()
|
||||
: undefined,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
// Add IO Intelligence if API key is provided.
|
||||
const ioIntelligenceApiKey = apiConfiguration.ioIntelligenceApiKey
|
||||
if (ioIntelligenceApiKey) {
|
||||
modelFetchPromises.push({
|
||||
key: "io-intelligence",
|
||||
options: { provider: "io-intelligence", apiKey: ioIntelligenceApiKey },
|
||||
})
|
||||
}
|
||||
|
||||
// Don't fetch Ollama and LM Studio models by default anymore.
|
||||
// They have their own specific handlers: requestOllamaModels and requestLmStudioModels.
|
||||
|
||||
const litellmApiKey = apiConfiguration.litellmApiKey || message?.values?.litellmApiKey
|
||||
const litellmBaseUrl = apiConfiguration.litellmBaseUrl || message?.values?.litellmBaseUrl
|
||||
if (litellmApiKey && litellmBaseUrl) {
|
||||
modelFetchPromises.push({
|
||||
key: "litellm",
|
||||
options: { provider: "litellm", apiKey: litellmApiKey, baseUrl: litellmBaseUrl },
|
||||
})
|
||||
}
|
||||
|
||||
const results = await Promise.allSettled(
|
||||
modelFetchPromises.map(async ({ key, options }) => {
|
||||
const models = await safeGetModels(options)
|
||||
return { key, models }
|
||||
}),
|
||||
)
|
||||
|
||||
results.forEach((result, index) => {
|
||||
const routerName = modelFetchPromises[index].key
|
||||
|
||||
if (result.status === "fulfilled") {
|
||||
routerModels[routerName] = result.value.models
|
||||
|
||||
// Ollama and LM Studio settings pages still need these events.
|
||||
if (routerName === "ollama" && Object.keys(result.value.models).length > 0) {
|
||||
provider.postMessageToWebview({
|
||||
type: "ollamaModels",
|
||||
ollamaModels: result.value.models,
|
||||
})
|
||||
} else if (routerName === "lmstudio" && Object.keys(result.value.models).length > 0) {
|
||||
provider.postMessageToWebview({
|
||||
type: "lmStudioModels",
|
||||
lmStudioModels: result.value.models,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
// Handle rejection: Post a specific error message for this provider.
|
||||
const errorMessage = result.reason instanceof Error ? result.reason.message : String(result.reason)
|
||||
provider.log(`Error fetching models for ${routerName}: ${errorMessage}`)
|
||||
|
||||
routerModels[routerName] = {}
|
||||
|
||||
provider.postMessageToWebview({
|
||||
type: "singleRouterModelFetchResponse",
|
||||
success: false,
|
||||
error: errorMessage,
|
||||
values: { provider: routerName },
|
||||
})
|
||||
}
|
||||
const { routerModels, errors } = await fetchRouterModels({
|
||||
apiConfiguration,
|
||||
activeProviderOnly: false,
|
||||
litellmOverrides: message?.values
|
||||
? {
|
||||
apiKey: message.values.litellmApiKey,
|
||||
baseUrl: message.values.litellmBaseUrl,
|
||||
}
|
||||
: undefined,
|
||||
})
|
||||
|
||||
// Send error notifications for failed providers
|
||||
errors.forEach((err) => {
|
||||
provider.log(`Error fetching models for ${err.provider}: ${err.error}`)
|
||||
provider.postMessageToWebview({
|
||||
type: "singleRouterModelFetchResponse",
|
||||
success: false,
|
||||
error: err.error,
|
||||
values: { provider: err.provider },
|
||||
})
|
||||
})
|
||||
|
||||
// Send ollama/lmstudio-specific messages if models were fetched
|
||||
if (routerModels.ollama && Object.keys(routerModels.ollama).length > 0) {
|
||||
provider.postMessageToWebview({
|
||||
type: "ollamaModels",
|
||||
ollamaModels: routerModels.ollama,
|
||||
})
|
||||
}
|
||||
if (routerModels.lmstudio && Object.keys(routerModels.lmstudio).length > 0) {
|
||||
provider.postMessageToWebview({
|
||||
type: "lmStudioModels",
|
||||
lmStudioModels: routerModels.lmstudio,
|
||||
})
|
||||
}
|
||||
|
||||
provider.postMessageToWebview({ type: "routerModels", routerModels: routerModels as RouterModels })
|
||||
break
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,266 @@
|
|||
import { describe, it, expect, beforeEach, vi } from "vitest"
|
||||
import type { Mock } from "vitest"
|
||||
import type { ProviderSettings } from "@roo-code/types"
|
||||
import { fetchRouterModels } from "../index"
|
||||
import { getModels } from "../../../api/providers/fetchers/modelCache"
|
||||
import { CloudService } from "@roo-code/cloud"
|
||||
|
||||
// Mock dependencies
|
||||
vi.mock("../../../api/providers/fetchers/modelCache")
|
||||
vi.mock("@roo-code/cloud")
|
||||
|
||||
const mockGetModels = getModels as Mock<typeof getModels>
|
||||
const mockCloudService = CloudService as any
|
||||
|
||||
describe("RouterModelsService", () => {
|
||||
const mockModels = {
|
||||
"test-model": {
|
||||
maxTokens: 4096,
|
||||
contextWindow: 8192,
|
||||
supportsPromptCache: false,
|
||||
description: "Test model",
|
||||
},
|
||||
}
|
||||
|
||||
const baseApiConfiguration: ProviderSettings = {
|
||||
apiProvider: "openrouter",
|
||||
openRouterApiKey: "test-key",
|
||||
requestyApiKey: "requesty-key",
|
||||
unboundApiKey: "unbound-key",
|
||||
ioIntelligenceApiKey: "io-key",
|
||||
deepInfraApiKey: "deepinfra-key",
|
||||
litellmApiKey: "litellm-key",
|
||||
litellmBaseUrl: "http://localhost:4000",
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
mockGetModels.mockResolvedValue(mockModels)
|
||||
mockCloudService.hasInstance = vi.fn().mockReturnValue(false)
|
||||
})
|
||||
|
||||
describe("fetchRouterModels", () => {
|
||||
it("fetches all providers when activeProviderOnly is false", async () => {
|
||||
const result = await fetchRouterModels({
|
||||
apiConfiguration: baseApiConfiguration,
|
||||
activeProviderOnly: false,
|
||||
})
|
||||
|
||||
// Should fetch all standard providers
|
||||
expect(mockGetModels).toHaveBeenCalledWith({ provider: "openrouter" })
|
||||
expect(mockGetModels).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ provider: "requesty", apiKey: "requesty-key" }),
|
||||
)
|
||||
expect(mockGetModels).toHaveBeenCalledWith({ provider: "glama" })
|
||||
expect(mockGetModels).toHaveBeenCalledWith({ provider: "unbound", apiKey: "unbound-key" })
|
||||
expect(mockGetModels).toHaveBeenCalledWith({ provider: "vercel-ai-gateway" })
|
||||
expect(mockGetModels).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ provider: "deepinfra", apiKey: "deepinfra-key" }),
|
||||
)
|
||||
expect(mockGetModels).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
provider: "roo",
|
||||
baseUrl: "https://api.roocode.com/proxy",
|
||||
}),
|
||||
)
|
||||
expect(mockGetModels).toHaveBeenCalledWith({ provider: "io-intelligence", apiKey: "io-key" })
|
||||
expect(mockGetModels).toHaveBeenCalledWith({
|
||||
provider: "litellm",
|
||||
apiKey: "litellm-key",
|
||||
baseUrl: "http://localhost:4000",
|
||||
})
|
||||
|
||||
// Should return models for all providers
|
||||
expect(result.routerModels).toHaveProperty("openrouter")
|
||||
expect(result.routerModels).toHaveProperty("requesty")
|
||||
expect(result.routerModels).toHaveProperty("glama")
|
||||
expect(result.errors).toEqual([])
|
||||
})
|
||||
|
||||
it("fetches only active provider when activeProviderOnly is true", async () => {
|
||||
const result = await fetchRouterModels({
|
||||
apiConfiguration: { ...baseApiConfiguration, apiProvider: "openrouter" },
|
||||
activeProviderOnly: true,
|
||||
})
|
||||
|
||||
// Should only fetch openrouter
|
||||
expect(mockGetModels).toHaveBeenCalledTimes(1)
|
||||
expect(mockGetModels).toHaveBeenCalledWith({ provider: "openrouter" })
|
||||
|
||||
// Should return models only for openrouter
|
||||
expect(result.routerModels.openrouter).toEqual(mockModels)
|
||||
expect(result.errors).toEqual([])
|
||||
})
|
||||
|
||||
it("includes ollama when it is the active provider", async () => {
|
||||
const config: ProviderSettings = {
|
||||
...baseApiConfiguration,
|
||||
apiProvider: "ollama",
|
||||
ollamaBaseUrl: "http://localhost:11434",
|
||||
}
|
||||
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: config,
|
||||
activeProviderOnly: true,
|
||||
})
|
||||
|
||||
expect(mockGetModels).toHaveBeenCalledWith({
|
||||
provider: "ollama",
|
||||
baseUrl: "http://localhost:11434",
|
||||
apiKey: undefined,
|
||||
})
|
||||
})
|
||||
|
||||
it("includes lmstudio when it is the active provider", async () => {
|
||||
const config: ProviderSettings = {
|
||||
...baseApiConfiguration,
|
||||
apiProvider: "lmstudio",
|
||||
lmStudioBaseUrl: "http://localhost:1234",
|
||||
}
|
||||
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: config,
|
||||
activeProviderOnly: true,
|
||||
})
|
||||
|
||||
expect(mockGetModels).toHaveBeenCalledWith({
|
||||
provider: "lmstudio",
|
||||
baseUrl: "http://localhost:1234",
|
||||
})
|
||||
})
|
||||
|
||||
it("includes huggingface when it is the active provider", async () => {
|
||||
const config: ProviderSettings = {
|
||||
...baseApiConfiguration,
|
||||
apiProvider: "huggingface",
|
||||
}
|
||||
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: config,
|
||||
activeProviderOnly: true,
|
||||
})
|
||||
|
||||
expect(mockGetModels).toHaveBeenCalledWith({
|
||||
provider: "huggingface",
|
||||
})
|
||||
})
|
||||
|
||||
it("uses litellmOverrides when provided", async () => {
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: { ...baseApiConfiguration, litellmApiKey: undefined, litellmBaseUrl: undefined },
|
||||
activeProviderOnly: false,
|
||||
litellmOverrides: {
|
||||
apiKey: "override-key",
|
||||
baseUrl: "http://override:5000",
|
||||
},
|
||||
})
|
||||
|
||||
expect(mockGetModels).toHaveBeenCalledWith({
|
||||
provider: "litellm",
|
||||
apiKey: "override-key",
|
||||
baseUrl: "http://override:5000",
|
||||
})
|
||||
})
|
||||
|
||||
it("handles provider fetch errors gracefully", async () => {
|
||||
mockGetModels
|
||||
.mockResolvedValueOnce(mockModels) // openrouter succeeds
|
||||
.mockRejectedValueOnce(new Error("Requesty API error")) // requesty fails
|
||||
.mockResolvedValueOnce(mockModels) // glama succeeds
|
||||
|
||||
const result = await fetchRouterModels({
|
||||
apiConfiguration: baseApiConfiguration,
|
||||
activeProviderOnly: false,
|
||||
})
|
||||
|
||||
// Should have errors for failed provider
|
||||
expect(result.errors).toHaveLength(1)
|
||||
expect(result.errors[0]).toEqual({
|
||||
provider: "requesty",
|
||||
error: "Requesty API error",
|
||||
})
|
||||
|
||||
// Should have empty object for failed provider
|
||||
expect(result.routerModels.requesty).toEqual({})
|
||||
|
||||
// Should have models for successful providers
|
||||
expect(result.routerModels.openrouter).toEqual(mockModels)
|
||||
})
|
||||
|
||||
it("skips litellm when no api key or base url provided", async () => {
|
||||
const config: ProviderSettings = {
|
||||
...baseApiConfiguration,
|
||||
litellmApiKey: undefined,
|
||||
litellmBaseUrl: undefined,
|
||||
}
|
||||
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: config,
|
||||
activeProviderOnly: false,
|
||||
})
|
||||
|
||||
// Should not call getModels for litellm
|
||||
expect(mockGetModels).not.toHaveBeenCalledWith(expect.objectContaining({ provider: "litellm" }))
|
||||
})
|
||||
|
||||
it("skips io-intelligence when no api key provided", async () => {
|
||||
const config: ProviderSettings = {
|
||||
...baseApiConfiguration,
|
||||
ioIntelligenceApiKey: undefined,
|
||||
}
|
||||
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: config,
|
||||
activeProviderOnly: false,
|
||||
})
|
||||
|
||||
// Should not call getModels for io-intelligence
|
||||
expect(mockGetModels).not.toHaveBeenCalledWith(expect.objectContaining({ provider: "io-intelligence" }))
|
||||
})
|
||||
|
||||
it("uses roo session token when CloudService is available", async () => {
|
||||
const mockAuthService = {
|
||||
getSessionToken: vi.fn().mockReturnValue("session-token-123"),
|
||||
}
|
||||
|
||||
vi.mocked(CloudService.hasInstance).mockReturnValue(true)
|
||||
Object.defineProperty(CloudService, "instance", {
|
||||
get: () => ({ authService: mockAuthService }),
|
||||
configurable: true,
|
||||
})
|
||||
|
||||
await fetchRouterModels({
|
||||
apiConfiguration: baseApiConfiguration,
|
||||
activeProviderOnly: false,
|
||||
})
|
||||
|
||||
expect(mockGetModels).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
provider: "roo",
|
||||
apiKey: "session-token-123",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("initializes all providers with empty objects", async () => {
|
||||
const result = await fetchRouterModels({
|
||||
apiConfiguration: { apiProvider: "openrouter" } as ProviderSettings,
|
||||
activeProviderOnly: true,
|
||||
})
|
||||
|
||||
// All providers should be initialized even if not fetched
|
||||
expect(result.routerModels).toHaveProperty("openrouter")
|
||||
expect(result.routerModels).toHaveProperty("requesty")
|
||||
expect(result.routerModels).toHaveProperty("glama")
|
||||
expect(result.routerModels).toHaveProperty("unbound")
|
||||
expect(result.routerModels).toHaveProperty("vercel-ai-gateway")
|
||||
expect(result.routerModels).toHaveProperty("deepinfra")
|
||||
expect(result.routerModels).toHaveProperty("roo")
|
||||
expect(result.routerModels).toHaveProperty("litellm")
|
||||
expect(result.routerModels).toHaveProperty("ollama")
|
||||
expect(result.routerModels).toHaveProperty("lmstudio")
|
||||
expect(result.routerModels).toHaveProperty("huggingface")
|
||||
expect(result.routerModels).toHaveProperty("io-intelligence")
|
||||
})
|
||||
})
|
||||
})
|
||||
171
src/services/router-models/index.ts
Normal file
171
src/services/router-models/index.ts
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
import type { ProviderSettings } from "@roo-code/types"
|
||||
import { CloudService } from "@roo-code/cloud"
|
||||
import type { RouterName, ModelRecord, GetModelsOptions } from "../../shared/api"
|
||||
import { getModels } from "../../api/providers/fetchers/modelCache"
|
||||
|
||||
export interface RouterModelsFetchOptions {
|
||||
apiConfiguration: ProviderSettings
|
||||
activeProviderOnly?: boolean
|
||||
litellmOverrides?: {
|
||||
apiKey?: string
|
||||
baseUrl?: string
|
||||
}
|
||||
}
|
||||
|
||||
export interface RouterModelsFetchResult {
|
||||
routerModels: Partial<Record<RouterName, ModelRecord>>
|
||||
errors: Array<{
|
||||
provider: RouterName
|
||||
error: string
|
||||
}>
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds the list of provider fetch options based on configuration and mode.
|
||||
*/
|
||||
function buildProviderFetchList(
|
||||
options: RouterModelsFetchOptions,
|
||||
): Array<{ key: RouterName; options: GetModelsOptions }> {
|
||||
const { apiConfiguration, activeProviderOnly, litellmOverrides } = options
|
||||
|
||||
const allFetches: Array<{ key: RouterName; options: GetModelsOptions }> = [
|
||||
{ key: "openrouter", options: { provider: "openrouter" } },
|
||||
{
|
||||
key: "requesty",
|
||||
options: {
|
||||
provider: "requesty",
|
||||
apiKey: apiConfiguration.requestyApiKey,
|
||||
baseUrl: apiConfiguration.requestyBaseUrl,
|
||||
},
|
||||
},
|
||||
{ key: "glama", options: { provider: "glama" } },
|
||||
{ key: "unbound", options: { provider: "unbound", apiKey: apiConfiguration.unboundApiKey } },
|
||||
{ key: "vercel-ai-gateway", options: { provider: "vercel-ai-gateway" } },
|
||||
{
|
||||
key: "deepinfra",
|
||||
options: {
|
||||
provider: "deepinfra",
|
||||
apiKey: apiConfiguration.deepInfraApiKey,
|
||||
baseUrl: apiConfiguration.deepInfraBaseUrl,
|
||||
},
|
||||
},
|
||||
{
|
||||
key: "roo",
|
||||
options: {
|
||||
provider: "roo",
|
||||
baseUrl: process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy",
|
||||
apiKey: CloudService.hasInstance() ? CloudService.instance.authService?.getSessionToken() : undefined,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
// Include local providers when in active-provider mode and they are selected
|
||||
if (activeProviderOnly) {
|
||||
const activeProvider = apiConfiguration.apiProvider
|
||||
|
||||
if (activeProvider === "ollama") {
|
||||
allFetches.push({
|
||||
key: "ollama",
|
||||
options: {
|
||||
provider: "ollama",
|
||||
baseUrl: apiConfiguration.ollamaBaseUrl,
|
||||
apiKey: apiConfiguration.ollamaApiKey,
|
||||
},
|
||||
})
|
||||
}
|
||||
if (activeProvider === "lmstudio") {
|
||||
allFetches.push({
|
||||
key: "lmstudio",
|
||||
options: {
|
||||
provider: "lmstudio",
|
||||
baseUrl: apiConfiguration.lmStudioBaseUrl,
|
||||
},
|
||||
})
|
||||
}
|
||||
if (activeProvider === "huggingface") {
|
||||
allFetches.push({
|
||||
key: "huggingface",
|
||||
options: {
|
||||
provider: "huggingface",
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Add IO Intelligence if API key is provided
|
||||
if (apiConfiguration.ioIntelligenceApiKey) {
|
||||
allFetches.push({
|
||||
key: "io-intelligence",
|
||||
options: { provider: "io-intelligence", apiKey: apiConfiguration.ioIntelligenceApiKey },
|
||||
})
|
||||
}
|
||||
|
||||
// Add LiteLLM if configured (with potential overrides from message)
|
||||
const litellmApiKey = apiConfiguration.litellmApiKey || litellmOverrides?.apiKey
|
||||
const litellmBaseUrl = apiConfiguration.litellmBaseUrl || litellmOverrides?.baseUrl
|
||||
if (litellmApiKey && litellmBaseUrl) {
|
||||
allFetches.push({
|
||||
key: "litellm",
|
||||
options: { provider: "litellm", apiKey: litellmApiKey, baseUrl: litellmBaseUrl },
|
||||
})
|
||||
}
|
||||
|
||||
return allFetches
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetches router models based on the provided options.
|
||||
* Can fetch all providers or only the active provider.
|
||||
*/
|
||||
export async function fetchRouterModels(options: RouterModelsFetchOptions): Promise<RouterModelsFetchResult> {
|
||||
const { apiConfiguration, activeProviderOnly } = options
|
||||
|
||||
// Initialize empty results for all providers
|
||||
const routerModels: Partial<Record<RouterName, ModelRecord>> = {
|
||||
openrouter: {},
|
||||
"vercel-ai-gateway": {},
|
||||
huggingface: {},
|
||||
litellm: {},
|
||||
deepinfra: {},
|
||||
"io-intelligence": {},
|
||||
requesty: {},
|
||||
unbound: {},
|
||||
glama: {},
|
||||
ollama: {},
|
||||
lmstudio: {},
|
||||
roo: {},
|
||||
}
|
||||
|
||||
const errors: Array<{ provider: RouterName; error: string }> = []
|
||||
|
||||
// Build fetch list
|
||||
const fetchList = buildProviderFetchList(options)
|
||||
|
||||
// Filter to active provider if requested
|
||||
const activeProvider = apiConfiguration.apiProvider as RouterName | undefined
|
||||
const modelFetchPromises =
|
||||
activeProviderOnly && activeProvider ? fetchList.filter(({ key }) => key === activeProvider) : fetchList
|
||||
|
||||
// Execute fetches
|
||||
const results = await Promise.allSettled(
|
||||
modelFetchPromises.map(async ({ key, options }) => {
|
||||
const models = await getModels(options)
|
||||
return { key, models }
|
||||
}),
|
||||
)
|
||||
|
||||
// Process results
|
||||
results.forEach((result, index) => {
|
||||
const routerName = modelFetchPromises[index].key
|
||||
|
||||
if (result.status === "fulfilled") {
|
||||
routerModels[routerName] = result.value.models
|
||||
} else {
|
||||
const errorMessage = result.reason instanceof Error ? result.reason.message : String(result.reason)
|
||||
routerModels[routerName] = {}
|
||||
errors.push({ provider: routerName, error: errorMessage })
|
||||
}
|
||||
})
|
||||
|
||||
return { routerModels, errors }
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue