Merge pull request #18 from Serverless-Devs/fix-integration-specific-model

feat(model): support specific model parameter in CommonModel

Author: OhYee

agentrun/integration/builtin/model.py

@@ -69,29 +69,25 @@ def model(
     backend_type = kwargs.get("backend_type")
     model = kwargs.get("model")
 
+    if isinstance(input, str):
+        from agentrun.model.client import ModelClient
+
+        client = ModelClient(config=config)
+        input = client.get(name=input, backend_type=backend_type, config=config)
+
     if isinstance(input, ModelProxy):
         return CommonModel(
-            model=model or "",
             model_obj=input,
             backend_type=BackendType.PROXY,
+            specific_model=model,
             config=config,
         )
     elif isinstance(input, ModelService):
         return CommonModel(
-            model=model or "",
             model_obj=input,
             backend_type=BackendType.SERVICE,
+            specific_model=model,
             config=config,
         )
-
-    from agentrun.model.client import ModelClient
-
-    client = ModelClient(config=config)
-    model_obj = client.get(name=input, backend_type=backend_type, config=config)
-
-    return CommonModel(
-        model=input,
-        model_obj=model_obj,
-        backend_type=backend_type,
-        config=config,
-    )
+    else:
+        raise TypeError("input must be str, ModelProxy or ModelService")

agentrun/integration/utils/model.py

@@ -19,14 +19,14 @@ class CommonModel:
 
     def __init__(
         self,
-        model: Optional[str],
         model_obj: Union[ModelService, ModelProxy],
         backend_type: Optional[BackendType] = None,
+        specific_model: Optional[str] = None,
         config: Optional[Config] = None,
     ):
-        self.model = model
         self.model_obj = model_obj
         self.backend_type = backend_type
+        self.specific_model = specific_model
         self.config = config or Config()
 
     def completions(self, *args, **kwargs):
@@ -40,7 +40,10 @@ def responses(self, *args, **kwargs):
     def get_model_info(self, config: Optional[Config] = None):
         """获取模型信息"""
         cfg = Config.with_configs(self.config, config)
-        return self.model_obj.model_info(config=cfg)
+        info = self.model_obj.model_info(config=cfg)
+        if self.specific_model:
+            info.model = self.specific_model
+        return info
 
     def __convert_model(self, adapter_name: str):
         try:

tests/unittests/integration/test_integration.py

@@ -86,6 +86,7 @@ def extract_payload(request):
         def build_response(request, route):
             payload = extract_payload(request)
             is_stream = payload.get("stream", False)
+            assert payload.get("model") == "mock-model-proxy"
             response_json = self._build_response(
                 payload.get("messages") or [], payload.get("tools")
             )
@@ -355,8 +356,6 @@ def get_mocked_model(
         )
         m = model("fake-llm-model")
 
-        assert m.model == "fake-llm-model"
-
         return m
 
     def test_langchain(self, monkeypatch, mock_llm_transport):