refactor(model): integrate ModelAPI into model proxy and service classes

OhYee · OhYee · commit cfffd23220b0 · 2026-01-13T12:04:18.000+08:00
This change introduces ModelAPI as a mixin to both ModelProxy and ModelService classes,
consolidating common API functionality. The completions and responses methods have been
removed from these classes as they are now inherited from ModelAPI, reducing code duplication
and improving maintainability.

The refactoring affects multiple files including __model_proxy_async_template.py,
__model_service_async_template.py, model_proxy.py, and model_service.py, along with
the new model_api.py module definition.

Two new files were created: agentrun/model/api/model_api.py which defines the ModelAPI
base class, and examples/embedding.py which provides usage examples for embedding functionality.

//

此更改将ModelAPI作为mixin引入到ModelProxy和ModelService类中，
整合了通用的API功能。completions和responses方法已从这些类中移除，
因为它们现在继承自ModelAPI，减少了代码重复并提高了可维护性。

重构影响了多个文件，包括__model_proxy_async_template.py、
__model_service_async_template.py、model_proxy.py和model_service.py，
以及新的model_api.py模块定义。

创建了两个新文件：agentrun/model/api/model_api.py定义了ModelAPI基类，
examples/embedding.py提供了嵌入功能的使用示例。

Change-Id: I81045039bef65df9d7daab6a56ff2e1442fddcb7
Signed-off-by: OhYee &lt;oyohyee@oyohyee.com&gt;
diff --git a/agentrun/model/__model_proxy_async_template.py b/agentrun/model/__model_proxy_async_template.py
@@ -9,6 +9,7 @@
 import pydash
 
 from agentrun.model.api.data import BaseInfo, ModelDataAPI
+from agentrun.model.api.model_api import ModelAPI
 from agentrun.utils.config import Config
 from agentrun.utils.model import Status
 from agentrun.utils.resource import ResourceBase
@@ -30,6 +31,7 @@ class ModelProxy(
     ModelProxyImmutableProps,
     ModelProxyMutableProps,
     ModelProxySystemProps,
+    ModelAPI,
     ResourceBase,
 ):
     """模型服务"""
@@ -230,41 +232,3 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
         )
 
         return self._data_client.model_info()
-
-    def completions(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        config: Optional[Config] = None,
-        **kwargs,
-    ):
-        self.model_info(config)
-        assert self._data_client
-
-        return self._data_client.completions(
-            **kwargs,
-            messages=messages,
-            model=model,
-            stream=stream,
-            config=config,
-        )
-
-    def responses(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        config: Optional[Config] = None,
-        **kwargs,
-    ):
-        self.model_info(config)
-        assert self._data_client
-
-        return self._data_client.responses(
-            **kwargs,
-            messages=messages,
-            model=model,
-            stream=stream,
-            config=config,
-        )
diff --git a/agentrun/model/__model_service_async_template.py b/agentrun/model/__model_service_async_template.py
@@ -6,7 +6,8 @@
 
 from typing import List, Optional
 
-from agentrun.model.api.data import BaseInfo, ModelCompletionAPI
+from agentrun.model.api.data import BaseInfo
+from agentrun.model.api.model_api import ModelAPI
 from agentrun.utils.config import Config
 from agentrun.utils.model import PageableInput
 from agentrun.utils.resource import ResourceBase
@@ -27,6 +28,7 @@ class ModelService(
     ModelServiceImmutableProps,
     ModelServiceMutableProps,
     ModelServicesSystemProps,
+    ModelAPI,
     ResourceBase,
 ):
     """模型服务"""
@@ -230,38 +232,3 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
             model=default_model,
             headers=cfg.get_headers(),
         )
-
-    def completions(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        **kwargs,
-    ):
-        info = self.model_info(config=kwargs.get("config"))
-
-        m = ModelCompletionAPI(
-            api_key=info.api_key or "",
-            base_url=info.base_url or "",
-            model=model or info.model or self.model_service_name or "",
-        )
-
-        return m.completions(**kwargs, messages=messages, stream=stream)
-
-    def responses(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        **kwargs,
-    ):
-        info = self.model_info(config=kwargs.get("config"))
-
-        m = ModelCompletionAPI(
-            api_key=info.api_key or "",
-            base_url=info.base_url or "",
-            model=model or info.model or self.model_service_name or "",
-            provider=(self.provider or "openai").lower(),
-        )
-
-        return m.responses(**kwargs, messages=messages, stream=stream)
diff --git a/agentrun/model/api/model_api.py b/agentrun/model/api/model_api.py
@@ -0,0 +1,156 @@
+from abc import ABC, abstractmethod
+from typing import Optional, TYPE_CHECKING, Union
+
+from .data import BaseInfo
+
+if TYPE_CHECKING:
+    from litellm import ResponseInputParam
+
+
+class ModelAPI(ABC):
+
+    @abstractmethod
+    def model_info(self) -> BaseInfo:
+        ...
+
+    def completions(
+        self,
+        **kwargs,
+    ):
+        """
+        Deprecated. Use completion() instead.
+        """
+        import warnings
+
+        warnings.warn(
+            "completions() is deprecated, use completion() instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.completion(**kwargs)
+
+    def completion(
+        self,
+        messages=[],
+        model: Optional[str] = None,
+        custom_llm_provider: Optional[str] = None,
+        **kwargs,
+    ):
+        from litellm import completion
+
+        info = self.model_info()
+        return completion(
+            **kwargs,
+            api_key=info.api_key,
+            base_url=info.base_url,
+            model=model or info.model or "",
+            custom_llm_provider=custom_llm_provider
+            or info.provider
+            or "openai",
+            messages=messages,
+        )
+
+    async def acompletion(
+        self,
+        messages=[],
+        model: Optional[str] = None,
+        custom_llm_provider: Optional[str] = None,
+        **kwargs,
+    ):
+        from litellm import acompletion
+
+        info = self.model_info()
+        return await acompletion(
+            **kwargs,
+            api_key=info.api_key,
+            base_url=info.base_url,
+            model=model or info.model or "",
+            custom_llm_provider=custom_llm_provider
+            or info.provider
+            or "openai",
+            messages=messages,
+        )
+
+    def responses(
+        self,
+        input: Union[str, "ResponseInputParam"],
+        model: Optional[str] = None,
+        custom_llm_provider: Optional[str] = None,
+        **kwargs,
+    ):
+        from litellm import responses
+
+        info = self.model_info()
+        return responses(
+            **kwargs,
+            api_key=info.api_key,
+            base_url=info.base_url,
+            model=model or info.model or "",
+            custom_llm_provider=custom_llm_provider
+            or info.provider
+            or "openai",
+            input=input,
+        )
+
+    async def aresponses(
+        self,
+        input: Union[str, "ResponseInputParam"],
+        model: Optional[str] = None,
+        custom_llm_provider: Optional[str] = None,
+        **kwargs,
+    ):
+        from litellm import aresponses
+
+        info = self.model_info()
+        return await aresponses(
+            **kwargs,
+            api_key=info.api_key,
+            base_url=info.base_url,
+            model=model or info.model or "",
+            custom_llm_provider=custom_llm_provider
+            or info.provider
+            or "openai",
+            input=input,
+        )
+
+    def embedding(
+        self,
+        input=[],
+        model: Optional[str] = None,
+        custom_llm_provider: Optional[str] = None,
+        **kwargs,
+    ):
+        from litellm import embedding
+
+        info = self.model_info()
+        return embedding(
+            **kwargs,
+            api_key=info.api_key,
+            api_base=info.base_url,
+            model=model or info.model or "",
+            custom_llm_provider=custom_llm_provider
+            or info.provider
+            or "openai",
+            input=input,
+        )
+
+    def aembedding(
+        self,
+        input=[],
+        model: Optional[str] = None,
+        custom_llm_provider: Optional[str] = None,
+        **kwargs,
+    ):
+        from litellm import aembedding
+
+        info = self.model_info()
+        return aembedding(
+            **kwargs,
+            api_key=info.api_key,
+            api_base=info.base_url,
+            model=model or info.model or "",
+            custom_llm_provider=custom_llm_provider
+            or info.provider
+            or "openai",
+            input=input,
+        )
diff --git a/agentrun/model/model_proxy.py b/agentrun/model/model_proxy.py
@@ -19,6 +19,7 @@
 import pydash
 
 from agentrun.model.api.data import BaseInfo, ModelDataAPI
+from agentrun.model.api.model_api import ModelAPI
 from agentrun.utils.config import Config
 from agentrun.utils.model import Status
 from agentrun.utils.resource import ResourceBase
@@ -40,6 +41,7 @@ class ModelProxy(
     ModelProxyImmutableProps,
     ModelProxyMutableProps,
     ModelProxySystemProps,
+    ModelAPI,
     ResourceBase,
 ):
     """模型服务"""
@@ -399,41 +401,3 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
         )
 
         return self._data_client.model_info()
-
-    def completions(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        config: Optional[Config] = None,
-        **kwargs,
-    ):
-        self.model_info(config)
-        assert self._data_client
-
-        return self._data_client.completions(
-            **kwargs,
-            messages=messages,
-            model=model,
-            stream=stream,
-            config=config,
-        )
-
-    def responses(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        config: Optional[Config] = None,
-        **kwargs,
-    ):
-        self.model_info(config)
-        assert self._data_client
-
-        return self._data_client.responses(
-            **kwargs,
-            messages=messages,
-            model=model,
-            stream=stream,
-            config=config,
-        )
diff --git a/agentrun/model/model_service.py b/agentrun/model/model_service.py
@@ -16,7 +16,8 @@
 
 from typing import List, Optional
 
-from agentrun.model.api.data import BaseInfo, ModelCompletionAPI
+from agentrun.model.api.data import BaseInfo
+from agentrun.model.api.model_api import ModelAPI
 from agentrun.utils.config import Config
 from agentrun.utils.model import PageableInput
 from agentrun.utils.resource import ResourceBase
@@ -37,6 +38,7 @@ class ModelService(
     ModelServiceImmutableProps,
     ModelServiceMutableProps,
     ModelServicesSystemProps,
+    ModelAPI,
     ResourceBase,
 ):
     """模型服务"""
@@ -401,38 +403,3 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
             model=default_model,
             headers=cfg.get_headers(),
         )
-
-    def completions(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        **kwargs,
-    ):
-        info = self.model_info(config=kwargs.get("config"))
-
-        m = ModelCompletionAPI(
-            api_key=info.api_key or "",
-            base_url=info.base_url or "",
-            model=model or info.model or self.model_service_name or "",
-        )
-
-        return m.completions(**kwargs, messages=messages, stream=stream)
-
-    def responses(
-        self,
-        messages: list,
-        model: Optional[str] = None,
-        stream: bool = False,
-        **kwargs,
-    ):
-        info = self.model_info(config=kwargs.get("config"))
-
-        m = ModelCompletionAPI(
-            api_key=info.api_key or "",
-            base_url=info.base_url or "",
-            model=model or info.model or self.model_service_name or "",
-            provider=(self.provider or "openai").lower(),
-        )
-
-        return m.responses(**kwargs, messages=messages, stream=stream)
diff --git a/examples/embedding.py b/examples/embedding.py
