Skip to content

Commit c3cf085

Browse files
brendan-kellambrendan-kellam
committed
add mechanism for specifying language model explicitly
1 parent 2f1b008 commit c3cf085

9 files changed

Lines changed: 143 additions & 20 deletions

File tree

‎docs/docs/features/mcp-server.mdx‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -206,15 +206,25 @@ Parameters:
206206
| `perPage` | no | Results per page for pagination (min 1, max 100, default: 50). |
207207

208208

209+
### `list_language_models`
210+
211+
Lists the available language models configured on the Sourcebot instance. Use this to discover which models can be specified when calling `ask_codebase`.
212+
213+
Parameters:
214+
215+
This tool takes no parameters.
216+
217+
209218
### `ask_codebase`
210219

211220
Ask a natural language question about the codebase. This tool uses an AI agent to autonomously search code, read files, and find symbol references/definitions to answer your question. Returns a detailed answer in markdown format with code references, plus a link to view the full research session in the Sourcebot web UI.
212221

213222
Parameters:
214-
| Name | Required | Description |
215-
|:--------|:---------|:------------------------------------------------------------------------------------------------------------------|
216-
| `query` | yes | The query to ask about the codebase. |
217-
| `repos` | no | The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible. |
223+
| Name | Required | Description |
224+
|:----------------|:---------|:-----------------------------------------------------------------------------------------------------------------------------------------------|
225+
| `query` | yes | The query to ask about the codebase. |
226+
| `repos` | no | The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible. |
227+
| `languageModel` | no | The language model to use for answering the question. Object with `provider` and `model`. If not provided, defaults to the first model in the config. Use `list_language_models` to see available options. |
218228

219229

220230
## Environment Variables

‎packages/mcp/README.md‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -234,17 +234,29 @@ Get a list of commits for a given repository.
234234

235235
</details>
236236

237+
### list_language_models
238+
239+
Lists the available language models configured on the Sourcebot instance. Use this to discover which models can be specified when calling `ask_codebase`.
240+
241+
<details>
242+
<summary>Parameters</summary>
243+
244+
This tool takes no parameters.
245+
246+
</details>
247+
237248
### ask_codebase
238249

239250
Ask a natural language question about the codebase. This tool uses an AI agent to autonomously search code, read files, and find symbol references/definitions to answer your question. Returns a detailed answer in markdown format with code references, plus a link to view the full research session in the Sourcebot web UI.
240251

241252
<details>
242253
<summary>Parameters</summary>
243254

244-
| Name | Required | Description |
245-
|:--------|:---------|:------------------------------------------------------------------------------------------------------------------|
246-
| `query` | yes | The query to ask about the codebase. |
247-
| `repos` | no | The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible. |
255+
| Name | Required | Description |
256+
|:----------------|:---------|:-----------------------------------------------------------------------------------------------------------------------------------------------|
257+
| `query` | yes | The query to ask about the codebase. |
258+
| `repos` | no | The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible. |
259+
| `languageModel` | no | The language model to use for answering the question. Object with `provider` and `model`. If not provided, defaults to the first model in the config. Use `list_language_models` to see available options. |
248260

249261
</details>
250262

‎packages/mcp/src/client.ts‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { env } from './env.js';
2-
import { listReposResponseSchema, searchResponseSchema, fileSourceResponseSchema, listCommitsResponseSchema, askCodebaseResponseSchema } from './schemas.js';
3-
import { AskCodebaseRequest, AskCodebaseResponse, FileSourceRequest, ListReposQueryParams, SearchRequest, ListCommitsQueryParamsSchema } from './types.js';
2+
import { listReposResponseSchema, searchResponseSchema, fileSourceResponseSchema, listCommitsResponseSchema, askCodebaseResponseSchema, listLanguageModelsResponseSchema } from './schemas.js';
3+
import { AskCodebaseRequest, AskCodebaseResponse, FileSourceRequest, ListReposQueryParams, SearchRequest, ListCommitsQueryParamsSchema, ListLanguageModelsResponse } from './types.js';
44
import { isServiceError, ServiceErrorException } from './utils.js';
55
import { z } from 'zod';
66

@@ -111,7 +111,7 @@ export const listCommits = async (queryParams: ListCommitsQueryParamsSchema) =>
111111
/**
112112
* Asks a natural language question about the codebase using the Sourcebot AI agent.
113113
* This is a blocking call that runs the full agent loop and returns when complete.
114-
*
114+
*
115115
* @param request - The question and optional repo filters
116116
* @returns The agent's answer, chat URL, sources, and metadata
117117
*/
@@ -128,3 +128,21 @@ export const askCodebase = async (request: AskCodebaseRequest): Promise<AskCodeb
128128

129129
return parseResponse(response, askCodebaseResponseSchema);
130130
}
131+
132+
/**
133+
* Lists the available language models configured on the Sourcebot instance.
134+
*
135+
* @returns Array of language model info objects
136+
*/
137+
export const listLanguageModels = async (): Promise<ListLanguageModelsResponse> => {
138+
const response = await fetch(`${env.SOURCEBOT_HOST}/api/models`, {
139+
method: 'GET',
140+
headers: {
141+
'Content-Type': 'application/json',
142+
'X-Sourcebot-Client-Source': 'mcp',
143+
...(env.SOURCEBOT_API_KEY ? { 'X-Sourcebot-Api-Key': env.SOURCEBOT_API_KEY } : {})
144+
},
145+
});
146+
147+
return parseResponse(response, listLanguageModelsResponseSchema);
148+
}

‎packages/mcp/src/index.ts‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
66
import _dedent from "dedent";
77
import escapeStringRegexp from 'escape-string-regexp';
88
import { z } from 'zod';
9-
import { askCodebase, getFileSource, listCommits, listRepos, search } from './client.js';
9+
import { askCodebase, getFileSource, listCommits, listLanguageModels, listRepos, search } from './client.js';
1010
import { env, numberSchema } from './env.js';
1111
import { askCodebaseRequestSchema, fileSourceRequestSchema, listCommitsQueryParamsSchema, listReposQueryParamsSchema } from './schemas.js';
1212
import { AskCodebaseRequest, FileSourceRequest, ListCommitsQueryParamsSchema, ListReposQueryParams, TextContent } from './types.js';
@@ -238,6 +238,22 @@ server.tool(
238238
}
239239
);
240240

241+
server.tool(
242+
"list_language_models",
243+
dedent`Lists the available language models configured on the Sourcebot instance. Use this to discover which models can be specified when calling ask_codebase.`,
244+
{},
245+
async () => {
246+
const models = await listLanguageModels();
247+
248+
return {
249+
content: [{
250+
type: "text",
251+
text: JSON.stringify(models),
252+
}],
253+
};
254+
}
255+
);
256+
241257
server.tool(
242258
"ask_codebase",
243259
dedent`
@@ -262,6 +278,7 @@ server.tool(
262278
263279
---
264280
**View full research session:** ${response.chatUrl}
281+
**Model used:** ${response.languageModel.model}
265282
`;
266283

267284
return {

‎packages/mcp/src/schemas.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,14 @@ export const listCommitsResponseSchema = z.array(z.object({
273273
author_email: z.string(),
274274
}));
275275

276+
export const languageModelInfoSchema = z.object({
277+
provider: z.string().describe("The model provider (e.g., 'anthropic', 'openai')"),
278+
model: z.string().describe("The model ID"),
279+
displayName: z.string().optional().describe("Optional display name for the model"),
280+
});
281+
282+
export const listLanguageModelsResponseSchema = z.array(languageModelInfoSchema);
283+
276284
export const askCodebaseRequestSchema = z.object({
277285
query: z
278286
.string()
@@ -281,6 +289,10 @@ export const askCodebaseRequestSchema = z.object({
281289
.array(z.string())
282290
.optional()
283291
.describe("The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible."),
292+
languageModel: languageModelInfoSchema
293+
.omit({ displayName: true })
294+
.optional()
295+
.describe("The language model to use for answering the question. If not provided, defaults to the first model in the config. Use list_language_models to see available options."),
284296
});
285297

286298
export const sourceSchema = z.object({
@@ -296,4 +308,9 @@ export const askCodebaseResponseSchema = z.object({
296308
answer: z.string().describe("The agent's final answer in markdown format"),
297309
chatId: z.string().describe("ID of the persisted chat session"),
298310
chatUrl: z.string().describe("URL to view the chat in the web UI"),
311+
languageModel: z.object({
312+
provider: z.string().describe("The model provider (e.g., 'anthropic', 'openai')"),
313+
model: z.string().describe("The model ID"),
314+
displayName: z.string().optional().describe("Optional display name for the model"),
315+
}).describe("The language model used to generate the response"),
299316
});

‎packages/mcp/src/types.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ import {
1414
listCommitsResponseSchema,
1515
askCodebaseRequestSchema,
1616
askCodebaseResponseSchema,
17+
languageModelInfoSchema,
18+
listLanguageModelsResponseSchema,
1719
} from "./schemas.js";
1820
import { z } from "zod";
1921

@@ -39,3 +41,6 @@ export type ListCommitsResponse = z.infer<typeof listCommitsResponseSchema>;
3941

4042
export type AskCodebaseRequest = z.infer<typeof askCodebaseRequestSchema>;
4143
export type AskCodebaseResponse = z.infer<typeof askCodebaseResponseSchema>;
44+
45+
export type LanguageModelInfo = z.infer<typeof languageModelInfoSchema>;
46+
export type ListLanguageModelsResponse = z.infer<typeof listLanguageModelsResponseSchema>;

‎packages/web/src/app/api/(server)/chat/blocking/route.ts‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { sew } from "@/actions";
22
import { _getConfiguredLanguageModelsFull, _getAISDKLanguageModelAndOptions, updateChatMessages, generateAndUpdateChatNameFromMessage } from "@/features/chat/actions";
3-
import { SBChatMessage, SearchScope } from "@/features/chat/types";
4-
import { convertLLMOutputToPortableMarkdown, getAnswerPartFromAssistantMessage } from "@/features/chat/utils";
3+
import { LanguageModelInfo, languageModelInfoSchema, SBChatMessage, SearchScope } from "@/features/chat/types";
4+
import { convertLLMOutputToPortableMarkdown, getAnswerPartFromAssistantMessage, getLanguageModelKey } from "@/features/chat/utils";
55
import { ErrorCode } from "@/lib/errorCodes";
66
import { requestBodySchemaValidationError, ServiceError, ServiceErrorException, serviceErrorResponse } from "@/lib/serviceError";
77
import { isServiceError } from "@/lib/utils";
@@ -32,6 +32,9 @@ const blockingChatRequestSchema = z.object({
3232
.array(z.string())
3333
.optional()
3434
.describe("The repositories that are accessible to the agent during the chat. If not provided, all repositories are accessible."),
35+
languageModel: languageModelInfoSchema
36+
.optional()
37+
.describe("The language model to use for the chat. If not provided, the first configured model is used."),
3538
});
3639

3740
/**
@@ -41,6 +44,7 @@ interface BlockingChatResponse {
4144
answer: string;
4245
chatId: string;
4346
chatUrl: string;
47+
languageModel: LanguageModelInfo;
4448
}
4549

4650
/**
@@ -60,7 +64,7 @@ export const POST = apiHandler(async (request: NextRequest) => {
6064
return serviceErrorResponse(requestBodySchemaValidationError(parsed.error));
6165
}
6266

63-
const { query, repos = [] } = parsed.data;
67+
const { query, repos = [], languageModel: requestedLanguageModel } = parsed.data;
6468

6569
const response: BlockingChatResponse | ServiceError = await sew(() =>
6670
withOptionalAuthV2(async ({ org, user, prisma }) => {
@@ -70,13 +74,25 @@ export const POST = apiHandler(async (request: NextRequest) => {
7074
return {
7175
statusCode: StatusCodes.BAD_REQUEST,
7276
errorCode: ErrorCode.INVALID_REQUEST_BODY,
73-
message: "No language models are configured. Please configure at least one language model.",
77+
message: "No language models are configured. Please configure at least one language model. See: https://docs.sourcebot.dev/docs/configuration/language-model-providers",
7478
} satisfies ServiceError;
7579
}
7680

77-
// @todo: we should probably have a option of passing the language model
78-
// into the request body. For now, just use the first configured model.
79-
const languageModelConfig = configuredModels[0];
81+
// Use the requested language model if provided, otherwise default to the first configured model
82+
let languageModelConfig = configuredModels[0];
83+
if (requestedLanguageModel) {
84+
const matchingModel = configuredModels.find(
85+
(m) => getLanguageModelKey(m) === getLanguageModelKey(requestedLanguageModel as LanguageModelInfo)
86+
);
87+
if (!matchingModel) {
88+
return {
89+
statusCode: StatusCodes.BAD_REQUEST,
90+
errorCode: ErrorCode.INVALID_REQUEST_BODY,
91+
message: `Language model '${requestedLanguageModel.provider}/${requestedLanguageModel.model}' is not configured.`,
92+
} satisfies ServiceError;
93+
}
94+
languageModelConfig = matchingModel;
95+
}
8096

8197
const { model, providerOptions } = await _getAISDKLanguageModelAndOptions(languageModelConfig);
8298
const modelName = languageModelConfig.displayName ?? languageModelConfig.model;
@@ -184,6 +200,11 @@ export const POST = apiHandler(async (request: NextRequest) => {
184200
answer: portableAnswer,
185201
chatId: chat.id,
186202
chatUrl,
203+
languageModel: {
204+
provider: languageModelConfig.provider,
205+
model: languageModelConfig.model,
206+
displayName: languageModelConfig.displayName,
207+
},
187208
} satisfies BlockingChatResponse;
188209
})
189210
);

‎packages/web/src/app/api/(server)/chat/route.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ import { z } from "zod";
3030
const logger = createLogger('chat-api');
3131

3232
const chatRequestSchema = z.object({
33-
// These paramt
3433
messages: z.array(z.any()),
3534
id: z.string(),
3635
...additionalChatRequestParamsSchema.shape,

‎packages/web/src/app/api/(server)/models/route.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { sew } from "@/actions";
2+
import { apiHandler } from "@/lib/apiHandler";
3+
import { getConfiguredLanguageModelsInfo } from "@/features/chat/actions";
4+
import { serviceErrorResponse } from "@/lib/serviceError";
5+
import { isServiceError } from "@/lib/utils";
6+
import { withOptionalAuthV2 } from "@/withAuthV2";
7+
8+
export const GET = apiHandler(async () => {
9+
const response = await sew(() =>
10+
withOptionalAuthV2(async () => {
11+
const models = await getConfiguredLanguageModelsInfo();
12+
return models;
13+
})
14+
);
15+
16+
if (isServiceError(response)) {
17+
return serviceErrorResponse(response);
18+
}
19+
20+
return new Response(JSON.stringify(response), {
21+
status: 200,
22+
headers: { 'Content-Type': 'application/json' },
23+
});
24+
});

0 commit comments

Comments
 (0)