File tree Expand file tree Collapse file tree 1 file changed +13
-1
lines changed Expand file tree Collapse file tree 1 file changed +13
-1
lines changed Original file line number Diff line number Diff line change 14
14
* See the License for the specific language governing permissions and
15
15
* limitations under the License.
16
16
*/
17
- import { ImagenModelParams , ModelParams , VertexAIErrorCode } from './types' ;
17
+ import {
18
+ ImagenModelParams ,
19
+ InferenceMode ,
20
+ ModelParams ,
21
+ VertexAIErrorCode
22
+ } from './types' ;
18
23
import { VertexAIError } from './errors' ;
19
24
import { ImagenModel , getGenerativeModel , getImagenModel } from './api' ;
20
25
import { expect } from 'chai' ;
@@ -101,6 +106,13 @@ describe('Top level API', () => {
101
106
expect ( genModel ) . to . be . an . instanceOf ( GenerativeModel ) ;
102
107
expect ( genModel . model ) . to . equal ( 'publishers/google/models/my-model' ) ;
103
108
} ) ;
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
+ } ) ;
104
116
it ( 'getImagenModel throws if no model is provided' , ( ) => {
105
117
try {
106
118
getImagenModel ( fakeVertexAI , { } as ImagenModelParams ) ;
You can’t perform that action at this time.
0 commit comments