@@ -28,7 +28,7 @@ def init_models(self, use_local_proxy: bool = False):
28
28
self ._register_google_models (use_local_proxy = use_local_proxy )
29
29
self ._register_qwen_models (use_local_proxy = use_local_proxy )
30
30
self ._register_langchain_models (use_local_proxy = use_local_proxy )
31
-
31
+ self . _register_vllm_models ( use_local_proxy = use_local_proxy )
32
32
def _check_local_api_key (self , local_api_key_name : str , remote_api_key_name : str ) -> str :
33
33
api_key = os .getenv (local_api_key_name , PLACEHOLDER )
34
34
if api_key == PLACEHOLDER :
@@ -383,4 +383,56 @@ def _register_langchain_models(self, use_local_proxy: bool = False):
383
383
api_key = api_key ,
384
384
base_url = api_base ,
385
385
)
386
- self .registed_models [model_name ] = model
386
+ self .registed_models [model_name ] = model
387
+ def _register_vllm_models (self , use_local_proxy : bool = False ):
388
+ # qwen
389
+ api_key = self ._check_local_api_key (local_api_key_name = "QWEN_API_KEY" ,
390
+ remote_api_key_name = "QWEN_API_KEY" )
391
+ api_base = self ._check_local_api_base (local_api_base_name = "QWEN_API_BASE" ,
392
+ remote_api_base_name = "QWEN_API_BASE" )
393
+ models = [
394
+ {
395
+ "model_name" : "Qwen" ,
396
+ "model_id" : "Qwen" ,
397
+ }
398
+ ]
399
+ for model in models :
400
+ model_name = model ["model_name" ]
401
+ model_id = model ["model_id" ]
402
+
403
+ client = OpenAI (
404
+ api_key = api_key ,
405
+ base_url = api_base ,
406
+ )
407
+ model = OpenAIServerModel (
408
+ model_id = model_id ,
409
+ http_client = client ,
410
+ custom_role_conversions = custom_role_conversions ,
411
+ )
412
+ self .registed_models [model_name ] = model
413
+
414
+ # Qwen-VL
415
+ api_key_VL = self ._check_local_api_key (local_api_key_name = "QWEN_VL_API_KEY" ,
416
+ remote_api_key_name = "QWEN_VL_API_KEY" )
417
+ api_base_VL = self ._check_local_api_base (local_api_base_name = "QWEN_VL_API_BASE" ,
418
+ remote_api_base_name = "QWEN_VL_API_BASE" )
419
+ models = [
420
+ {
421
+ "model_name" : "Qwen-VL" ,
422
+ "model_id" : "Qwen-VL" ,
423
+ }
424
+ ]
425
+ for model in models :
426
+ model_name = model ["model_name" ]
427
+ model_id = model ["model_id" ]
428
+
429
+ client = OpenAI (
430
+ api_key = api_key_VL ,
431
+ base_url = api_base_VL ,
432
+ )
433
+ model = OpenAIServerModel (
434
+ model_id = model_id ,
435
+ http_client = client ,
436
+ custom_role_conversions = custom_role_conversions ,
437
+ )
438
+ self .registed_models [model_name ] = model
0 commit comments