chmantha/zenml
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

star

Browse Source
This commit is contained in:
Christian Mantha
2026-03-02 19:10:52 -05:00
commit 2ca0b9ef7c
28907 changed files with 5233713 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

602
venv/lib/python3.9/site-packages/mlflow/llama_index/__init__.py Normal file
View File

@@ -0,0 +1,602 @@
import logging
import os
import tempfile
from typing import Any, Optional, Union
import yaml
import mlflow
from mlflow import pyfunc
from mlflow.entities.model_registry.prompt import Prompt
from mlflow.exceptions import MlflowException
from mlflow.llama_index.pyfunc_wrapper import create_pyfunc_wrapper
from mlflow.models import Model, ModelInputExample, ModelSignature
from mlflow.models.model import MLMODEL_FILE_NAME, MODEL_CODE_PATH, MODEL_CONFIG
from mlflow.models.signature import _infer_signature_from_input_example
from mlflow.models.utils import (
_load_model_code_path,
_save_example,
_validate_and_get_model_code_path,
)
from mlflow.tracing.provider import trace_disabled
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.utils.annotations import experimental
from mlflow.utils.autologging_utils import autologging_integration
from mlflow.utils.docstring_utils import LOG_MODEL_PARAM_DOCS, format_docstring
from mlflow.utils.environment import (
_CONDA_ENV_FILE_NAME,
_CONSTRAINTS_FILE_NAME,
_PYTHON_ENV_FILE_NAME,
_REQUIREMENTS_FILE_NAME,
_mlflow_conda_env,
_process_conda_env,
_process_pip_requirements,
_PythonEnv,
_validate_env_arguments,
)
from mlflow.utils.file_utils import get_total_file_size, write_to
from mlflow.utils.model_utils import (
_add_code_from_conf_to_system_path,
_get_flavor_configuration,
_validate_and_copy_code_paths,
_validate_and_copy_file_to_directory,
_validate_and_get_model_config_from_file,
_validate_and_prepare_target_save_path,
)
from mlflow.utils.requirements_utils import _get_pinned_requirement
FLAVOR_NAME = "llama_index"
_INDEX_PERSIST_FOLDER = "index"
_SETTINGS_FILE = "settings.json"
_logger = logging.getLogger(__name__)
def get_default_pip_requirements():
"""
Returns:
A list of default pip requirements for MLflow Models produced by this flavor.
Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment
that, at a minimum, contains these requirements.
"""
return [_get_pinned_requirement("llama-index")]
def get_default_conda_env():
"""
Returns:
The default Conda environment for MLflow Models produced by calls to
:func:`save_model()` and :func:`log_model()`.
"""
return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements())
def _validate_engine_type(engine_type: str):
from mlflow.llama_index.pyfunc_wrapper import SUPPORTED_ENGINES
if engine_type not in SUPPORTED_ENGINES:
raise ValueError(
f"Currently mlflow only supports the following engine types: "
f"{SUPPORTED_ENGINES}. {engine_type} is not supported, so please "
"use one of the above types."
)
def _get_llama_index_version() -> str:
try:
import llama_index.core
return llama_index.core.__version__
except ImportError:
raise MlflowException(
"The llama_index module is not installed. "
"Please install it via `pip install llama-index`."
)
def _supported_classes():
from llama_index.core.base.base_query_engine import BaseQueryEngine
from llama_index.core.chat_engine.types import BaseChatEngine
from llama_index.core.indices.base import BaseIndex
from llama_index.core.retrievers import BaseRetriever
supported = (BaseIndex, BaseChatEngine, BaseQueryEngine, BaseRetriever)
try:
from llama_index.core.workflow import Workflow
supported += (Workflow,)
except ImportError:
pass
return supported
@experimental
@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
@trace_disabled # Suppress traces while loading model
def save_model(
llama_index_model,
path: str,
engine_type: Optional[str] = None,
model_config: Optional[Union[str, dict[str, Any]]] = None,
code_paths=None,
mlflow_model: Optional[Model] = None,
signature: Optional[ModelSignature] = None,
input_example: Optional[ModelInputExample] = None,
pip_requirements: Optional[Union[list[str], str]] = None,
extra_pip_requirements: Optional[Union[list[str], str]] = None,
conda_env=None,
metadata: Optional[dict[str, Any]] = None,
) -> None:
"""
Save a LlamaIndex model to a path on the local file system.
.. attention::
Saving a non-index object is only supported in the 'Model-from-Code' saving mode.
Please refer to the `Models From Code Guide <https://www.mlflow.org/docs/latest/model/models-from-code.html>`_
for more information.
.. note::
When logging a model, MLflow will automatically save the state of the ``Settings``
object so that you can use the same settings at inference time. However, please
note that some information in the ``Settings`` object will not be saved, including:
- API keys for avoiding key leakage.
- Function objects which are not serializable.
Args:
llama_index_model: A LlamaIndex object to be saved. Supported model types are:
1. An Index object.
2. An Engine object e.g. ChatEngine, QueryEngine, Retriever.
3. A `Workflow <https://docs.llamaindex.ai/en/stable/module_guides/workflow/>`_ object.
4. A string representing the path to a script contains LlamaIndex model definition
of the one of the above types.
path: Local path where the serialized model (as YAML) is to be saved.
engine_type: Required when saving an Index object to determine the inference interface
for the index when loaded as a pyfunc model. This field is **not** required when
saving other LlamaIndex objects. The supported values are as follows:
- ``"chat"``: load the index as an instance of the LlamaIndex
`ChatEngine <https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/>`_.
- ``"query"``: load the index as an instance of the LlamaIndex
`QueryEngine <https://docs.llamaindex.ai/en/stable/module_guides/deploying/query_engine/>`_.
- ``"retriever"``: load the index as an instance of the LlamaIndex
`Retriever <https://docs.llamaindex.ai/en/stable/module_guides/querying/retriever/>`_.
model_config: The model configuration to apply when loading the model back with
``mlflow.pyfunc.load_model()``. It will be applied in a different way depending on the
model type and saving method. See the docstring of :func:`log_model` for more details
and usage examples.
code_paths: {{ code_paths }}
mlflow_model: An MLflow model object that specifies the flavor that this model is being
added to.
signature: A Model Signature object that describes the input and output Schema of the
model. The model signature can be inferred using ``infer_signature`` function
of ``mlflow.models.signature``.
input_example: {{ input_example }}
pip_requirements: {{ pip_requirements }}
extra_pip_requirements: {{ extra_pip_requirements }}
conda_env: {{ conda_env }}
metadata: {{ metadata }}
"""
from llama_index.core.indices.base import BaseIndex
from mlflow.llama_index.serialize_objects import serialize_settings
# TODO: make this logic cleaner and maybe a util
with tempfile.TemporaryDirectory() as temp_dir:
model_or_code_path = _validate_and_prepare_llama_index_model_or_path(
llama_index_model, temp_dir
)
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
path = os.path.abspath(path)
_validate_and_prepare_target_save_path(path)
if isinstance(model_config, str):
model_config = _validate_and_get_model_config_from_file(model_config)
model_code_path = None
if isinstance(model_or_code_path, str):
model_code_path = model_or_code_path
llama_index_model = _load_model_code_path(model_code_path, model_config)
_validate_and_copy_file_to_directory(model_code_path, path, "code")
# Warn when user provides `engine_type` argument while saving an engine directly
if not isinstance(llama_index_model, BaseIndex) and engine_type is not None:
_logger.warning(
"The `engine_type` argument is ignored when saving a non-index object."
)
elif isinstance(model_or_code_path, BaseIndex):
_validate_engine_type(engine_type)
llama_index_model = model_or_code_path
elif isinstance(model_or_code_path, _supported_classes()):
raise MlflowException.invalid_parameter_value(
"Saving a non-index object is only supported in the 'Model-from-Code' saving mode. "
"The legacy serialization method is exclusively for saving index objects. Please "
"pass the path to the script containing the model definition to save a non-index "
"object. For more information, see "
"https://www.mlflow.org/docs/latest/model/models-from-code.html",
)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
saved_example = _save_example(mlflow_model, input_example, path)
if signature is None and saved_example is not None:
wrapped_model = create_pyfunc_wrapper(llama_index_model, engine_type, model_config)
signature = _infer_signature_from_input_example(saved_example, wrapped_model)
elif signature is False:
signature = None
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if metadata is not None:
mlflow_model.metadata = metadata
# NB: llama_index.core.Settings is a singleton that manages the storage/service context
# for a given llama_index application. Given it holds the required objects for most of
# the index's functionality, we look to serialize the entire object. For components of
# the object that are not serializable, we log a warning.
settings_path = os.path.join(path, _SETTINGS_FILE)
serialize_settings(settings_path)
# Do not save the index/engine object in model-from-code saving mode
if not isinstance(model_code_path, str) and isinstance(llama_index_model, BaseIndex):
_save_index(llama_index_model, path)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.llama_index",
conda_env=_CONDA_ENV_FILE_NAME,
python_env=_PYTHON_ENV_FILE_NAME,
code=code_dir_subpath,
model_code_path=model_code_path,
model_config=model_config,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
llama_index_version=_get_llama_index_version(),
code=code_dir_subpath,
engine_type=engine_type,
)
if size := get_total_file_size(path):
mlflow_model.model_size_bytes = size
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
default_reqs = None
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
inferred_reqs = mlflow.models.infer_pip_requirements(
str(path), FLAVOR_NAME, fallback=default_reqs
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
_PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))
@experimental
@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
@trace_disabled # Suppress traces while loading model
def log_model(
llama_index_model,
artifact_path: str,
engine_type: Optional[str] = None,
model_config: Optional[dict[str, Any]] = None,
code_paths: Optional[list[str]] = None,
registered_model_name: Optional[str] = None,
signature: Optional[ModelSignature] = None,
input_example: Optional[ModelInputExample] = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
pip_requirements: Optional[Union[list[str], str]] = None,
extra_pip_requirements: Optional[Union[list[str], str]] = None,
conda_env=None,
metadata: Optional[dict[str, Any]] = None,
prompts: Optional[list[Union[str, Prompt]]] = None,
**kwargs,
):
"""
Log a LlamaIndex model as an MLflow artifact for the current run.
.. attention::
Saving a non-index object is only supported in the 'Model-from-Code' saving mode.
Please refer to the `Models From Code Guide <https://www.mlflow.org/docs/latest/model/models-from-code.html>`_
for more information.
.. note::
When logging a model, MLflow will automatically save the state of the ``Settings``
object so that you can use the same settings at inference time. However, please
note that some information in the ``Settings`` object will not be saved, including:
- API keys for avoiding key leakage.
- Function objects which are not serializable.
Args:
llama_index_model: A LlamaIndex object to be saved. Supported model types are:
1. An Index object.
2. An Engine object e.g. ChatEngine, QueryEngine, Retriever.
3. A `Workflow <https://docs.llamaindex.ai/en/stable/module_guides/workflow/>`_ object.
4. A string representing the path to a script contains LlamaIndex model definition
of the one of the above types.
artifact_path: Local path where the serialized model (as YAML) is to be saved.
engine_type: Required when saving an Index object to determine the inference interface
for the index when loaded as a pyfunc model. This field is **not** required when
saving other LlamaIndex objects. The supported values are as follows:
- ``"chat"``: load the index as an instance of the LlamaIndex
`ChatEngine <https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/>`_.
- ``"query"``: load the index as an instance of the LlamaIndex
`QueryEngine <https://docs.llamaindex.ai/en/stable/module_guides/deploying/query_engine/>`_.
- ``"retriever"``: load the index as an instance of the LlamaIndex
`Retriever <https://docs.llamaindex.ai/en/stable/module_guides/querying/retriever/>`_.
model_config: The model configuration to apply when loading the model back with
``mlflow.pyfunc.load_model()``. It will be applied in a different way depending on the
model type and saving method:
For in-memory Index objects saved directly, it will be passed as keyword arguments to
instantiate the LlamaIndex engine with the specified engine type at logging.
.. code-block:: python
with mlflow.start_run() as run:
model_info = mlflow.llama_index.log_model(
index,
artifact_path="index",
engine_type="chat",
model_config={"top_k": 10},
)
# When loading back, MLflow will call ``index.as_chat_engine(top_k=10)``
engine = mlflow.pyfunc.load_model(model_info.model_uri)
For other model types saved with the `Model-from-Code <https://www.mlflow.org/docs/latest/model/models-from-code.html>`
method, the config will be accessed via the :py:class`~mlflow.models.ModelConfig`
object within your model code.
.. code-block:: python
with mlflow.start_run() as run:
model_info = mlflow.llama_index.log_model(
"model.py",
artifact_path="model",
model_config={"qdrant_host": "localhost", "qdrant_port": 6333},
)
model.py:
.. code-block:: python
import mlflow
from llama_index.vector_stores.qdrant import QdrantVectorStore
import qdrant_client
# The model configuration is accessible via the ModelConfig singleton
model_config = mlflow.models.ModelConfig()
qdrant_host = model_config.get("top_k", 5)
qdrant_port = model_config.get("qdrant_port", 6333)
client = qdrant_client.Client(host=qdrant_host, port=qdrant_port)
vectorstore = QdrantVectorStore(client)
# the rest of the model definition...
code_paths: {{ code_paths }}
registered_model_name: This argument may change or be removed in a
future release without warning. If given, create a model
version under ``registered_model_name``, also creating a
registered model if one with the given name does not exist.
signature: A Model Signature object that describes the input and output Schema of the
model. The model signature can be inferred using ``infer_signature`` function
of `mlflow.models.signature`.
input_example: {{ input_example }}
await_registration_for: Number of seconds to wait for the model version
to finish being created and is in ``READY`` status.
By default, the function waits for five minutes.
Specify 0 or None to skip waiting.
pip_requirements: {{ pip_requirements }}
extra_pip_requirements: {{ extra_pip_requirements }}
conda_env: {{ conda_env }}
metadata: {{ metadata }}
prompts: {{ prompts }}
kwargs: Additional arguments for :py:class:`mlflow.models.model.Model`
"""
return Model.log(
artifact_path=artifact_path,
engine_type=engine_type,
model_config=model_config,
flavor=mlflow.llama_index,
registered_model_name=registered_model_name,
llama_index_model=llama_index_model,
conda_env=conda_env,
code_paths=code_paths,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
metadata=metadata,
prompts=prompts,
**kwargs,
)
def _validate_and_prepare_llama_index_model_or_path(llama_index_model, temp_dir=None):
if isinstance(llama_index_model, str):
return _validate_and_get_model_code_path(llama_index_model, temp_dir)
if not isinstance(llama_index_model, _supported_classes()):
supported_cls_names = [cls.__name__ for cls in _supported_classes()]
raise MlflowException.invalid_parameter_value(
message=f"The provided object of type {type(llama_index_model).__name__} is not "
"supported. MLflow llama-index flavor only supports saving LlamaIndex objects "
f"subclassed from one of the following classes: {supported_cls_names}.",
)
return llama_index_model
def _save_index(index, path):
"""Serialize the index."""
index_path = os.path.join(path, _INDEX_PERSIST_FOLDER)
index.storage_context.persist(persist_dir=index_path)
def _load_llama_model(path, flavor_conf):
"""Load the LlamaIndex index/engine/workflow from either model code or serialized index."""
from llama_index.core import StorageContext, load_index_from_storage
_add_code_from_conf_to_system_path(path, flavor_conf)
# Handle model-from-code
pyfunc_flavor_conf = _get_flavor_configuration(model_path=path, flavor_name=pyfunc.FLAVOR_NAME)
if model_code_path := pyfunc_flavor_conf.get(MODEL_CODE_PATH):
# TODO: The code path saved in the MLModel file is the local absolute path to the code
# file when it is saved. We should update the relative path in artifact directory.
model_code_path = os.path.join(path, os.path.basename(model_code_path))
model_config = pyfunc_flavor_conf.get(MODEL_CONFIG) or flavor_conf.get(MODEL_CONFIG, {})
if isinstance(model_config, str):
config_path = os.path.join(path, os.path.basename(model_config))
model_config = _validate_and_get_model_config_from_file(config_path)
return _load_model_code_path(model_code_path, model_config)
else:
# Use default vector store when loading from the serialized index
index_path = os.path.join(path, _INDEX_PERSIST_FOLDER)
storage_context = StorageContext.from_defaults(persist_dir=index_path)
return load_index_from_storage(storage_context)
@experimental
@trace_disabled # Suppress traces while loading model
def load_model(model_uri, dst_path=None):
"""
Load a LlamaIndex index/engine/workflow from a local file or a run.
Args:
model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``mlflow-artifacts:/path/to/model``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
artifact-locations>`_.
dst_path: The local filesystem path to utilize for downloading the model artifact.
This directory must already exist if provided. If unspecified, a local output
path will be created.
Returns:
A LlamaIndex index object.
"""
from mlflow.llama_index.serialize_objects import deserialize_settings
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
settings_path = os.path.join(local_model_path, _SETTINGS_FILE)
# NB: Settings is a singleton and can be loaded via llama_index.core.Settings
deserialize_settings(settings_path)
return _load_llama_model(local_model_path, flavor_conf)
def _load_pyfunc(path, model_config: Optional[dict[str, Any]] = None):
from mlflow.llama_index.pyfunc_wrapper import create_pyfunc_wrapper
index = load_model(path)
flavor_conf = _get_flavor_configuration(model_path=path, flavor_name=FLAVOR_NAME)
engine_type = flavor_conf.pop(
"engine_type", None
) # Not present when saving an non-index object
return create_pyfunc_wrapper(index, engine_type, model_config)
@experimental
def autolog(
log_traces: bool = True,
disable: bool = False,
silent: bool = False,
):
"""
Enables (or disables) and configures autologging from LlamaIndex to MLflow. Currently, MLflow
only supports autologging for tracing.
Args:
log_traces: If ``True``, traces are logged for LlamaIndex models by using. If ``False``,
no traces are collected during inference. Default to ``True``.
disable: If ``True``, disables the LlamaIndex autologging integration. If ``False``,
enables the LlamaIndex autologging integration.
silent: If ``True``, suppress all event logs and warnings from MLflow during LlamaIndex
autologging. If ``False``, show all events and warnings.
"""
from mlflow.llama_index.tracer import remove_llama_index_tracer, set_llama_index_tracer
# NB: The @autologging_integration annotation is used for adding shared logic. However, one
# caveat is that the wrapped function is NOT executed when disable=True is passed. This prevents
# us from running cleaning up logging when autologging is turned off. To workaround this, we
# annotate _autolog() instead of this entrypoint, and define the cleanup logic outside it.
# TODO: since this implementation is inconsistent, explore a universal way to solve the issue.
if log_traces and not disable:
set_llama_index_tracer()
else:
remove_llama_index_tracer()
_autolog(log_traces=log_traces, disable=disable, silent=silent)
# This is required by mlflow.autolog()
autolog.integration_name = FLAVOR_NAME
@autologging_integration(FLAVOR_NAME)
def _autolog(
log_traces: bool,
disable: bool = False,
silent: bool = False,
):
"""
TODO: Implement patching logic for autologging models and artifacts.
"""

BIN
venv/lib/python3.9/site-packages/mlflow/llama_index/__pycache__/__init__.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
venv/lib/python3.9/site-packages/mlflow/llama_index/__pycache__/chat.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
venv/lib/python3.9/site-packages/mlflow/llama_index/__pycache__/pyfunc_wrapper.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
venv/lib/python3.9/site-packages/mlflow/llama_index/__pycache__/serialize_objects.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
venv/lib/python3.9/site-packages/mlflow/llama_index/__pycache__/tracer.cpython-39.pyc Normal file
View File

Binary file not shown.

43
venv/lib/python3.9/site-packages/mlflow/llama_index/chat.py Normal file
View File

@@ -0,0 +1,43 @@
from llama_index.core.base.llms.types import ChatMessage as LLamaChatMessage
from llama_index.core.instrumentation.events import BaseEvent
from llama_index.core.instrumentation.events.llm import (
LLMChatEndEvent,
LLMChatStartEvent,
LLMCompletionEndEvent,
LLMCompletionStartEvent,
)
# llama-index includes llama-index-llms-openai in its requirements
# https://github.com/run-llama/llama_index/blob/663e1700f58c2414e549b9f5005abe87a275dd77/pyproject.toml#L52
from llama_index.llms.openai.utils import to_openai_message_dict
from mlflow.types.chat import ChatMessage
from mlflow.utils.pydantic_utils import model_dump_compat
def get_chat_messages_from_event(event: BaseEvent) -> list[ChatMessage]:
"""
Extract chat messages from the LlamaIndex callback event.
"""
if isinstance(event, LLMCompletionStartEvent):
return [ChatMessage(role="user", content=event.prompt)]
elif isinstance(event, LLMCompletionEndEvent):
return [ChatMessage(role="assistant", content=event.response.text)]
elif isinstance(event, LLMChatStartEvent):
return [_convert_message_to_mlflow_chat(msg) for msg in event.messages]
elif isinstance(event, LLMChatEndEvent):
message = event.response.message
return [_convert_message_to_mlflow_chat(message)]
raise ValueError(f"Unsupported event type for chat attribute extraction: {type(event)}")
def _convert_message_to_mlflow_chat(message: LLamaChatMessage) -> ChatMessage:
"""Convert a message object from LlamaIndex to MLflow's standard format."""
message = to_openai_message_dict(message, drop_none=False)
# tool calls are pydantic models in llama-index
if tool_calls := message.get("tool_calls"):
message["tool_calls"] = [model_dump_compat(tool) for tool in tool_calls]
return ChatMessage.validate_compat(message)

325
venv/lib/python3.9/site-packages/mlflow/llama_index/pyfunc_wrapper.py Normal file
View File

@@ -0,0 +1,325 @@
import asyncio
import threading
from typing import TYPE_CHECKING, Any, Optional, Union
if TYPE_CHECKING:
from llama_index.core import QueryBundle
from mlflow.models.utils import _convert_llm_input_data
CHAT_ENGINE_NAME = "chat"
QUERY_ENGINE_NAME = "query"
RETRIEVER_ENGINE_NAME = "retriever"
SUPPORTED_ENGINES = {CHAT_ENGINE_NAME, QUERY_ENGINE_NAME, RETRIEVER_ENGINE_NAME}
_CHAT_MESSAGE_HISTORY_PARAMETER_NAME = "chat_history"
def _convert_llm_input_data_with_unwrapping(data):
"""
Transforms the input data to the format expected by the LlamaIndex engine.
TODO: Migrate the unwrapping logic to mlflow.evaluate() function or _convert_llm_input_data,
# because it is not specific to LlamaIndex.
"""
data = _convert_llm_input_data(data)
# For mlflow.evaluate() call, the input dataset will be a pandas DataFrame. The DF should have
# a column named "inputs" which contains the actual query data. After the preprocessing, the
# each row will be passed here as a dictionary with the key "inputs". Therefore, we need to
# extract the actual query data from the dictionary.
if isinstance(data, dict) and ("inputs" in data):
data = data["inputs"]
return data
def _format_predict_input_query_engine_and_retriever(data) -> "QueryBundle":
"""Convert pyfunc input to a QueryBundle."""
from llama_index.core import QueryBundle
data = _convert_llm_input_data_with_unwrapping(data)
if isinstance(data, str):
return QueryBundle(query_str=data)
elif isinstance(data, dict):
return QueryBundle(**data)
elif isinstance(data, list):
# NB: handle pandas returning lists when there is a single row
prediction_input = [_format_predict_input_query_engine_and_retriever(d) for d in data]
return prediction_input if len(prediction_input) > 1 else prediction_input[0]
else:
raise ValueError(
f"Unsupported input type: {type(data)}. It must be one of "
"[str, dict, list, numpy.ndarray, pandas.DataFrame]"
)
class _LlamaIndexModelWrapperBase:
def __init__(
self,
llama_model, # Engine or Workflow
model_config: Optional[dict[str, Any]] = None,
):
self._llama_model = llama_model
self.model_config = model_config or {}
@property
def index(self):
return self._llama_model.index
def get_raw_model(self):
return self._llama_model
def _predict_single(self, *args, **kwargs) -> Any:
raise NotImplementedError
def _format_predict_input(self, data):
raise NotImplementedError
def _do_inference(self, input, params: Optional[dict[str, Any]]) -> dict:
"""
Perform engine inference on a single engine input e.g. not an iterable of
engine inputs. The engine inputs must already be preprocessed/cleaned.
"""
if isinstance(input, dict):
return self._predict_single(**input, **(params or {}))
else:
return self._predict_single(input, **(params or {}))
def predict(self, data, params: Optional[dict[str, Any]] = None) -> Union[list[str], str]:
data = self._format_predict_input(data)
if isinstance(data, list):
return [self._do_inference(x, params) for x in data]
else:
return self._do_inference(data, params)
class ChatEngineWrapper(_LlamaIndexModelWrapperBase):
@property
def engine_type(self):
return CHAT_ENGINE_NAME
def _predict_single(self, *args, **kwargs) -> str:
return self._llama_model.chat(*args, **kwargs).response
@staticmethod
def _convert_chat_message_history_to_chat_message_objects(data: dict) -> dict:
from llama_index.core.llms import ChatMessage
if chat_message_history := data.get(_CHAT_MESSAGE_HISTORY_PARAMETER_NAME):
if isinstance(chat_message_history, list):
if all(isinstance(message, dict) for message in chat_message_history):
data[_CHAT_MESSAGE_HISTORY_PARAMETER_NAME] = [
ChatMessage(**message) for message in chat_message_history
]
else:
raise ValueError(
f"Unsupported input type: {type(chat_message_history)}. "
"It must be a list of dicts."
)
return data
def _format_predict_input(self, data) -> Union[str, dict, list]:
data = _convert_llm_input_data_with_unwrapping(data)
if isinstance(data, str):
return data
elif isinstance(data, dict):
return self._convert_chat_message_history_to_chat_message_objects(data)
elif isinstance(data, list):
# NB: handle pandas returning lists when there is a single row
prediction_input = [self._format_predict_input(d) for d in data]
return prediction_input if len(prediction_input) > 1 else prediction_input[0]
else:
raise ValueError(
f"Unsupported input type: {type(data)}. It must be one of "
"[str, dict, list, numpy.ndarray, pandas.DataFrame]"
)
class QueryEngineWrapper(_LlamaIndexModelWrapperBase):
@property
def engine_type(self):
return QUERY_ENGINE_NAME
def _predict_single(self, *args, **kwargs) -> str:
return self._llama_model.query(*args, **kwargs).response
def _format_predict_input(self, data) -> "QueryBundle":
return _format_predict_input_query_engine_and_retriever(data)
class RetrieverEngineWrapper(_LlamaIndexModelWrapperBase):
@property
def engine_type(self):
return RETRIEVER_ENGINE_NAME
def _predict_single(self, *args, **kwargs) -> list[dict]:
response = self._llama_model.retrieve(*args, **kwargs)
return [node.dict() for node in response]
def _format_predict_input(self, data) -> "QueryBundle":
return _format_predict_input_query_engine_and_retriever(data)
class WorkflowWrapper(_LlamaIndexModelWrapperBase):
@property
def index(self):
raise NotImplementedError("LlamaIndex Workflow does not have an index")
@property
def engine_type(self):
raise NotImplementedError("LlamaIndex Workflow is not an engine")
def predict(self, data, params: Optional[dict[str, Any]] = None) -> Union[list[str], str]:
inputs = self._format_predict_input(data, params)
# LlamaIndex Workflow runs async but MLflow pyfunc doesn't support async inference yet.
predictions = self._wait_async_task(self._run_predictions(inputs))
# Even if the input is single instance, the signature enforcement convert it to a Pandas
# DataFrame with a single row. In this case, we should unwrap the result (list) so it
# won't be inconsistent with the output without signature enforcement.
should_unwrap = len(data) == 1 and isinstance(predictions, list)
return predictions[0] if should_unwrap else predictions
def _format_predict_input(self, data, params: Optional[dict[str, Any]] = None) -> list[dict]:
inputs = _convert_llm_input_data_with_unwrapping(data)
params = params or {}
if isinstance(inputs, dict):
return [{**inputs, **params}]
return [{**x, **params} for x in inputs]
async def _run_predictions(self, inputs: list[dict[str, Any]]) -> asyncio.Future:
tasks = [self._predict_single(x) for x in inputs]
return await asyncio.gather(*tasks)
async def _predict_single(self, x: dict[str, Any]) -> Any:
if not isinstance(x, dict):
raise ValueError(f"Unsupported input type: {type(x)}. It must be a dictionary.")
return await self._llama_model.run(**x)
def _wait_async_task(self, task: asyncio.Future) -> Any:
"""
A utility function to run async tasks in a blocking manner.
If there is no event loop running already, for example, in a model serving endpoint,
we can simply create a new event loop and run the task there. However, in a notebook
environment (or pytest with asyncio decoration), there is already an event loop running
at the root level and we cannot start a new one.
"""
if not self._is_event_loop_running():
return asyncio.new_event_loop().run_until_complete(task)
else:
# NB: The popular way to run async task where an event loop is already running is to
# use nest_asyncio. However, nest_asyncio.apply() breaks the async OpenAI client
# somehow, which is used for the most of LLM calls in LlamaIndex including Databricks
# LLMs. Therefore, we use a hacky workaround that creates a new thread and run the
# new event loop there. This may degrade the performance compared to the native
# asyncio, but it should be fine because this is only used in the notebook env.
results = None
exception = None
def _run():
nonlocal results, exception
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
results = loop.run_until_complete(task)
except Exception as e:
exception = e
finally:
loop.close()
thread = threading.Thread(target=_run)
thread.start()
thread.join()
if exception:
raise exception
return results
def _is_event_loop_running(self) -> bool:
try:
loop = asyncio.get_running_loop()
return loop is not None
except Exception:
return False
def create_pyfunc_wrapper(
model: Any,
engine_type: Optional[str] = None,
model_config: Optional[dict[str, Any]] = None,
):
"""
A factory function that creates a Pyfunc wrapper around a LlamaIndex index/engine/workflow.
Args:
model: A LlamaIndex index/engine/workflow.
engine_type: The type of the engine. Only required if `model` is an index
and must be one of [chat, query, retriever].
model_config: A dictionary of model configuration parameters.
"""
try:
from llama_index.core.workflow import Workflow
if isinstance(model, Workflow):
return _create_wrapper_from_workflow(model, model_config)
except ImportError:
pass
from llama_index.core.indices.base import BaseIndex
if isinstance(model, BaseIndex):
return _create_wrapper_from_index(model, engine_type, model_config)
else:
# Engine does not have a common base class so we assume
# everything else is an engine
return _create_wrapper_from_engine(model, model_config)
def _create_wrapper_from_index(
index, engine_type: str, model_config: Optional[dict[str, Any]] = None
):
model_config = model_config or {}
if engine_type == QUERY_ENGINE_NAME:
engine = index.as_query_engine(**model_config)
return QueryEngineWrapper(engine, model_config)
elif engine_type == CHAT_ENGINE_NAME:
engine = index.as_chat_engine(**model_config)
return ChatEngineWrapper(engine, model_config)
elif engine_type == RETRIEVER_ENGINE_NAME:
engine = index.as_retriever(**model_config)
return RetrieverEngineWrapper(engine, model_config)
else:
raise ValueError(
f"Unsupported engine type: {engine_type}. It must be one of {SUPPORTED_ENGINES}"
)
def _create_wrapper_from_engine(engine: Any, model_config: Optional[dict[str, Any]] = None):
from llama_index.core.base.base_query_engine import BaseQueryEngine
from llama_index.core.chat_engine.types import BaseChatEngine
from llama_index.core.retrievers import BaseRetriever
if isinstance(engine, BaseChatEngine):
return ChatEngineWrapper(engine, model_config)
elif isinstance(engine, BaseQueryEngine):
return QueryEngineWrapper(engine, model_config)
elif isinstance(engine, BaseRetriever):
return RetrieverEngineWrapper(engine, model_config)
else:
raise ValueError(
f"Unsupported engine type: {type(engine)}. It must be one of {SUPPORTED_ENGINES}"
)
def _create_wrapper_from_workflow(workflow: Any, model_config: Optional[dict[str, Any]] = None):
return WorkflowWrapper(workflow, model_config)

188
venv/lib/python3.9/site-packages/mlflow/llama_index/serialize_objects.py Normal file
View File

@@ -0,0 +1,188 @@
import importlib
import inspect
import json
import logging
from typing import Any, Callable
from llama_index.core import PromptTemplate
from llama_index.core.base.embeddings.base import BaseEmbedding
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.schema import BaseComponent
_logger = logging.getLogger(__name__)
def _get_object_import_path(o: object) -> str:
if not inspect.isclass(o):
o = o.__class__
module_name = inspect.getmodule(o).__name__
class_name = o.__qualname__
# Validate the import
module = importlib.import_module(module_name)
if not hasattr(module, class_name):
raise ValueError(f"Module {module} does not have {class_name}")
return f"{module_name}.{class_name}"
def _sanitize_api_key(object_as_dict: dict[str, str]) -> dict[str, str]:
return {k: v for k, v in object_as_dict.items() if "api_key" not in k.lower()}
def object_to_dict(o: object):
if isinstance(o, (list, tuple)):
return [object_to_dict(v) for v in o]
if isinstance(o, BaseComponent):
# we can't serialize callables in the model fields
callable_fields = set()
fields = o.model_fields if hasattr(o, "model_fields") else o.__fields__
for k, v in fields.items():
field_val = getattr(o, k, None)
if field_val != v.default and callable(field_val):
callable_fields.add(k)
# exclude default values from serialization to avoid
# unnecessary clutter in the serialized object
o_state_as_dict = o.to_dict(exclude=callable_fields)
if o_state_as_dict != {}:
o_state_as_dict = _sanitize_api_key(o_state_as_dict)
o_state_as_dict.pop("class_name")
else:
return o_state_as_dict
return {
"object_constructor": _get_object_import_path(o),
"object_kwargs": o_state_as_dict,
}
else:
return None
def _construct_prompt_template_object(
constructor: Callable, kwargs: dict[str, Any]
) -> PromptTemplate:
"""Construct a PromptTemplate object based on the constructor and kwargs.
This method is necessary because the `template_vars` cannot be passed directly to the
constructor and needs to be set on an instantiated object.
"""
if template := kwargs.pop("template", None):
prompt_template = constructor(template)
for k, v in kwargs.items():
setattr(prompt_template, k, v)
return prompt_template
else:
raise ValueError(
"'template' is a required kwargs and is not present in the prompt template kwargs."
)
def dict_to_object(object_representation: dict[str, Any]) -> object:
if "object_constructor" not in object_representation:
raise ValueError("'object_constructor' key not found in dict.")
if "object_kwargs" not in object_representation:
raise ValueError("'object_kwargs' key not found in dict.")
constructor_str = object_representation["object_constructor"]
kwargs = object_representation["object_kwargs"]
import_path, class_name = constructor_str.rsplit(".", 1)
module = importlib.import_module(import_path)
if isinstance(module, PromptTemplate):
return _construct_prompt_template_object(module, kwargs)
else:
object_class = getattr(module, class_name)
# Many embeddings model accepts parameter `model`, while BaseEmbedding accepts `model_name`.
# Both parameters will be serialized as kwargs, but passing both to the constructor will
# raise duplicate argument error. Some class like OpenAIEmbedding handles this in its
# constructor, but not all integrations do. Therefore, we have to handle it here.
# E.g. https://github.com/run-llama/llama_index/blob/2b18eb4654b14c68d63f6239cddb10740668fbc8/llama-index-integrations/embeddings/llama-index-embeddings-openai/llama_index/embeddings/openai/base.py#L316-L320
if (
issubclass(object_class, BaseEmbedding)
and (model := kwargs.get("model"))
and (model_name := kwargs.get("model_name"))
and model == model_name
):
kwargs.pop("model_name")
return object_class.from_dict(kwargs)
def _deserialize_dict_of_objects(path: str) -> dict[str, Any]:
with open(path) as f:
to_deserialize = json.load(f)
output = {}
for k, v in to_deserialize.items():
if isinstance(v, list):
output.update({k: [dict_to_object(vv) for vv in v]})
else:
output.update({k: dict_to_object(v)})
return output
def serialize_settings(path: str) -> None:
"""Serialize the global LlamaIndex Settings object to a JSON file at the given path."""
from llama_index.core import Settings
_logger.info(
"API key(s) will be removed from the global Settings object during serialization "
"to protect against key leakage. At inference time, the key(s) must be passed as "
"environment variables."
)
to_serialize = {}
unsupported_objects = []
for k, v in Settings.__dict__.items():
if v is None:
continue
# Setting.callback_manager is default to an empty CallbackManager instance.
if (k == "_callback_manager") and isinstance(v, CallbackManager) and v.handlers == []:
continue
def _convert(obj):
object_json = object_to_dict(obj)
if object_json is None:
prop_name = k[1:] if k.startswith("_") else k
unsupported_objects.append((prop_name, v))
return object_json
if isinstance(v, list):
to_serialize[k] = [_convert(obj) for obj in v if v is not None]
else:
if (object_json := _convert(v)) and (object_json is not None):
to_serialize[k] = object_json
if unsupported_objects:
msg = (
"The following objects in Settings are not supported for serialization and will not "
"be logged with your model. MLflow only supports serialization of objects that inherit "
"from llama_index.core.schema.BaseComponent.\n"
)
msg += "\n".join(f" - {type(v).__name__} for Settings.{k}" for k, v in unsupported_objects)
_logger.info(msg)
with open(path, "w") as f:
json.dump(to_serialize, f, indent=2)
def deserialize_settings(path: str):
"""Deserialize the global LlamaIndex Settings object from a JSON file at the given path."""
settings_dict = _deserialize_dict_of_objects(path)
from llama_index.core import Settings
for k, v in settings_dict.items():
# To use the property setter rather than directly setting the private attribute e.g. _llm
if k.startswith("_"):
k = k[1:]
setattr(Settings, k, v)

559
venv/lib/python3.9/site-packages/mlflow/llama_index/tracer.py Normal file
View File

@@ -0,0 +1,559 @@
import inspect
import json
import logging
from functools import singledispatchmethod
from typing import Any, Generator, Optional, Union
import llama_index.core
import pydantic
from llama_index.core.base.agent.types import BaseAgent, BaseAgentWorker, TaskStepOutput
from llama_index.core.base.base_retriever import BaseRetriever
from llama_index.core.base.embeddings.base import BaseEmbedding
from llama_index.core.base.llms.base import BaseLLM
from llama_index.core.base.llms.types import ChatResponse, CompletionResponse
from llama_index.core.base.response.schema import AsyncStreamingResponse, StreamingResponse
from llama_index.core.chat_engine.types import StreamingAgentChatResponse
from llama_index.core.instrumentation.event_handlers import BaseEventHandler
from llama_index.core.instrumentation.events import BaseEvent
from llama_index.core.instrumentation.events.agent import AgentToolCallEvent
from llama_index.core.instrumentation.events.embedding import EmbeddingStartEvent
from llama_index.core.instrumentation.events.exception import ExceptionEvent
from llama_index.core.instrumentation.events.llm import (
LLMChatEndEvent,
LLMChatStartEvent,
LLMCompletionEndEvent,
LLMCompletionStartEvent,
LLMPredictStartEvent,
)
from llama_index.core.instrumentation.events.rerank import ReRankStartEvent
from llama_index.core.instrumentation.span.base import BaseSpan
from llama_index.core.instrumentation.span_handlers import BaseSpanHandler
from llama_index.core.multi_modal_llms import MultiModalLLM
from llama_index.core.schema import NodeWithScore
from llama_index.core.tools import BaseTool
from packaging.version import Version
import mlflow
from mlflow.entities import LiveSpan, SpanEvent, SpanType
from mlflow.entities.document import Document
from mlflow.entities.span_status import SpanStatusCode
from mlflow.llama_index.chat import get_chat_messages_from_event
from mlflow.tracing.constant import SpanAttributeKey
from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
from mlflow.tracking.client import MlflowClient
from mlflow.utils.pydantic_utils import model_dump_compat
_logger = logging.getLogger(__name__)
def _get_llama_index_version() -> Version:
return Version(llama_index.core.__version__)
def set_llama_index_tracer():
"""
Set the MlflowSpanHandler and MlflowEventHandler to the global dispatcher.
If the handlers are already set, skip setting.
"""
from llama_index.core.instrumentation import get_dispatcher
dsp = get_dispatcher()
span_handler = None
for handler in dsp.span_handlers:
if isinstance(handler, MlflowSpanHandler):
_logger.debug("MlflowSpanHandler is already set to the dispatcher. Skip setting.")
span_handler = handler
break
else:
span_handler = MlflowSpanHandler()
dsp.add_span_handler(span_handler)
for handler in dsp.event_handlers:
if isinstance(handler, MlflowEventHandler):
_logger.debug("MlflowEventHandler is already set to the dispatcher. Skip setting.")
break
else:
dsp.add_event_handler(MlflowEventHandler(span_handler))
def remove_llama_index_tracer():
"""
Remove the MlflowSpanHandler and MlflowEventHandler from the global dispatcher.
"""
from llama_index.core.instrumentation import get_dispatcher
dsp = get_dispatcher()
dsp.span_handlers = [h for h in dsp.span_handlers if h.class_name() != "MlflowSpanHandler"]
dsp.event_handlers = [h for h in dsp.event_handlers if h.class_name() != "MlflowEventHandler"]
class _LlamaSpan(BaseSpan, extra="allow"):
_mlflow_span: LiveSpan = pydantic.PrivateAttr()
def __init__(self, id_: str, parent_id: Optional[str], mlflow_span: LiveSpan):
super().__init__(id_=id_, parent_id=parent_id)
self._mlflow_span = mlflow_span
def _end_span(span: LiveSpan, status=SpanStatusCode.OK, outputs=None, token=None):
"""An utility function to end the span or trace."""
if isinstance(outputs, (StreamingResponse, AsyncStreamingResponse, StreamingAgentChatResponse)):
_logger.warning(
"Trying to record streaming response to the MLflow trace. This may consume "
"the generator and result in an empty response."
)
# for retriever spans, convert the outputs to Document objects
# so they can be rendered in a more user-friendly way in the UI
if (
span.span_type == SpanType.RETRIEVER
and isinstance(outputs, list)
and all(isinstance(item, NodeWithScore) for item in outputs)
):
try:
outputs = [Document.from_llama_index_node_with_score(node) for node in outputs]
except Exception as e:
_logger.debug(
f"Failed to convert NodeWithScore to Document objects: {e}", exc_info=True
)
if outputs is None:
outputs = span.outputs
try:
if span.parent_id is None:
# NB: Initiate the new client every time to handle tracking URI updates.
MlflowClient().end_trace(span.request_id, status=status, outputs=outputs)
else:
MlflowClient().end_span(span.request_id, span.span_id, status=status, outputs=outputs)
finally:
# We should detach span even when end_span / end_trace API call fails
if token:
detach_span_from_context(token)
class MlflowSpanHandler(BaseSpanHandler[_LlamaSpan], extra="allow"):
def __init__(self):
super().__init__()
self._span_id_to_token = {}
self._stream_resolver = StreamResolver()
self._pending_spans: dict[str, _LlamaSpan] = {}
@classmethod
def class_name(cls) -> str:
return "MlflowSpanHandler"
def get_span_for_event(self, event: BaseEvent) -> LiveSpan:
llama_span = self.open_spans.get(event.span_id) or self._pending_spans.get(event.span_id)
return llama_span._mlflow_span if llama_span else None
def new_span(
self,
id_: str,
bound_args: inspect.BoundArguments,
instance: Optional[Any] = None,
parent_span_id: Optional[str] = None,
**kwargs: Any,
) -> _LlamaSpan:
with self.lock:
parent = self.open_spans.get(parent_span_id) if parent_span_id else None
parent_span = parent._mlflow_span if parent else mlflow.get_current_active_span()
try:
input_args = bound_args.arguments
attributes = self._get_instance_attributes(instance)
span_type = self._get_span_type(instance) or SpanType.UNKNOWN
if parent_span:
# NB: Initiate the new client every time to handle tracking URI updates.
span = MlflowClient().start_span(
request_id=parent_span.request_id,
parent_id=parent_span.span_id,
name=id_.partition("-")[0],
span_type=span_type,
inputs=input_args,
attributes=attributes,
)
else:
span = MlflowClient().start_trace(
name=id_.partition("-")[0],
span_type=span_type,
inputs=input_args,
attributes=attributes,
)
token = set_span_in_context(span)
self._span_id_to_token[span.span_id] = token
# NB: The tool definition is passed to LLM via kwargs, but it is not set
# to the LLM/Chat start event. Therefore, we need to handle it here.
tools = input_args.get("kwargs", {}).get("tools")
if tools and span_type in [SpanType.LLM, SpanType.CHAT_MODEL]:
try:
set_span_chat_tools(span, tools)
except Exception as e:
_logger.debug(f"Failed to set tools for {span}: {e}")
return _LlamaSpan(id_=id_, parent_id=parent_span_id, mlflow_span=span)
except BaseException as e:
_logger.debug(f"Failed to create a new span: {e}", exc_info=True)
def prepare_to_exit_span(
self,
id_: str,
result: Optional[Any] = None,
**kwargs: Any,
) -> _LlamaSpan:
try:
with self.lock:
llama_span = self.open_spans.get(id_)
if not llama_span:
return
span = llama_span._mlflow_span
token = self._span_id_to_token.pop(span.span_id, None)
if self._stream_resolver.is_streaming_result(result):
# If the result is a generator, we keep the span in progress for streaming
# and end it when the generator is exhausted.
is_pended = self._stream_resolver.register_stream_span(span, result)
if is_pended:
self._pending_spans[id_] = llama_span
# We still need to detach the span from the context, otherwise it will
# be considered as "active"
detach_span_from_context(token)
else:
# If the span is not pended successfully, end it immediately
_end_span(span=span, outputs=result, token=token)
else:
_end_span(span=span, outputs=result, token=token)
return llama_span
except BaseException as e:
_logger.debug(f"Failed to end a span: {e}", exc_info=True)
def resolve_pending_stream_span(self, span: LiveSpan, event: Any):
"""End the pending streaming span(s)"""
self._stream_resolver.resolve(span, event)
self._pending_spans.pop(event.span_id, None)
def prepare_to_drop_span(self, id_: str, err: Optional[Exception], **kwargs) -> _LlamaSpan:
"""Logic for handling errors during the model execution."""
with self.lock:
llama_span = self.open_spans.get(id_)
span = llama_span._mlflow_span
token = self._span_id_to_token.pop(span.span_id, None)
if _get_llama_index_version() >= Version("0.10.59"):
# LlamaIndex determines if a workflow is terminated or not by propagating an special
# exception WorkflowDone. We should treat this exception as a successful termination.
from llama_index.core.workflow.errors import WorkflowDone
if err and isinstance(err, WorkflowDone):
return _end_span(span=span, status=SpanStatusCode.OK, token=token)
span.add_event(SpanEvent.from_exception(err))
_end_span(span=span, status="ERROR", token=token)
return llama_span
def _get_span_type(self, instance: Any) -> SpanType:
"""
Map LlamaIndex instance type to MLflow span type. Some span type cannot be determined
by instance type alone, rather need event info e.g. ChatModel, ReRanker
"""
if isinstance(instance, (BaseLLM, MultiModalLLM)):
return SpanType.LLM
elif isinstance(instance, BaseRetriever):
return SpanType.RETRIEVER
elif isinstance(instance, (BaseAgent, BaseAgentWorker)):
return SpanType.AGENT
elif isinstance(instance, BaseEmbedding):
return SpanType.EMBEDDING
elif isinstance(instance, BaseTool):
return SpanType.TOOL
else:
return SpanType.CHAIN
@singledispatchmethod
def _get_instance_attributes(self, instance: Any) -> dict[str, Any]:
"""
Extract span attributes from LlamaIndex objects.
NB: There are some overlap between attributes extracted from instance metadata and the
events. For example, model name for an LLM is available in both. However, events might
not always be triggered (e.g. 3P llm integration doesn't implement the event logic),
so the instance metadata serves as a fallback source of information.
"""
# TODO: Union type hint doesn't work with singledispatchmethod, so we have to define
# two separate methods for BaseLLM and MultiModalLLM. Once we upgrade to Python 3.10,
# we can use `BaseLLM | MultiModelLLM` type hint and it works with singledispatchmethod.
@_get_instance_attributes.register
def _(self, instance: BaseLLM):
return self._get_llm_attributes(instance)
@_get_instance_attributes.register
def _(self, instance: MultiModalLLM):
return self._get_llm_attributes(instance)
def _get_llm_attributes(self, instance) -> dict[str, Any]:
attr = {}
if metadata := instance.metadata:
attr["model_name"] = metadata.model_name
if params_str := metadata.json(exclude_unset=True):
attr["invocation_params"] = json.loads(params_str)
return attr
@_get_instance_attributes.register
def _(self, instance: BaseEmbedding):
return {
"model_name": instance.model_name,
"embed_batch_size": instance.embed_batch_size,
}
@_get_instance_attributes.register
def _(self, instance: BaseTool):
metadata = instance.metadata
attributes = {"description": metadata.description}
try:
attributes["name"] = metadata.name
except ValueError:
# ToolMetadata.get_name() raises ValueError if name is None
pass
try:
attributes["parameters"] = json.loads(metadata.fn_schema_str)
except ValueError:
# ToolMetadata.get_fn_schema_str() raises ValueError if fn_schema is None
pass
return attributes
class MlflowEventHandler(BaseEventHandler, extra="allow"):
"""
Event handler processes various events that are triggered during execution.
Events are used as supplemental source for recording additional metadata to the span,
such as model name, parameters to the span, because they are not available in the inputs
and outputs in SpanHandler.
"""
_span_handler: MlflowSpanHandler
@classmethod
def class_name(cls) -> str:
return "MlflowEventHandler"
def __init__(self, _span_handler):
super().__init__()
self._span_handler = _span_handler
def handle(self, event: BaseEvent) -> Any:
try:
if span := self._span_handler.get_span_for_event(event):
self._handle_event(event, span)
except Exception as e:
_logger.debug(f"Failed to handle event: {e}", exc_info=True)
@singledispatchmethod
def _handle_event(self, event: BaseEvent, span: LiveSpan):
# Pass through the events we are not interested in
pass
@_handle_event.register
def _(self, event: AgentToolCallEvent, span: LiveSpan):
span.set_attribute("name", event.tool.name)
span.set_attribute("description", event.tool.description)
span.set_attribute("parameters", event.tool.get_parameters_dict())
@_handle_event.register
def _(self, event: EmbeddingStartEvent, span: LiveSpan):
span.set_attribute("model_dict", event.model_dict)
@_handle_event.register
def _(self, event: LLMPredictStartEvent, span: LiveSpan):
"""
An event triggered when LLM's predict() is called.
In LlamaIndex, predict() is a gateway method that dispatch the request to
either chat() or completion() method depending on the model type, as well
as crafting prompt from the template.
"""
template = event.template
template_args = {
**template.kwargs,
**(event.template_args if event.template_args else {}),
}
span.set_attributes(
{
"prmopt_template": template.get_template(),
"template_arguments": {var: template_args.get(var) for var in template_args},
}
)
@_handle_event.register
def _(self, event: LLMCompletionStartEvent, span: LiveSpan):
span.set_attribute("prompt", event.prompt)
span.set_attribute("model_dict", event.model_dict)
self._extract_and_set_chat_messages(span, event)
@_handle_event.register
def _(self, event: LLMCompletionEndEvent, span: LiveSpan):
span.set_attribute("usage", self._extract_token_usage(event.response))
self._extract_and_set_chat_messages(span, event)
self._span_handler.resolve_pending_stream_span(span, event)
@_handle_event.register
def _(self, event: LLMChatStartEvent, span: LiveSpan):
span.set_attribute(SpanAttributeKey.SPAN_TYPE, SpanType.CHAT_MODEL)
span.set_attribute("model_dict", event.model_dict)
self._extract_and_set_chat_messages(span, event)
@_handle_event.register
def _(self, event: LLMChatEndEvent, span: LiveSpan):
span.set_attribute("usage", self._extract_token_usage(event.response))
self._extract_and_set_chat_messages(span, event)
self._span_handler.resolve_pending_stream_span(span, event)
@_handle_event.register
def _(self, event: ReRankStartEvent, span: LiveSpan):
span.set_attribute(SpanAttributeKey.SPAN_TYPE, SpanType.RERANKER)
span.set_attributes(
{
"model_name": event.model_name,
"top_n": event.top_n,
}
)
@_handle_event.register
def _(self, event: ExceptionEvent, span: LiveSpan):
"""
Handle an exception event for stream spans.
For non-stream spans, exception is processed by the prepare_to_drop_span() handler of
the span handler. However, for stream spans, the exception may raised during the
streaming after it exit. Therefore, we need to resolve the span here.
"""
self._span_handler.resolve_pending_stream_span(span, event)
def _extract_token_usage(
self, response: Union[ChatResponse, CompletionResponse]
) -> dict[str, int]:
if raw := response.raw:
# The raw response can be a Pydantic model or a dictionary
if isinstance(raw, pydantic.BaseModel):
raw = model_dump_compat(raw)
if usage := raw.get("usage"):
return usage
# If the usage is not found in the raw response, look for token counts
# in additional_kwargs of the completion payload
usage = {}
if additional_kwargs := getattr(response, "additional_kwargs", None):
for k in ["prompt_tokens", "completion_tokens", "total_tokens"]:
if (v := additional_kwargs.get(k)) is not None:
usage[k] = v
return usage
def _extract_and_set_chat_messages(self, span: LiveSpan, event: BaseEvent):
try:
messages = get_chat_messages_from_event(event)
set_span_chat_messages(span, messages, append=True)
except Exception as e:
_logger.debug(f"Failed to set chat messages to the span: {e}", exc_info=True)
_StreamEndEvent = Union[LLMChatEndEvent, LLMCompletionEndEvent, ExceptionEvent]
class StreamResolver:
"""
A class is responsible for closing the pending streaming spans that are waiting
for the stream to be exhausted. Once the associated stream is exhausted, this
class will resolve the span, as well as recursively resolve the parent spans
that returns the same (or derived) stream.
"""
def __init__(self):
self._span_id_to_span_and_gen: dict[str, tuple[LiveSpan, Generator]] = {}
def is_streaming_result(self, result: Any) -> bool:
return (
inspect.isgenerator(result) # noqa: SIM101
or isinstance(result, (StreamingResponse, AsyncStreamingResponse))
or isinstance(result, StreamingAgentChatResponse)
or (isinstance(result, TaskStepOutput) and self.is_streaming_result(result.output))
)
def register_stream_span(self, span: LiveSpan, result: Any) -> bool:
"""
Register the pending streaming span with the associated generator.
Args:
span: The span that has a streaming output.
result: The streaming result that is being processed.
Returns:
True if the span is registered successfully, False otherwise.
"""
if inspect.isgenerator(result):
stream = result
elif isinstance(result, (StreamingResponse, AsyncStreamingResponse)):
stream = result.response_gen
elif isinstance(result, StreamingAgentChatResponse):
stream = result.chat_stream
elif isinstance(result, TaskStepOutput):
stream = result.output.chat_stream
else:
raise ValueError(f"Unsupported streaming response type: {type(result)}")
if inspect.getgeneratorstate(stream) == inspect.GEN_CLOSED:
# Not registering the span because the generator is already exhausted.
# It's counter-intuitive that the generator is closed before the response
# is returned, but it can happen because some agents run streaming request
# in a separate thread. In this case, the generator can be closed before
# the response is returned in the main thread.
return False
self._span_id_to_span_and_gen[span.span_id] = (span, stream)
return True
def resolve(self, span: LiveSpan, event: _StreamEndEvent):
"""
Finish the streaming span and recursively resolve the parent spans that
returns the same (or derived) stream.
"""
_, stream = self._span_id_to_span_and_gen.pop(span.span_id, (None, None))
if not stream:
return
if isinstance(event, (LLMChatEndEvent, LLMCompletionEndEvent)):
outputs = event.response
status = SpanStatusCode.OK
elif isinstance(event, ExceptionEvent):
outputs = None
status = SpanStatusCode.ERROR
span.add_event(SpanEvent.from_exception(event.exception))
else:
raise ValueError(f"Unsupported event type to resolve streaming: {type(event)}")
_end_span(span=span, status=status, outputs=outputs)
# Extract the complete text from the event.
if isinstance(outputs, ChatResponse):
output_text = outputs.message.content
elif isinstance(outputs, CompletionResponse):
output_text = outputs.response.text
else:
output_text = None
# Recursively resolve the parent spans that are also waiting for the same token
# stream to be exhausted.
while span.parent_id in self._span_id_to_span_and_gen:
if span_and_stream := self._span_id_to_span_and_gen.pop(span.parent_id, None):
span, stream = span_and_stream
# We reuse the same output text for parent spans. This may not be 100% correct
# as token stream can be modified by callers. However, it is technically
# challenging to track the modified stream across multiple spans.
_end_span(span=span, status=status, outputs=output_text)