@@ -75,7 +75,7 @@ func (s *composeService) ensureModels(ctx context.Context, project *types.Projec
75
75
76
76
eg , gctx := errgroup .WithContext (ctx )
77
77
eg .Go (func () error {
78
- return s .setModelEndpointVariable (gctx , dockerModel , project )
78
+ return s .setModelVariables (gctx , dockerModel , project )
79
79
})
80
80
81
81
for name , config := range project .Models {
@@ -161,7 +161,7 @@ func (s *composeService) configureModel(ctx context.Context, dockerModel *manage
161
161
return cmd .Run ()
162
162
}
163
163
164
- func (s * composeService ) setModelEndpointVariable (ctx context.Context , dockerModel * manager.Plugin , project * types.Project ) error {
164
+ func (s * composeService ) setModelVariables (ctx context.Context , dockerModel * manager.Plugin , project * types.Project ) error {
165
165
cmd := exec .CommandContext (ctx , dockerModel .Path , "status" , "--json" )
166
166
s .setupChildProcess (ctx , cmd )
167
167
statusOut , err := cmd .CombinedOutput ()
@@ -179,12 +179,21 @@ func (s *composeService) setModelEndpointVariable(ctx context.Context, dockerMod
179
179
}
180
180
181
181
for _ , service := range project .Services {
182
- for model , modelConfig := range service .Models {
182
+ for ref , modelConfig := range service .Models {
183
+ model := project .Models [ref ]
184
+ varPrefix := strings .ReplaceAll (strings .ToUpper (ref ), "-" , "_" )
183
185
var variable string
186
+ if modelConfig != nil && modelConfig .ModelVariable != "" {
187
+ variable = modelConfig .ModelVariable
188
+ } else {
189
+ variable = varPrefix
190
+ }
191
+ service .Environment [variable ] = & model .Model
192
+
184
193
if modelConfig != nil && modelConfig .EndpointVariable != "" {
185
194
variable = modelConfig .EndpointVariable
186
195
} else {
187
- variable = strings . ToUpper ( model ) + "_URL"
196
+ variable = varPrefix + "_URL"
188
197
}
189
198
service .Environment [variable ] = & status .Endpoint
190
199
}
0 commit comments