Skip to content
Closed
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/provider-settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ const openAiSchema = baseProviderSettingsSchema.extend({
const ollamaSchema = baseProviderSettingsSchema.extend({
ollamaModelId: z.string().optional(),
ollamaBaseUrl: z.string().optional(),
ollamaContextWindow: z.number().optional(),
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we add validation constraints here? Something like:

Suggested change
ollamaContextWindow: z.number().optional(),
ollamaContextWindow: z.number().positive().optional(),

This would prevent invalid values like 0 or negative numbers from being set.

})

const vsCodeLmSchema = baseProviderSettingsSchema.extend({
Expand Down
98 changes: 98 additions & 0 deletions src/api/providers/__tests__/native-ollama.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,103 @@ describe("NativeOllamaHandler", () => {
expect(results.some((r) => r.type === "reasoning")).toBe(true)
expect(results.some((r) => r.type === "text")).toBe(true)
})

describe("context window configuration", () => {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Excellent test coverage! These tests thoroughly verify both the custom context window usage and the fallback behavior. Consider adding an edge case test for invalid values (0 or negative) once we add schema validation.

it("should use custom context window when ollamaContextWindow is provided", async () => {
const customContextWindow = 48000
const optionsWithCustomContext: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
ollamaBaseUrl: "http://localhost:11434",
ollamaContextWindow: customContextWindow,
}
handler = new NativeOllamaHandler(optionsWithCustomContext)

// Mock the chat response
mockChat.mockImplementation(async function* () {
yield {
message: { content: "Test response" },
eval_count: 10,
prompt_eval_count: 5,
}
})

// Create a message to trigger the chat call
const generator = handler.createMessage("System prompt", [{ role: "user", content: "Test message" }])

// Consume the generator
const results = []
for await (const chunk of generator) {
results.push(chunk)
}

// Verify that chat was called with the custom context window
expect(mockChat).toHaveBeenCalledWith(
expect.objectContaining({
options: expect.objectContaining({
num_ctx: customContextWindow,
}),
}),
)
})

it("should use model's default context window when ollamaContextWindow is not provided", async () => {
// Mock the chat response
mockChat.mockImplementation(async function* () {
yield {
message: { content: "Test response" },
eval_count: 10,
prompt_eval_count: 5,
}
})

// Create a message to trigger the chat call
const generator = handler.createMessage("System prompt", [{ role: "user", content: "Test message" }])

// Consume the generator
const results = []
for await (const chunk of generator) {
results.push(chunk)
}

// Verify that chat was called with the model's default context window (4096)
expect(mockChat).toHaveBeenCalledWith(
expect.objectContaining({
options: expect.objectContaining({
num_ctx: 4096,
}),
}),
)
})

it("should use custom context window in completePrompt method", async () => {
const customContextWindow = 48000
const optionsWithCustomContext: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
ollamaBaseUrl: "http://localhost:11434",
ollamaContextWindow: customContextWindow,
}
handler = new NativeOllamaHandler(optionsWithCustomContext)

// Mock the chat response
mockChat.mockResolvedValue({
message: { content: "Test response" },
})

// Call completePrompt
await handler.completePrompt("Test prompt")

// Verify that chat was called with the custom context window
expect(mockChat).toHaveBeenCalledWith(
expect.objectContaining({
options: expect.objectContaining({
num_ctx: customContextWindow,
}),
}),
)
})
})
})

describe("completePrompt", () => {
Expand All @@ -115,6 +212,7 @@ describe("NativeOllamaHandler", () => {
messages: [{ role: "user", content: "Tell me a joke" }],
stream: false,
options: {
num_ctx: 4096,
temperature: 0,
},
})
Expand Down
5 changes: 3 additions & 2 deletions src/api/providers/native-ollama.ts
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
messages: ollamaMessages,
stream: true,
options: {
num_ctx: modelInfo.contextWindow,
num_ctx: this.options.ollamaContextWindow || modelInfo.contextWindow,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice implementation using the fallback pattern! The optional chaining with || fallback to modelInfo.contextWindow is clean and follows the existing patterns in the codebase.

temperature: this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
},
})
Expand Down Expand Up @@ -262,14 +262,15 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
async completePrompt(prompt: string): Promise<string> {
try {
const client = this.ensureClient()
const { id: modelId } = await this.fetchModel()
const { id: modelId, info: modelInfo } = await this.fetchModel()
const useR1Format = modelId.toLowerCase().includes("deepseek-r1")

const response = await client.chat({
model: modelId,
messages: [{ role: "user", content: prompt }],
stream: false,
options: {
num_ctx: this.options.ollamaContextWindow || modelInfo.contextWindow,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good consistency - same pattern applied to the non-streaming method. Though I wonder if we should add more specific error handling here to distinguish between model fetch failures and context window configuration issues?

temperature: this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
},
})
Expand Down
Loading