44 lines
1.8 KiB
Python
44 lines
1.8 KiB
Python
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)
|