[MLflow Tracing](/genai/tracing) allows you to monitor and log the execution of your model's methods, providing valuable insights during debugging and performance optimization.

In our example `BasicAgent` implementation we utilize two separate APIs for the initiation of trace spans: the decorator API and the fluent API.

#### Decorator API

```python
@mlflow.trace
def _get_system_message(self, role: str) -> Dict:
    if role not in self.models:
        raise ValueError(f"Unknown role: {role}")

    instruction = self.models[role]["instruction"]
    return ChatMessage(role="system", content=instruction).to_dict()
```

Using the <APILink fn="mlflow.trace">`@mlflow.trace`</APILink> tracing decorator is the simplest way to add tracing functionality to functions and methods. By default, a span that is generated from
the application of this decorator will utilize the name of the function as the name of the span. It is possible to override this naming, as well as
other parameters associated with the span, as follows:

```python
@mlflow.trace(name="custom_span_name", attributes={"key": "value"}, span_type="func")
def _get_system_message(self, role: str) -> Dict:
    if role not in self.models:
        raise ValueError(f"Unknown role: {role}")

    instruction = self.models[role]["instruction"]
    return ChatMessage(role="system", content=instruction).to_dict()
```

:::tip
It is always advised to set a human-readable name for any span that you generate, particularly if you are instrumenting private or generically
named functions or methods. The MLflow Trace UI will display the name of the function or method by default, which can be confusing to follow
if your functions and methods are ambiguously named.
:::

#### Fluent API

The <APILink fn="mlflow.start_span">`fluent APIs`</APILink> context handler implementation for initiating spans is useful when you need full control of the logging of each aspect of the span's data.

The example from our application for ensuring that we're capturing the parameters that are set when loading the model via the `load_context` method is
shown below. We are pulling from the instance attributes `self.models_config` and `self.models` to set the attributes of the span.

```python
with mlflow.start_span("Audit Agent") as root_span:
    root_span.set_inputs(messages)
    attributes = {**params.to_dict(), **self.models_config, **self.models}
    root_span.set_attributes(attributes)
    # More span manipulation...
```

#### Traces in the MLflow UI

After running our example that includes these combined usage patterns for trace span generation and instrumentation,

![Traces in the MLflow UI for the Agent example](/images/llms/chat-model-guide/agent-trace-ui.png)

In order to control the behavior of our `BasicAgent` model without having to hard-code configuration values directly into our model logic, specifying
configurations within the `model_config` parameter when logging the model gives some flexibility and versatility to our model definition.

This functionality allows us to:

- **Rapidly test** different configurations without having to make changes to source code
- **See the configuration** that was used when logging different iterations directly in the MLflow UI
- **Simplify the model code** by decoupling the configuration from the implementation

:::note
In our example model, we set a standard set of configurations that control the behavior of the ``BasicAgent``. The configuration
structure expected by the code is a dictionary with the following components:

- `models`: Defines the per-agent configurations.
- `(model_name)`: Represents the role of the agent. This section contains:
    - `endpoint`: The specific model type being used by the agent.
    - `instruction`: The prompt given to the model, describing its role and responsibilities.
    - `temperature`: The temperature setting controlling response variability.
    - `max_tokens`: The maximum token limit for generating responses.
- `configuration`: Contains miscellaneous settings for the agent application.
    - `user_response_instruction`: Provides context for the second agent by simulating a user response based on the first agent's output.
:::

This configuration structure definition will be:

- **Defined when logging the model** and structured to support the needs of the model's behavior
- **Used by the load_context method** and applied to the model when loading
- **Logged within the MLmodel file** and will be visible within the artifact viewer in the MLflow UI

The `model_config` values that are submitted for our `BasicAgent` example within this tutorial can be seen within the logged model's
`MLmodel` file in the UI:

<div className="center-div" style={{ width: "50%" }}>
![Model configuration in the MLflow UI](/images/llms/chat-model-guide/model-config-in-ui.png)
</div>

One of the more complex tasks associated with deploying an LLM application or AI agent with MLflow arises when attempting to build a custom implementation
that is based on subclassing the <APILink fn="mlflow.pyfunc.PythonModel">`mlflow.pyfunc.PythonModel`</APILink> abstraction.

While `PythonModel` is recommended for custom Deep Learning and traditional Machine Learning models (such as `sklearn` or `torch` models that require
additional processing logic apart from that of a base model), there are internal manipulations of the input data that occur
when serving these models that introduce unnecessary complications with LLM applications and AI agents.

Due to the fact that DL and traditional ML models largely rely on structured data, when input data is passed via a REST interface for model serving,
the `PythonModel` implementation will convert JSON data into `pandas.DataFrame` or `numpy` objects. This conversion creates a confusing and difficult to
debug scenario when using LLM and AI agent models. These implementations generally deal exclusively with JSON-conformant data structures and have no tabular
representation that makes intuitive sense, thereby creating a frustrating and complex conversion interface needed to make application deployment function
correctly.

To simplify this problem, the <APILink fn="mlflow.pyfunc.ChatModel">`mlflow.pyfunc.ChatModel`</APILink> class was created to provide a simpler interface for handling of the data
passed into and returned from a call to the `predict()` method on custom Python models serving LLM and AI agent use cases.

In the example tutorial code below, we subclass `ChatModel` in order to utilize this simplified interface with its immutable input and output
formats. Because of this immutability, we don't have to reason about model signatures, and can instead directly use API standards that have
been broadly accepted throughout the AI industry.

To illustrate why it is preferred to use `ChatModel` as a super class to custom LLM and AI agent implementations in MLflow, here is the signature that
would otherwise need to be defined and supplied during model logging to conform to the `OpenAI` API spec as of September 2024:

**Input Schema** as a `dict`:

```python
[
    {
        "type": "array",
        "items": {
            "type": "object",
            "properties": {
                "content": {"type": "string", "required": True},
                "name": {"type": "string", "required": False},
                "role": {"type": "string", "required": True},
            },
        },
        "name": "messages",
        "required": True,
    },
    {"type": "double", "name": "temperature", "required": False},
    {"type": "long", "name": "max_tokens", "required": False},
    {"type": "array", "items": {"type": "string"}, "name": "stop", "required": False},
    {"type": "long", "name": "n", "required": False},
    {"type": "boolean", "name": "stream", "required": False},
    {"type": "double", "name": "top_p", "required": False},
    {"type": "long", "name": "top_k", "required": False},
    {"type": "double", "name": "frequency_penalty", "required": False},
    {"type": "double", "name": "presence_penalty", "required": False},
]
```

:::note
Agent-based (tool-calling) schemas are significantly more complex than the simpler chat interface example shown above. As LLM and AI agent frameworks and services
evolve with increasingly sophisticated capabilities and features, the complexity of these interfaces will grow, making manual schema definitions a
challenging and time-consuming task. The structured input validation provided by the MLflow <APILink fn="mlflow.pyfunc.ChatModel">`mlflow.pyfunc.ChatModel`</APILink> interface removes the burden of defining and
managing these intricate signatures manually. By leveraging these pre-defined schemas, you gain robust input type safety and validation, ensuring your
deployed applications handle inputs consistently and correctly without additional effort. This approach not only reduces the risk of errors but also
streamlines the development process, allowing you to focus on building impactful AI solutions without the overhead of managing complex input specifications.
:::

By using <APILink fn="mlflow.pyfunc.ChatModel">`mlflow.pyfunc.ChatModel`</APILink> to base a custom implementation off of, we don't have to reason about this complex signature.
It is provided for us.

The only two considerations to be aware of when interfacing with the static signatures of `ChatModel` are:

- If the service that your custom implementation is interfacing with doesn't adhere to the `OpenAI` spec, you will need to extract data from the
  standard structure of <APILink fn="mlflow.types.llm.ChatMessage">`mlflow.types.llm.ChatMessage`</APILink> and <APILink fn="mlflow.types.llm.ChatParams">`mlflow.types.llm.ChatParams`</APILink> and ensure that it conforms to what
  your service is expecting.
- The returned response from `predict` should adhere to the output structure defined within the `ChatModel` output signature:
  <APILink fn="mlflow.types.llm.ChatCompletionResponse">`mlflow.types.llm.ChatCompletionResponse`</APILink>.

There are a number of ways that building a custom implementation for an LLM or AI agent use case can be frustrating or not intuitive. Here are some of the
most common that we've heard from our users:

#### Not using a supported flavor

If you're working with a library that is natively supported in MLflow, leveraging the built-in support for logging and loading your implementation
will always be easier than implementing a custom model. It is recommended to check the [supported flavors](/genai/flavors)
to see if there is a built-in solution that will meet your use case needs in one of the many integrations that are available.

#### Misinterpreting what `load_context` does

While subclassing one of the base model types for a custom model, it may appear that the class definition is a "what you see is what you get" standard
Python class. However, when loading your custom model instance, the `load_context` method is actually called by another loader object.

Because of the implementation, you **cannot have direct assignment of undefined instance attributes** within `load_context`.

For example, this does not work:

```python
from mlflow.pyfunc import ChatModel


class MyModel(ChatModel):
    def __init__(self):
        self.state = []

    def load_context(self, context):
        # This will fail on load as the instance attribute self.my_model_config is not defined
        self.my_model_config = context.get("my_model_config")
```

Instead, ensure that any instance attributes that are set by the `load_context` method are defined in the class constructor with a
placeholder value:

```python
from mlflow.pyfunc import ChatModel


class MyModel(ChatModel):
    def __init__(self):
        self.state = []
        self.my_model_config = None  # Define the attribute here

    def load_context(self, context):
        self.my_model_config = context.get("my_model_config")
```

#### Failing to Handle Secrets securely

It might be tempting to simplify your model's deployment by specifying authentication secrets within a configuration. However, any configuration
data that is defined within your `model_config` parameters **is directly visible in the MLflow UI** and is not stored securely.

The recommended approach for handling sensitive configuration data such as API keys or access tokens is to utilize a Secret Manager.
The configuration for **what to fetch** from your secrets management system can be stored within the `model_config` definition and
your deployment environment can utilize a secure means of accessing the key reference for your secrets management service.

An effective place to handle secrets assignment (generally set as environment variables or passed as a part of request headers) is to
handle the acquisition and per-session setting within `load_context`. If you have rotating tokens, it is worthwhile to embed the acquisition
of secrets and re-fetching of them upon expiry as part of a retry mechanism within the call stack of `predict`.

#### Failing to use `input_example`

While it may seem that providing an `input_example` when logging a model in MLflow is purely for cosmetic purposes within the artifact view
display within the MLflow UI, there is an additional bit of functionality that makes providing this data very useful, particularly for LLM and AI agent
use cases.

When an `input_example` is provided, MLflow will call your model's `predict` method with the example data to validate that the input is
compatible with the model object that you are logging. If there are any failures that occur, you will receive an error message detailing
what is wrong with the input syntax. This is very beneficial to ensure that, at the point of logging, you can ensure that your expected
input interface structure is what will be allowable for the deployed model, thereby saving you hours of debugging and troubleshooting later
when attempting to deploy your solution.

It is **highly recommended** to supply this example during logging.

#### Failing to handle retries for Rate Limits being hit

Nearly all LLM provider services impose rate limits and token-based usage limits to prevent disruption to their service or to help protect
users from unexpected bills. When limits are reached, it is important that your prediction logic is robust to handle these failures to ensure
that a user of your deployed application understands why their request was not successful.

It can be beneficial to introduce retry logic for certain errors, particularly those involving transient connection issues or per-unit-of-time
request limits.

#### Not validating before deployment

The process of deploying an LLM application or AI agent can take a significant amount of time. When an implementation is finally ready to be submitted to a
serving environment, the last thing that you want to deal with is a model that is incapable of being served due to some issue with a decoded
JSON payload being submitted to your model's `predict()` method.

MLflow offers the <APILink fn="mlflow.models.validate_serving_input" /> API to ensure that the model that you have logged is capable of being interacted
with by emulating the data processing that occurs with a deployed model.

To use this API, simply navigate to your logged model with the MLflow UI's artifact viewer. The model display pane on the right side of
the artifact viewer contains the code snippet that you can execute in an interactive environment to ensure that your model is ready to
deploy.

For the example in this tutorial, this is the generated code that is copied from the artifact viewer display:

```python
from mlflow.models import validate_serving_input

model_uri = "runs:/8935b7aff5a84f559b5fcc2af3e2ea31/model"

# The model is logged with an input example. MLflow converts
# it into the serving payload format for the deployed model endpoint,
# and saves it to 'serving_input_payload.json'
serving_payload = """{
"messages": [
    {
    "role": "user",
    "content": "What is a good recipe for baking scones that doesn't require a lot of skill?"
    }
],
"temperature": 1.0,
"n": 1,
"stream": false
}"""

# Validate the serving payload works on the model
validate_serving_input(model_uri, serving_payload)
```

import mlflow
from mlflow.types.llm import ChatCompletionResponse, ChatMessage, ChatParams, ChatChoice
from mlflow.pyfunc import ChatModel
from mlflow import deployments
from typing import List, Optional, Dict


class BasicAgent(ChatModel):
    def __init__(self):
        """Initialize the BasicAgent with placeholder values."""
        self.deploy_client = None
        self.models = {}
        self.models_config = {}
        self.conversation_history = []

    def load_context(self, context):
        """Initialize the connectors and model configurations."""
        self.deploy_client = deployments.get_deploy_client("databricks")
        self.models = context.model_config.get("models", {})
        self.models_config = context.model_config.get("configuration", {})

    def _get_system_message(self, role: str) -> Dict:
        """
        Get the system message configuration for the specified role.

        Args:
            role (str): The role of the agent (e.g., "oracle" or "judge").

        Returns:
            dict: The system message for the given role.
        """
        if role not in self.models:
            raise ValueError(f"Unknown role: {role}")

        instruction = self.models[role]["instruction"]
        return ChatMessage(role="system", content=instruction).to_dict()

    @mlflow.trace(name="Raw Agent Response")
    def _get_agent_response(
        self, message_list: List[Dict], endpoint: str, params: Optional[dict] = None
    ) -> Dict:
        """
        Call the agent endpoint to get a response.

        Args:
            message_list (List[Dict]): List of messages for the agent.
            endpoint (str): The agent's endpoint.
            params (Optional[dict]): Additional parameters for the call.

        Returns:
            dict: The response from the agent.
        """
        response = self.deploy_client.predict(
            endpoint=endpoint, inputs={"messages": message_list, **(params or {})}
        )
        return response["choices"][0]["message"]

    @mlflow.trace(name="Agent Call")
    def _call_agent(self, message: ChatMessage, role: str, params: Optional[dict] = None) -> Dict:
        """
        Prepares and sends the request to a specific agent based on the role.

        Args:
            message (ChatMessage): The message to be processed.
            role (str): The role of the agent (e.g., "oracle" or "judge").
            params (Optional[dict]): Additional parameters for the call.

        Returns:
            dict: The response from the agent.
        """
        system_message = self._get_system_message(role)
        message_list = self._prepare_message_list(system_message, message)

        # Fetch agent response
        agent_config = self.models[role]
        response = self._get_agent_response(message_list, agent_config["endpoint"], params)

        # Update conversation history
        self.conversation_history.extend([message.to_dict(), response])
        return response

    @mlflow.trace(name="Assemble Conversation")
    def _prepare_message_list(self, system_message: Dict, user_message: ChatMessage) -> List[Dict]:
        """
        Prepare the list of messages to send to the agent.

        Args:
            system_message (dict): The system message dictionary.
            user_message (ChatMessage): The user message.

        Returns:
            List[dict]: The complete list of messages to send.
        """
        user_prompt = {
            "role": "user",
            "content": self.models_config.get(
                "user_response_instruction", "Can you make the answer better?"
            ),
        }
        if self.conversation_history:
            return [system_message, *self.conversation_history, user_prompt]
        else:
            return [system_message, user_message.to_dict()]

    def predict(
        self, context, messages: List[ChatMessage], params: Optional[ChatParams] = None
    ) -> ChatCompletionResponse:
        """
        Predict method to handle agent conversation.

        Args:
            context: The MLflow context.
            messages (List[ChatMessage]): List of messages to process.
            params (Optional[ChatParams]): Additional parameters for the conversation.

        Returns:
            ChatCompletionResponse: The structured response object.
        """
        # Use the fluent API context handler to have added control over what is included in the span
        with mlflow.start_span(name="Audit Agent") as root_span:
            # Add the user input to the root span
            root_span.set_inputs(messages)

            # Add attributes to the root span
            attributes = {**params.to_dict(), **self.models_config, **self.models}
            root_span.set_attributes(attributes)

            # Initiate the conversation with the oracle
            oracle_params = self._get_model_params("oracle")
            oracle_response = self._call_agent(messages[0], "oracle", oracle_params)

            # Process the response with the judge
            judge_params = self._get_model_params("judge")
            judge_response = self._call_agent(ChatMessage(**oracle_response), "judge", judge_params)

            # Reset the conversation history and return the final response
            self.conversation_history = []

            output = ChatCompletionResponse(
                choices=[ChatChoice(index=0, message=ChatMessage(**judge_response))],
                usage={},
                model=judge_params.get("endpoint", "unknown"),
            )

            root_span.set_outputs(output)

        return output

    def _get_model_params(self, role: str) -> dict:
        """
        Retrieves model parameters for a given role.

        Args:
            role (str): The role of the agent (e.g., "oracle" or "judge").

        Returns:
            dict: A dictionary of parameters for the agent.
        """
        role_config = self.models.get(role, {})

        return {
            "temperature": role_config.get("temperature", 0.5),
            "max_tokens": role_config.get("max_tokens", 500),
        }


# IMPORTANT: specifies the Python ChatModel instance to use for inference requests when
# the model is loaded back
agent = BasicAgent()
mlflow.models.set_model(agent)

model_config = {
    "models": {
        "judge": {
            "endpoint": "databricks-meta-llama-3-1-405b-instruct",
            "instruction": (
                "You are an evaluator of answers provided by others. Based on the context of both the question and the answer, "
                "provide a corrected answer if it is incorrect; otherwise, enhance the answer with additional context and explanation."
            ),
            "temperature": 0.5,
            "max_tokens": 2000,
        },
        "oracle": {
            "endpoint": "databricks-mixtral-8x7b-instruct",
            "instruction": (
                "You are a knowledgeable source of information that excels at providing detailed, but brief answers to questions. "
                "Provide an answer to the question based on the information provided."
            ),
            "temperature": 0.9,
            "max_tokens": 5000,
        },
    },
    "configuration": {
        "user_response_instruction": "Can you evaluate and enhance this answer with the provided contextual history?"
    },
}

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is a good recipe for baking scones that doesn't require a lot of skill?",
        }
    ]
}

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="model",
        # If needed, update `python_model` to the Python file containing your agent code
        python_model="basic_agent.py",
        model_config=model_config,
        input_example=input_example,
    )

loaded = mlflow.pyfunc.load_model(model_info.model_uri)

response = loaded.predict({
    "messages": [
        {
            "role": "user",
            "content": "What is the best material to make a baseball bat out of?",
        }
    ]
})