Skip to content

Commit 569b276

Browse files
andrewshu2000Brendan-Zdaniel-lxs
authored
feat: adding default headers and testing for litellm fetcher (RooCodeInc#5242)
* chore: adding x-title header and testing for litellm * chore: indentation fi and headers order fix * chore: spacing fix * chore: removed white space * fix: allow user headers to override default headers and clean up formatting - Reorder header spread in router-provider.ts so user-provided openAiHeaders can override DEFAULT_HEADERS - Remove unnecessary blank lines after imports for consistency - This matches the pattern used in openai.ts where DEFAULT_HEADERS come first --------- Co-authored-by: Brendan-Z <[email protected]> Co-authored-by: Daniel Riccio <[email protected]>
1 parent 18c3f5d commit 569b276

File tree

3 files changed

+16
-1
lines changed

3 files changed

+16
-1
lines changed

src/api/providers/fetchers/__tests__/litellm.spec.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ vi.mock("axios")
44
import type { Mock } from "vitest"
55
import axios from "axios"
66
import { getLiteLLMModels } from "../litellm"
7+
import { DEFAULT_HEADERS } from "../../constants"
78

89
const mockedAxios = axios as typeof axios & {
910
get: Mock
@@ -32,6 +33,7 @@ describe("getLiteLLMModels", () => {
3233
headers: {
3334
Authorization: "Bearer test-api-key",
3435
"Content-Type": "application/json",
36+
...DEFAULT_HEADERS,
3537
},
3638
timeout: 5000,
3739
})
@@ -83,6 +85,7 @@ describe("getLiteLLMModels", () => {
8385
headers: {
8486
Authorization: "Bearer test-api-key",
8587
"Content-Type": "application/json",
88+
...DEFAULT_HEADERS,
8689
},
8790
timeout: 5000,
8891
})
@@ -125,6 +128,7 @@ describe("getLiteLLMModels", () => {
125128
expect(mockedAxios.get).toHaveBeenCalledWith("http://localhost:4000/v1/model/info", {
126129
headers: {
127130
"Content-Type": "application/json",
131+
...DEFAULT_HEADERS,
128132
},
129133
timeout: 5000,
130134
})

src/api/providers/fetchers/litellm.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { LITELLM_COMPUTER_USE_MODELS } from "@roo-code/types"
44

55
import type { ModelRecord } from "../../../shared/api"
66

7+
import { DEFAULT_HEADERS } from "../constants"
78
/**
89
* Fetches available models from a LiteLLM server
910
*
@@ -16,6 +17,7 @@ export async function getLiteLLMModels(apiKey: string, baseUrl: string): Promise
1617
try {
1718
const headers: Record<string, string> = {
1819
"Content-Type": "application/json",
20+
...DEFAULT_HEADERS,
1921
}
2022

2123
if (apiKey) {

src/api/providers/router-provider.ts

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import { ApiHandlerOptions, RouterName, ModelRecord } from "../../shared/api"
77
import { BaseProvider } from "./base-provider"
88
import { getModels } from "./fetchers/modelCache"
99

10+
import { DEFAULT_HEADERS } from "./constants"
11+
1012
type RouterProviderOptions = {
1113
name: RouterName
1214
baseURL: string
@@ -43,7 +45,14 @@ export abstract class RouterProvider extends BaseProvider {
4345
this.defaultModelId = defaultModelId
4446
this.defaultModelInfo = defaultModelInfo
4547

46-
this.client = new OpenAI({ baseURL, apiKey })
48+
this.client = new OpenAI({
49+
baseURL,
50+
apiKey,
51+
defaultHeaders: {
52+
...DEFAULT_HEADERS,
53+
...(options.openAiHeaders || {}),
54+
},
55+
})
4756
}
4857

4958
public async fetchModel() {

0 commit comments

Comments
 (0)