@@ -141,26 +141,40 @@ async def update_endpoint(
141
141
except Exception as err :
142
142
raise ValueError ("Unable to get models from provider: {}" .format (str (err )))
143
143
144
- # Reset all provider models.
145
- await self ._db_writer .delete_provider_models (str (endpoint .id ))
144
+ models_set = set (models )
146
145
147
- for model in models :
146
+ # Get the models from the provider
147
+ models_in_db = await self ._db_reader .get_provider_models_by_provider_id (str (endpoint .id ))
148
+
149
+ models_in_db_set = set (model .name for model in models_in_db )
150
+
151
+ # Add the models that are in the provider but not in the DB
152
+ for model in models_set - models_in_db_set :
148
153
await self ._db_writer .add_provider_model (
149
154
dbmodels .ProviderModel (
150
155
provider_endpoint_id = founddbe .id ,
151
156
name = model ,
152
157
)
153
158
)
154
159
160
+ # Remove the models that are in the DB but not in the provider
161
+ for model in models_in_db_set - models_set :
162
+ await self ._db_writer .delete_provider_model (
163
+ founddbe .id ,
164
+ model ,
165
+ )
166
+
155
167
dbendpoint = await self ._db_writer .update_provider_endpoint (endpoint .to_db_model ())
156
168
157
- await self ._db_writer .push_provider_auth_material (
158
- dbmodels .ProviderAuthMaterial (
159
- provider_endpoint_id = dbendpoint .id ,
160
- auth_type = endpoint .auth_type ,
161
- auth_blob = endpoint .api_key if endpoint .api_key else "" ,
169
+ # If an API key was provided or we've changed the auth type, we update the auth material
170
+ if endpoint .auth_type != founddbe .auth_type or endpoint .api_key :
171
+ await self ._db_writer .push_provider_auth_material (
172
+ dbmodels .ProviderAuthMaterial (
173
+ provider_endpoint_id = dbendpoint .id ,
174
+ auth_type = endpoint .auth_type ,
175
+ auth_blob = endpoint .api_key if endpoint .api_key else "" ,
176
+ )
162
177
)
163
- )
164
178
165
179
return apimodelsv1 .ProviderEndpoint .from_db_model (dbendpoint )
166
180
0 commit comments