Skip to content

Add support for LLM routing using developer preferences #6075

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 22 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/types/src/global-settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ export const SECRET_STATE_KEYS = [
"xaiApiKey",
"groqApiKey",
"chutesApiKey",
"archgwApiKey",
"litellmApiKey",
"codeIndexOpenAiKey",
"codeIndexQdrantApiKey",
Expand Down
12 changes: 12 additions & 0 deletions packages/types/src/provider-settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ export const providerNames = [
"mistral",
"moonshot",
"deepseek",
"archgw",
"unbound",
"requesty",
"human-relay",
Expand Down Expand Up @@ -195,6 +196,14 @@ const moonshotSchema = apiModelIdProviderModelSchema.extend({
moonshotApiKey: z.string().optional(),
})

const archgwSchema = apiModelIdProviderModelSchema.extend({
archgwBaseUrl: z.string().optional(),
archgwApiKey: z.string().optional(),
archgwModelId: z.string().optional(),
archgwUsePreferences: z.boolean().optional(),
archgwPreferenceConfig: z.string().optional(),
})

const unboundSchema = baseProviderSettingsSchema.extend({
unboundApiKey: z.string().optional(),
unboundModelId: z.string().optional(),
Expand Down Expand Up @@ -250,6 +259,7 @@ export const providerSettingsSchemaDiscriminated = z.discriminatedUnion("apiProv
mistralSchema.merge(z.object({ apiProvider: z.literal("mistral") })),
deepSeekSchema.merge(z.object({ apiProvider: z.literal("deepseek") })),
moonshotSchema.merge(z.object({ apiProvider: z.literal("moonshot") })),
archgwSchema.merge(z.object({ apiProvider: z.literal("archgw") })),
unboundSchema.merge(z.object({ apiProvider: z.literal("unbound") })),
requestySchema.merge(z.object({ apiProvider: z.literal("requesty") })),
humanRelaySchema.merge(z.object({ apiProvider: z.literal("human-relay") })),
Expand Down Expand Up @@ -279,6 +289,7 @@ export const providerSettingsSchema = z.object({
...mistralSchema.shape,
...deepSeekSchema.shape,
...moonshotSchema.shape,
...archgwSchema.shape,
...unboundSchema.shape,
...requestySchema.shape,
...humanRelaySchema.shape,
Expand All @@ -304,6 +315,7 @@ export const MODEL_ID_KEYS: Partial<keyof ProviderSettings>[] = [
"unboundModelId",
"requestyModelId",
"litellmModelId",
"archgwModelId",
]

export const getModelId = (settings: ProviderSettings): string | undefined => {
Expand Down
13 changes: 13 additions & 0 deletions packages/types/src/providers/archgw.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import type { ModelInfo } from "../model.js"

export const archgwDefaultModelId = "openai/gpt-4.1"

export const archgwDefaultModelInfo: ModelInfo = {
maxTokens: 32_768,
contextWindow: 1_047_576,
supportsImages: true,
supportsPromptCache: true,
inputPrice: 2,
outputPrice: 8,
cacheReadsPrice: 0.5,
}
1 change: 1 addition & 0 deletions packages/types/src/providers/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,3 +18,4 @@ export * from "./unbound.js"
export * from "./vertex.js"
export * from "./vscode-llm.js"
export * from "./xai.js"
export * from "./archgw.js"
5 changes: 4 additions & 1 deletion pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 5 additions & 5 deletions scripts/update-contributors.js
Original file line number Diff line number Diff line change
Expand Up @@ -183,14 +183,14 @@ async function readReadme() {
* @param {Array} contributors Array of contributor objects from GitHub API
* @returns {string} HTML for contributors section
*/
const EXCLUDED_LOGIN_SUBSTRINGS = ['[bot]', 'R00-B0T'];
const EXCLUDED_LOGIN_EXACTS = ['cursor', 'roomote'];
const EXCLUDED_LOGIN_SUBSTRINGS = ["[bot]", "R00-B0T"]
const EXCLUDED_LOGIN_EXACTS = ["cursor", "roomote"]

function formatContributorsSection(contributors) {
// Filter out GitHub Actions bot, cursor, and roomote
const filteredContributors = contributors.filter((c) =>
!EXCLUDED_LOGIN_SUBSTRINGS.some(sub => c.login.includes(sub)) &&
!EXCLUDED_LOGIN_EXACTS.includes(c.login)
const filteredContributors = contributors.filter(
(c) =>
!EXCLUDED_LOGIN_SUBSTRINGS.some((sub) => c.login.includes(sub)) && !EXCLUDED_LOGIN_EXACTS.includes(c.login),
)

// Start building with Markdown table format
Expand Down
3 changes: 3 additions & 0 deletions src/api/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import {
GroqHandler,
ChutesHandler,
LiteLLMHandler,
ArchGwHandler,
ClaudeCodeHandler,
} from "./providers"

Expand Down Expand Up @@ -92,6 +93,8 @@ export function buildApiHandler(configuration: ProviderSettings): ApiHandler {
return new DeepSeekHandler(options)
case "moonshot":
return new MoonshotHandler(options)
case "archgw":
return new ArchGwHandler(options)
case "vscode-lm":
return new VsCodeLmHandler(options)
case "mistral":
Expand Down
214 changes: 214 additions & 0 deletions src/api/providers/__tests__/archgw.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,214 @@
// npx vitest run src/api/providers/__tests__/archgw.spec.ts

import { Anthropic } from "@anthropic-ai/sdk"
import { ArchGwHandler } from "../archgw"
import { ApiHandlerOptions } from "../../../shared/api"

const mockCreate = vitest.fn()

vitest.mock("openai", () => {
return {
__esModule: true,
default: vitest.fn().mockImplementation(() => ({
chat: {
completions: {
create: mockCreate.mockImplementation((options) => {
if (!options.stream) {
return Promise.resolve({
id: "test-completion",
choices: [
{
message: { role: "assistant", content: "ArchGW response" },
finish_reason: "stop",
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
},
})
}

// Streaming: return an object with withResponse() method
return {
withResponse: () => ({
data: {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [
{
delta: { content: "ArchGW stream" },
index: 0,
},
],
usage: null,
}
yield {
choices: [
{
delta: {},
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
},
}
},
},
}),
}
}),
},
},
})),
}
})

describe("ArchGwHandler", () => {
let handler: ArchGwHandler
let mockOptions: ApiHandlerOptions

beforeEach(() => {
mockOptions = {
archgwModelId: "arch-model",
archgwBaseUrl: "http://localhost:12000/v1",
archgwApiKey: "test-key",
archgwPreferenceConfig: "test-pref",
}
handler = new ArchGwHandler(mockOptions)
mockCreate.mockClear()
})

describe("constructor", () => {
it("should initialize with provided options", () => {
expect(handler).toBeInstanceOf(ArchGwHandler)
expect(handler.preferenceConfig).toBe("test-pref")
})

it("should use default base URL if not provided", () => {
const handlerWithoutUrl = new ArchGwHandler({
archgwModelId: "arch-model",
archgwApiKey: "test-key",
})
expect(handlerWithoutUrl).toBeInstanceOf(ArchGwHandler)
})

it("should not set preferenceConfig if not provided", () => {
const handlerNoPref = new ArchGwHandler({
archgwModelId: "arch-model",
archgwApiKey: "test-key",
})
expect(handlerNoPref.preferenceConfig).toBeUndefined()
})

it("should set preferenceConfig to empty string if provided as empty", () => {
const handlerEmptyPref = new ArchGwHandler({
archgwModelId: "arch-model",
archgwApiKey: "test-key",
archgwPreferenceConfig: "",
})
expect(handlerEmptyPref.preferenceConfig).toBe("")
})
})

describe("createMessage", () => {
const systemPrompt = "You are a helpful assistant."
const messages: Anthropic.Messages.MessageParam[] = [
{
role: "user",
content: "Hello!",
},
]

it("should handle streaming responses", async () => {
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
const textChunks = chunks.filter((chunk) => chunk.type === "text")
expect(textChunks).toHaveLength(1)
expect(textChunks[0].text).toBe("ArchGW stream")
})

it("should handle API errors", async () => {
mockCreate.mockImplementationOnce(() => ({
withResponse: () => {
throw new Error("API Error")
},
}))

const stream = handler.createMessage(systemPrompt, messages)

await expect(async () => {
for await (const _chunk of stream) {
// Should not reach here
}
}).rejects.toThrow("archgw streaming error: API Error")
})
})

describe("completePrompt", () => {
it("should complete prompt successfully", async () => {
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("ArchGW response")
expect(mockCreate).toHaveBeenCalledWith({
model: mockOptions.archgwModelId,
messages: [{ role: "user", content: "Test prompt" }],
temperature: 0,
max_tokens: 32768,
metadata: { archgw_preference_config: "test-pref" },
})
})

it("should not include archgw_preference_config in metadata if preferenceConfig is not set", async () => {
const handlerNoPref = new ArchGwHandler({
archgwModelId: "arch-model",
archgwApiKey: "test-key",
})
mockCreate.mockClear()
await handlerNoPref.completePrompt("Test prompt")
const call = mockCreate.mock.calls[0][0]
expect(call.metadata).toBeUndefined()
})

it("should include archgw_preference_config in metadata if preferenceConfig is empty string", async () => {
const preferenceConfig = `- name: code generation
model: gpt-4o-mini
usage: generating new code snippets
- name: code understanding
model: gpt-4.1
usage: understand and explain existing code snippets
`
const handlerEmptyPref = new ArchGwHandler({
archgwModelId: "arch-model",
archgwApiKey: "test-key",
archgwPreferenceConfig: preferenceConfig,
})
mockCreate.mockClear()
await handlerEmptyPref.completePrompt("Test prompt")
const call = mockCreate.mock.calls[0][0]
expect(call.metadata).toEqual({ archgw_preference_config: preferenceConfig })
})

it("should handle API errors", async () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))
await expect(handler.completePrompt("Test prompt")).rejects.toThrow("archgw completion error: API Error")
})

it("should handle empty response", async () => {
mockCreate.mockResolvedValueOnce({
choices: [{ message: { content: "" } }],
})
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("")
})
})
})
Loading