chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
from typing import Any
|
||||
|
||||
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: dict[str, Any] | None = 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, context)
|
||||
else:
|
||||
return _PythonModelPyfuncWrapper(model, context, signature)
|
||||
Reference in New Issue
Block a user