feat(ask_sb): Add headers param to config to allow users to specify custom headers (#449)

This commit is contained in:
Brendan Kellam 2025-08-08 14:49:00 -07:00 committed by GitHub
parent a9a61e7338
commit 111e1c3cee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 4720 additions and 192 deletions

View file

@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Added
- Added multi-branch indexing support for Gerrit. [#433](https://github.com/sourcebot-dev/sourcebot/pull/433)
- [ask sb] Added `reasoningEffort` option to OpenAI provider. [#446](https://github.com/sourcebot-dev/sourcebot/pull/446)
- [ask db] Added `headers` option to all providers. [#449](https://github.com/sourcebot-dev/sourcebot/pull/449)
### Fixed
- Removed prefix from structured log output. [#443](https://github.com/sourcebot-dev/sourcebot/pull/443)

View file

@ -324,8 +324,34 @@ The OpenAI compatible provider allows you to use any model that is compatible wi
}
```
# Custom headers
## Schema reference
You can pass custom headers to the language model provider by using the `headers` parameter. Header values can either be a string or a environment variable. Headers are supported for all providers.
```json wrap icon="code" Example config with custom headers
{
"$schema": "https://raw.githubusercontent.com/sourcebot-dev/sourcebot/main/schemas/v3/index.json",
"models": [
{
// ... provider, model, displayName, etc...
// Key-value pairs of headers
"headers": {
// Header values can be passed as a environment variable...
"my-secret-header": {
"env": "MY_SECRET_HEADER_ENV_VAR"
},
// ... or directly as a string.
"my-non-secret-header": "plaintextvalue"
}
}
]
}
```
# Schema reference
<Accordion title="Reference">
[schemas/v3/languageModel.json](https://github.com/sourcebot-dev/sourcebot/blob/main/schemas/v3/languageModel.json)

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -74,6 +74,50 @@
}
},
"additionalProperties": false
},
"LanguageModelHeaders": {
"type": "object",
"description": "Optional headers to use with the model.",
"patternProperties": {
"^[!#$%&'*+\\-.^_`|~0-9A-Za-z]+$": {
"anyOf": [
{
"type": "string"
},
{
"anyOf": [
{
"type": "object",
"properties": {
"secret": {
"type": "string",
"description": "The name of the secret that contains the token."
}
},
"required": [
"secret"
],
"additionalProperties": false
},
{
"type": "object",
"properties": {
"env": {
"type": "string",
"description": "The name of the environment variable that contains the token. Only supported in declarative connection configs."
}
},
"required": [
"env"
],
"additionalProperties": false
}
]
}
]
}
},
"additionalProperties": false
}
}
}

File diff suppressed because it is too large Load diff

View file

@ -496,6 +496,32 @@ export interface AmazonBedrockLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
/**
* Optional headers to use with the model.
*/
export interface LanguageModelHeaders {
/**
* This interface was referenced by `LanguageModelHeaders`'s JSON-Schema definition
* via the `patternProperty` "^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$".
*/
[k: string]:
| string
| (
| {
/**
* The name of the secret that contains the token.
*/
secret: string;
}
| {
/**
* The name of the environment variable that contains the token. Only supported in declarative connection configs.
*/
env: string;
}
);
}
export interface AnthropicLanguageModel {
/**
@ -530,6 +556,7 @@ export interface AnthropicLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface AzureLanguageModel {
/**
@ -572,6 +599,7 @@ export interface AzureLanguageModel {
* Use a different URL prefix for API calls. Either this or `resourceName` can be used.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface DeepSeekLanguageModel {
/**
@ -606,6 +634,7 @@ export interface DeepSeekLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface GoogleGenerativeAILanguageModel {
/**
@ -640,6 +669,7 @@ export interface GoogleGenerativeAILanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface GoogleVertexAnthropicLanguageModel {
/**
@ -682,6 +712,7 @@ export interface GoogleVertexAnthropicLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface GoogleVertexLanguageModel {
/**
@ -724,6 +755,7 @@ export interface GoogleVertexLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface MistralLanguageModel {
/**
@ -758,6 +790,7 @@ export interface MistralLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface OpenAILanguageModel {
/**
@ -796,6 +829,7 @@ export interface OpenAILanguageModel {
* The reasoning effort to use with the model. Defaults to `medium`. See https://platform.openai.com/docs/guides/reasoning#get-started-with-reasonings
*/
reasoningEffort?: string;
headers?: LanguageModelHeaders;
}
export interface OpenAICompatibleLanguageModel {
/**
@ -830,6 +864,7 @@ export interface OpenAICompatibleLanguageModel {
* Base URL of the OpenAI-compatible chat completions API endpoint.
*/
baseUrl: string;
headers?: LanguageModelHeaders;
}
export interface OpenRouterLanguageModel {
/**
@ -864,6 +899,7 @@ export interface OpenRouterLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface XaiLanguageModel {
/**
@ -898,4 +934,5 @@ export interface XaiLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}

File diff suppressed because it is too large Load diff

View file

@ -67,6 +67,32 @@ export interface AmazonBedrockLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
/**
* Optional headers to use with the model.
*/
export interface LanguageModelHeaders {
/**
* This interface was referenced by `LanguageModelHeaders`'s JSON-Schema definition
* via the `patternProperty` "^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$".
*/
[k: string]:
| string
| (
| {
/**
* The name of the secret that contains the token.
*/
secret: string;
}
| {
/**
* The name of the environment variable that contains the token. Only supported in declarative connection configs.
*/
env: string;
}
);
}
export interface AnthropicLanguageModel {
/**
@ -101,6 +127,7 @@ export interface AnthropicLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface AzureLanguageModel {
/**
@ -143,6 +170,7 @@ export interface AzureLanguageModel {
* Use a different URL prefix for API calls. Either this or `resourceName` can be used.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface DeepSeekLanguageModel {
/**
@ -177,6 +205,7 @@ export interface DeepSeekLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface GoogleGenerativeAILanguageModel {
/**
@ -211,6 +240,7 @@ export interface GoogleGenerativeAILanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface GoogleVertexAnthropicLanguageModel {
/**
@ -253,6 +283,7 @@ export interface GoogleVertexAnthropicLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface GoogleVertexLanguageModel {
/**
@ -295,6 +326,7 @@ export interface GoogleVertexLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface MistralLanguageModel {
/**
@ -329,6 +361,7 @@ export interface MistralLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface OpenAILanguageModel {
/**
@ -367,6 +400,7 @@ export interface OpenAILanguageModel {
* The reasoning effort to use with the model. Defaults to `medium`. See https://platform.openai.com/docs/guides/reasoning#get-started-with-reasonings
*/
reasoningEffort?: string;
headers?: LanguageModelHeaders;
}
export interface OpenAICompatibleLanguageModel {
/**
@ -401,6 +435,7 @@ export interface OpenAICompatibleLanguageModel {
* Base URL of the OpenAI-compatible chat completions API endpoint.
*/
baseUrl: string;
headers?: LanguageModelHeaders;
}
export interface OpenRouterLanguageModel {
/**
@ -435,6 +470,7 @@ export interface OpenRouterLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}
export interface XaiLanguageModel {
/**
@ -469,4 +505,5 @@ export interface XaiLanguageModel {
* Optional base URL.
*/
baseUrl?: string;
headers?: LanguageModelHeaders;
}

View file

@ -73,6 +73,50 @@ const schema = {
}
},
"additionalProperties": false
},
"LanguageModelHeaders": {
"type": "object",
"description": "Optional headers to use with the model.",
"patternProperties": {
"^[!#$%&'*+\\-.^_`|~0-9A-Za-z]+$": {
"anyOf": [
{
"type": "string"
},
{
"anyOf": [
{
"type": "object",
"properties": {
"secret": {
"type": "string",
"description": "The name of the secret that contains the token."
}
},
"required": [
"secret"
],
"additionalProperties": false
},
{
"type": "object",
"properties": {
"env": {
"type": "string",
"description": "The name of the environment variable that contains the token. Only supported in declarative connection configs."
}
},
"required": [
"env"
],
"additionalProperties": false
}
]
}
]
}
},
"additionalProperties": false
}
}
} as const;

View file

@ -37,3 +37,16 @@ export interface GitRevisions {
*/
tags?: string[];
}
/**
* Optional headers to use with the model.
*
* This interface was referenced by `Shared`'s JSON-Schema
* via the `definition` "LanguageModelHeaders".
*/
export interface LanguageModelHeaders {
/**
* This interface was referenced by `LanguageModelHeaders`'s JSON-Schema definition
* via the `patternProperty` "^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$".
*/
[k: string]: string | Token;
}

View file

@ -88,7 +88,7 @@ export async function POST(req: Request) {
});
}
const { model, providerOptions, headers } = await _getAISDKLanguageModelAndOptions(languageModelConfig, org.id);
const { model, providerOptions } = await _getAISDKLanguageModelAndOptions(languageModelConfig, org.id);
return createMessageStreamResponse({
messages,
@ -97,7 +97,6 @@ export async function POST(req: Request) {
model,
modelName: languageModelConfig.displayName ?? languageModelConfig.model,
modelProviderOptions: providerOptions,
modelHeaders: headers,
domain,
orgId: org.id,
});
@ -129,7 +128,6 @@ interface CreateMessageStreamResponseProps {
model: AISDKLanguageModelV2;
modelName: string;
modelProviderOptions?: Record<string, Record<string, JSONValue>>;
modelHeaders?: Record<string, string>;
domain: string;
orgId: number;
}
@ -141,7 +139,6 @@ const createMessageStreamResponse = async ({
model,
modelName,
modelProviderOptions,
modelHeaders,
domain,
orgId,
}: CreateMessageStreamResponseProps) => {
@ -210,7 +207,6 @@ const createMessageStreamResponse = async ({
const researchStream = await createAgentStream({
model,
providerOptions: modelProviderOptions,
headers: modelHeaders,
inputMessages: messageHistory,
inputSources: sources,
searchScopeRepoNames: expandedRepos,

View file

@ -20,8 +20,8 @@ import { LanguageModelV2 as AISDKLanguageModelV2 } from "@ai-sdk/provider";
import { createXai } from '@ai-sdk/xai';
import { createOpenRouter } from '@openrouter/ai-sdk-provider';
import { getTokenFromConfig } from "@sourcebot/crypto";
import { ChatVisibility, OrgRole, Prisma } from "@sourcebot/db";
import { LanguageModel } from "@sourcebot/schemas/v3/languageModel.type";
import { ChatVisibility, OrgRole, Prisma, PrismaClient } from "@sourcebot/db";
import { LanguageModel, LanguageModelHeaders } from "@sourcebot/schemas/v3/languageModel.type";
import { loadConfig } from "@sourcebot/shared";
import { generateText, JSONValue } from "ai";
import fs from 'fs';
@ -375,7 +375,6 @@ export const _getConfiguredLanguageModelsFull = async (): Promise<LanguageModel[
export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, orgId: number): Promise<{
model: AISDKLanguageModelV2,
providerOptions?: Record<string, Record<string, JSONValue>>,
headers?: Record<string, string>,
}> => {
const { provider, model: modelId } = config;
@ -390,6 +389,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
secretAccessKey: config.accessKeySecret
? await getTokenFromConfig(config.accessKeySecret, orgId, prisma)
: env.AWS_SECRET_ACCESS_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -402,6 +404,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: env.ANTHROPIC_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -414,10 +419,6 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
}
} satisfies AnthropicProviderOptions,
},
headers: {
// @see: https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#interleaved-thinking
'anthropic-beta': 'interleaved-thinking-2025-05-14',
},
};
}
case 'azure': {
@ -426,6 +427,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token ? (await getTokenFromConfig(config.token, orgId, prisma)) : env.AZURE_API_KEY,
apiVersion: config.apiVersion,
resourceName: config.resourceName ?? env.AZURE_RESOURCE_NAME,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -436,6 +440,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
const deepseek = createDeepSeek({
baseURL: config.baseUrl,
apiKey: config.token ? (await getTokenFromConfig(config.token, orgId, prisma)) : env.DEEPSEEK_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -448,6 +455,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: env.GOOGLE_GENERATIVE_AI_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -463,6 +473,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
keyFilename: await getTokenFromConfig(config.credentials, orgId, prisma),
}
} : {}),
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -486,6 +499,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
keyFilename: await getTokenFromConfig(config.credentials, orgId, prisma),
}
} : {}),
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -498,6 +514,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: env.MISTRAL_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -510,6 +529,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: env.OPENAI_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -528,6 +550,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: undefined,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -540,6 +565,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: env.OPENROUTER_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -552,6 +580,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
apiKey: config.token
? await getTokenFromConfig(config.token, orgId, prisma)
: env.XAI_API_KEY,
headers: config.headers
? await extractLanguageModelHeaders(config.headers, orgId, prisma)
: undefined,
});
return {
@ -559,4 +590,28 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or
};
}
}
}
}
const extractLanguageModelHeaders = async (
headers: LanguageModelHeaders,
orgId: number,
db: PrismaClient,
): Promise<Record<string, string>> => {
const resolvedHeaders: Record<string, string> = {};
if (!headers) {
return resolvedHeaders;
}
for (const [headerName, headerValue] of Object.entries(headers)) {
if (typeof headerValue === "string") {
resolvedHeaders[headerName] = headerValue;
continue;
}
const value = await getTokenFromConfig(headerValue, orgId, db);
resolvedHeaders[headerName] = value;
}
return resolvedHeaders;
}

View file

@ -32,7 +32,6 @@ const stepCountIsGTE = (stepCount: number): StopCondition<any> => {
export const createAgentStream = async ({
model,
providerOptions,
headers,
inputMessages,
inputSources,
searchScopeRepoNames,
@ -46,7 +45,6 @@ export const createAgentStream = async ({
const stream = streamText({
model,
providerOptions,
headers,
system: baseSystemPrompt,
messages: inputMessages,
tools: {

View file

@ -39,6 +39,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -71,6 +74,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -111,6 +117,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Use a different URL prefix for API calls. Either this or `resourceName` can be used."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -143,6 +152,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -175,6 +187,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -223,6 +238,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -273,6 +291,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -305,6 +326,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -353,6 +377,9 @@
"medium",
"high"
]
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -388,6 +415,9 @@
"examples": [
"http://localhost:8080/v1"
]
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -421,6 +451,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [
@ -457,6 +490,9 @@
"format": "url",
"pattern": "^https?:\\/\\/[^\\s/$.?#].[^\\s]*$",
"description": "Optional base URL."
},
"headers": {
"$ref": "./shared.json#/definitions/LanguageModelHeaders"
}
},
"required": [

View file

@ -72,6 +72,23 @@
}
},
"additionalProperties": false
},
"LanguageModelHeaders": {
"type": "object",
"description": "Optional headers to use with the model.",
"patternProperties": {
"^[!#$%&'*+\\-.^_`|~0-9A-Za-z]+$": {
"anyOf": [
{
"type": "string"
},
{
"$ref": "#/definitions/Token"
}
]
}
},
"additionalProperties": false
}
}
}