32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
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)
|