Skip to content

Commit 69213fb

Browse files
committed
fix(vertexai): pass GenerativeModel's BaseParams to ChatSession
1 parent 9952dbc commit 69213fb

File tree

2 files changed

+36
-0
lines changed

2 files changed

+36
-0
lines changed

packages/vertexai/src/models/generative-model.test.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,35 @@ describe('GenerativeModel', () => {
165165
);
166166
restore();
167167
});
168+
it('passes base model params through to ChatSession when there are no startChatParams', async () => {
169+
const genModel = new GenerativeModel(fakeVertexAI, {
170+
model: 'my-model',
171+
generationConfig: {
172+
topK: 1
173+
}
174+
});
175+
const chatSession = genModel.startChat();
176+
expect(chatSession.params?.generationConfig).to.deep.equal({
177+
topK: 1
178+
});
179+
restore();
180+
});
181+
it('overrides base model params with startChatParams', () => {
182+
const genModel = new GenerativeModel(fakeVertexAI, {
183+
model: 'my-model',
184+
generationConfig: {
185+
topK: 1
186+
}
187+
});
188+
const chatSession = genModel.startChat({
189+
generationConfig: {
190+
topK: 2
191+
}
192+
});
193+
expect(chatSession.params?.generationConfig).to.deep.equal({
194+
topK: 2
195+
});
196+
});
168197
it('passes params through to chat.sendMessage', async () => {
169198
const genModel = new GenerativeModel(fakeVertexAI, {
170199
model: 'my-model',

packages/vertexai/src/models/generative-model.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,13 @@ export class GenerativeModel extends VertexAIModel {
132132
tools: this.tools,
133133
toolConfig: this.toolConfig,
134134
systemInstruction: this.systemInstruction,
135+
generationConfig: this.generationConfig,
136+
safetySettings: this.safetySettings,
137+
/**
138+
* Overrides params inherited from GenerativeModel with those explicitly set in the
139+
* StartChatParams. For example, if startChatParams.generationConfig is set, it'll override
140+
* this.generationConfig.
141+
*/
135142
...startChatParams
136143
},
137144
this.requestOptions

0 commit comments

Comments
 (0)