Skip to content

Commit 117b1e2

Browse files
authored
Merge pull request #1058 from narengogi/chore/gemini-thinking-model-changes
add support for gemini thinking input params
2 parents 51430c9 + 3773892 commit 117b1e2

File tree

4 files changed

+25
-0
lines changed

4 files changed

+25
-0
lines changed

src/providers/google-vertex-ai/chatComplete.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,10 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
337337
labels: {
338338
param: 'labels',
339339
},
340+
thinking: {
341+
param: 'generationConfig',
342+
transform: (params: Params) => transformGenerationConfig(params),
343+
},
340344
};
341345

342346
interface AnthorpicTextContentItem {

src/providers/google-vertex-ai/transformGenerationConfig.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,5 +50,12 @@ export function transformGenerationConfig(params: Params) {
5050
generationConfig['responseSchema'] = schema;
5151
}
5252

53+
if (params?.thinking) {
54+
const thinkingConfig: Record<string, any> = {};
55+
thinkingConfig['include_thoughts'] = true;
56+
thinkingConfig['thinking_budget'] = params.thinking.budget_tokens;
57+
generationConfig['thinking_config'] = thinkingConfig;
58+
}
59+
5360
return generationConfig;
5461
}

src/providers/google/chatComplete.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,12 @@ const transformGenerationConfig = (params: Params) => {
7171
}
7272
generationConfig['responseSchema'] = schema;
7373
}
74+
if (params?.thinking) {
75+
const thinkingConfig: Record<string, any> = {};
76+
thinkingConfig['include_thoughts'] = true;
77+
thinkingConfig['thinking_budget'] = params.thinking.budget_tokens;
78+
generationConfig['thinking_config'] = thinkingConfig;
79+
}
7480
return generationConfig;
7581
};
7682

@@ -405,6 +411,10 @@ export const GoogleChatCompleteConfig: ProviderConfig = {
405411
}
406412
},
407413
},
414+
thinking: {
415+
param: 'generationConfig',
416+
transform: (params: Params) => transformGenerationConfig(params),
417+
},
408418
};
409419

410420
export interface GoogleErrorResponse {

src/types/requestBody.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,10 @@ export interface Params {
404404
// Anthropic specific
405405
anthropic_beta?: string;
406406
anthropic_version?: string;
407+
thinking?: {
408+
type?: string;
409+
budget_tokens: number;
410+
};
407411
}
408412

409413
interface Examples {

0 commit comments

Comments
 (0)