追踪多轮对话会话

在对话式AI应用中，用户通常在同一个对话会话中与模型进行多次交互。由于在典型的MLflow设置中，每次交互都会生成一个追踪记录，因此将这些追踪记录分组以分析整个对话非常有用。您可以通过将会话ID作为标签附加到每个追踪记录来实现这一点。

以下示例展示了如何在使用了mlflow.pyfunc.ChatModel()类实现的聊天模型中使用会话ID。有关如何在追踪记录上设置标签的更多信息，请参阅操作指南中的设置追踪标签部分。

import mlflow
from mlflow.entities import SpanType
from mlflow.types.llm import ChatMessage, ChatParams, ChatCompletionResponse

import openai
from typing import Optional

mlflow.set_experiment("Tracing Session ID Demo")


class ChatModelWithSession(mlflow.pyfunc.ChatModel):
    @mlflow.trace(span_type=SpanType.CHAT_MODEL)
    def predict(
        self, context, messages: list[ChatMessage], params: Optional[ChatParams] = None
    ) -> ChatCompletionResponse:
        if session_id := (params.custom_inputs or {}).get("session_id"):
            # Set session ID tag on the current trace
            mlflow.update_current_trace(tags={"session_id": session_id})

        response = openai.OpenAI().chat.completions.create(
            messages=[m.to_dict() for m in messages],
            model="gpt-4o-mini",
        )

        return ChatCompletionResponse.from_dict(response.to_dict())


model = ChatModelWithSession()

# Invoke the chat model multiple times with the same session ID
session_id = "123"
messages = [ChatMessage(role="user", content="What is MLflow Tracing?")]
response = model.predict(
    None, messages, ChatParams(custom_inputs={"session_id": session_id})
)

# Invoke again with the same session ID
messages.append(
    ChatMessage(role="assistant", content=response.choices[0].message.content)
)
messages.append(ChatMessage(role="user", content="How to get started?"))
response = model.predict(
    None, messages, ChatParams(custom_inputs={"session_id": session_id})
)

上述代码创建了两个具有相同会话ID标签的新追踪记录。在MLflow UI中，您可以使用tag.session_id = '123'搜索具有此定义会话ID的追踪记录。

另外，您可以使用mlflow.search_traces()函数以编程方式获取这些追踪记录。有关更多详细信息，请参阅搜索和检索追踪记录。

traces = mlflow.search_traces(filter_string="tag.session_id = '123456'")