1  from typing import Any
 2  
 3  from mlflow.pyfunc.loaders.chat_agent import _ChatAgentPyfuncWrapper
 4  from mlflow.pyfunc.loaders.chat_model import _ChatModelPyfuncWrapper
 5  from mlflow.pyfunc.model import (
 6      ChatAgent,
 7      ChatModel,
 8      _load_context_model_and_signature,
 9      _PythonModelPyfuncWrapper,
10  )
11  
12  try:
13      from mlflow.pyfunc.model import ResponsesAgent
14  
15      IS_RESPONSES_AGENT_AVAILABLE = True
16  except ImportError:
17      IS_RESPONSES_AGENT_AVAILABLE = False
18  
19  
20  def _load_pyfunc(local_path: str, model_config: dict[str, Any] | None = None):
21      context, model, signature = _load_context_model_and_signature(local_path, model_config)
22      if isinstance(model, ChatModel):
23          return _ChatModelPyfuncWrapper(model, context, signature)
24      elif isinstance(model, ChatAgent):
25          return _ChatAgentPyfuncWrapper(model)
26      elif IS_RESPONSES_AGENT_AVAILABLE and isinstance(model, ResponsesAgent):
27          from mlflow.pyfunc.loaders.responses_agent import _ResponsesAgentPyfuncWrapper
28  
29          return _ResponsesAgentPyfuncWrapper(model, context)
30      else:
31          return _PythonModelPyfuncWrapper(model, context, signature)