Skip to content

Commit 0d628cf

Browse files
committed
Assert HybridParams sets the model name
1 parent 7752cb0 commit 0d628cf

File tree

1 file changed

+13
-1
lines changed

1 file changed

+13
-1
lines changed

packages/vertexai/src/api.test.ts

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,12 @@
1414
* See the License for the specific language governing permissions and
1515
* limitations under the License.
1616
*/
17-
import { ImagenModelParams, ModelParams, VertexAIErrorCode } from './types';
17+
import {
18+
ImagenModelParams,
19+
InferenceMode,
20+
ModelParams,
21+
VertexAIErrorCode
22+
} from './types';
1823
import { VertexAIError } from './errors';
1924
import { ImagenModel, getGenerativeModel, getImagenModel } from './api';
2025
import { expect } from 'chai';
@@ -101,6 +106,13 @@ describe('Top level API', () => {
101106
expect(genModel).to.be.an.instanceOf(GenerativeModel);
102107
expect(genModel.model).to.equal('publishers/google/models/my-model');
103108
});
109+
it('getGenerativeModel with HybridParams sets the model', () => {
110+
const genModel = getGenerativeModel(fakeVertexAI, {
111+
mode: InferenceMode.ONLY_ON_CLOUD,
112+
onCloudParams: { model: 'my-model' }
113+
});
114+
expect(genModel.model).to.equal('publishers/google/models/my-model');
115+
});
104116
it('getImagenModel throws if no model is provided', () => {
105117
try {
106118
getImagenModel(fakeVertexAI, {} as ImagenModelParams);

0 commit comments

Comments
 (0)