star

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/code_model.py

@@ -0,0 +1,31 @@
+from typing import Any, Optional
+
+from mlflow.pyfunc.loaders.chat_agent import _ChatAgentPyfuncWrapper
+from mlflow.pyfunc.loaders.chat_model import _ChatModelPyfuncWrapper
+from mlflow.pyfunc.model import (
+    ChatAgent,
+    ChatModel,
+    _load_context_model_and_signature,
+    _PythonModelPyfuncWrapper,
+)
+
+try:
+    from mlflow.pyfunc.model import ResponsesAgent
+
+    IS_RESPONSES_AGENT_AVAILABLE = True
+except ImportError:
+    IS_RESPONSES_AGENT_AVAILABLE = False
+
+
+def _load_pyfunc(local_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, model, signature = _load_context_model_and_signature(local_path, model_config)
+    if isinstance(model, ChatModel):
+        return _ChatModelPyfuncWrapper(model, context, signature)
+    elif isinstance(model, ChatAgent):
+        return _ChatAgentPyfuncWrapper(model)
+    elif IS_RESPONSES_AGENT_AVAILABLE and isinstance(model, ResponsesAgent):
+        from mlflow.pyfunc.loaders.responses_agent import _ResponsesAgentPyfuncWrapper
+
+        return _ResponsesAgentPyfuncWrapper(model)
+    else:
+        return _PythonModelPyfuncWrapper(model, context, signature)