663 lines
24 KiB
Python
663 lines
24 KiB
Python
import ast
|
|
import logging
|
|
from contextvars import ContextVar
|
|
from typing import Any, Optional, Sequence, Union
|
|
from uuid import UUID
|
|
|
|
import pydantic
|
|
from langchain_core.agents import AgentAction, AgentFinish
|
|
from langchain_core.callbacks.base import BaseCallbackHandler
|
|
from langchain_core.documents import Document
|
|
from langchain_core.load.dump import dumps
|
|
from langchain_core.messages import BaseMessage
|
|
from langchain_core.outputs import (
|
|
ChatGenerationChunk,
|
|
GenerationChunk,
|
|
LLMResult,
|
|
)
|
|
from tenacity import RetryCallState
|
|
|
|
import mlflow
|
|
from mlflow import MlflowClient
|
|
from mlflow.entities import Document as MlflowDocument
|
|
from mlflow.entities import LiveSpan, SpanEvent, SpanStatus, SpanStatusCode, SpanType
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.langchain.utils.chat import (
|
|
convert_lc_generation_to_chat_message,
|
|
convert_lc_message_to_chat_message,
|
|
)
|
|
from mlflow.pyfunc.context import Context, maybe_set_prediction_context
|
|
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.tracing.utils.token import SpanWithToken
|
|
from mlflow.types.chat import ChatMessage, ChatTool, FunctionToolDefinition
|
|
from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
|
|
from mlflow.utils.autologging_utils import ExceptionSafeAbstractClass
|
|
from mlflow.utils.autologging_utils.config import AutoLoggingConfig
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
_should_attach_span_to_context = ContextVar("should_attach_span_to_context", default=True)
|
|
|
|
|
|
def patched_callback_manager_init(original, self, *args, **kwargs):
|
|
original(self, *args, **kwargs)
|
|
|
|
if not AutoLoggingConfig.init(mlflow.langchain.FLAVOR_NAME).log_traces:
|
|
return
|
|
|
|
for handler in self.inheritable_handlers:
|
|
if isinstance(handler, MlflowLangchainTracer):
|
|
return
|
|
|
|
_handler = MlflowLangchainTracer()
|
|
self.add_handler(_handler, inherit=True)
|
|
|
|
|
|
def patched_callback_manager_merge(original, self, *args, **kwargs):
|
|
"""
|
|
Patch BaseCallbackManager.merge to avoid a duplicated callback issue.
|
|
|
|
In the above patched __init__, we check `inheritable_handlers` to see if the MLflow tracer
|
|
is already propagated. This works when the `inheritable_handlers` is specified as constructor
|
|
arguments. However, in the `merge` method, LangChain does not use constructor but set
|
|
callbacks via the setter method. This causes duplicated callbacks injection.
|
|
https://github.com/langchain-ai/langchain/blob/d9a069c414a321e7a3f3638a32ecf8a37ec2d188/libs/core/langchain_core/callbacks/base.py#L962-L982
|
|
"""
|
|
# Get the MLflow callback inherited from parent
|
|
inherited = self.inheritable_handlers + args[0].inheritable_handlers
|
|
inherited_mlflow_cb = next(
|
|
(cb for cb in inherited if isinstance(cb, MlflowLangchainTracer)), None
|
|
)
|
|
|
|
if not inherited_mlflow_cb:
|
|
return original(self, *args, **kwargs)
|
|
|
|
merged = original(self, *args, **kwargs)
|
|
# If a new MLflow callback is generated inside __init__, remove it
|
|
duplicate_mlflow_cbs = [
|
|
cb
|
|
for cb in merged.inheritable_handlers
|
|
if isinstance(cb, MlflowLangchainTracer) and cb != inherited_mlflow_cb
|
|
]
|
|
for cb in duplicate_mlflow_cbs:
|
|
merged.remove_handler(cb)
|
|
|
|
return merged
|
|
|
|
|
|
def patched_runnable_sequence_batch(original, self, *args, **kwargs):
|
|
"""
|
|
Patch to terminate span context attachment during batch execution.
|
|
|
|
RunnableSequence's batch() methods are implemented in a peculiar way
|
|
that iterates on steps->items sequentially within the same thread. For example, if a
|
|
sequence has 2 steps and the batch size is 3, the execution flow will be:
|
|
- Step 1 for item 1
|
|
- Step 1 for item 2
|
|
- Step 1 for item 3
|
|
- Step 2 for item 1
|
|
- Step 2 for item 2
|
|
- Step 2 for item 3
|
|
Due to this behavior, we cannot attach the span to the context for this particular
|
|
API, otherwise spans for different inputs will be mixed up.
|
|
"""
|
|
original_state = _should_attach_span_to_context.get()
|
|
_should_attach_span_to_context.set(False)
|
|
try:
|
|
return original(self, *args, **kwargs)
|
|
finally:
|
|
_should_attach_span_to_context.set(original_state)
|
|
|
|
|
|
class MlflowLangchainTracer(BaseCallbackHandler, metaclass=ExceptionSafeAbstractClass):
|
|
"""
|
|
Callback for auto-logging traces.
|
|
We need to inherit ExceptionSafeAbstractClass to avoid invalid new
|
|
input arguments added to original function call.
|
|
|
|
Args:
|
|
prediction_context: Optional prediction context object to be set for the
|
|
thread-local context. Occasionally this has to be passed manually because
|
|
the callback may be invoked asynchronously and Langchain doesn't correctly
|
|
propagate the thread-local context.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
prediction_context: Optional[Context] = None,
|
|
):
|
|
# NB: The tracer can handle multiple traces in parallel under multi-threading scenarios.
|
|
# DO NOT use instance variables to manage the state of single trace.
|
|
super().__init__()
|
|
self._mlflow_client = MlflowClient()
|
|
# run_id: (LiveSpan, OTel token)
|
|
self._run_span_mapping: dict[str, SpanWithToken] = {}
|
|
self._prediction_context = prediction_context
|
|
|
|
def _get_span_by_run_id(self, run_id: UUID) -> Optional[LiveSpan]:
|
|
if span_with_token := self._run_span_mapping.get(str(run_id), None):
|
|
return span_with_token.span
|
|
raise MlflowException(f"Span for run_id {run_id!s} not found.")
|
|
|
|
def _serialize_invocation_params(
|
|
self, attributes: Optional[dict[str, Any]]
|
|
) -> Optional[dict[str, Any]]:
|
|
"""
|
|
Serialize the 'invocation_params' in the attributes dictionary.
|
|
If 'invocation_params' contains a key 'response_format' whose value is a subclass
|
|
of pydantic.BaseModel, replace it with its JSON schema.
|
|
"""
|
|
if not attributes:
|
|
return attributes
|
|
|
|
invocation_params = attributes.get("invocation_params")
|
|
if not isinstance(invocation_params, dict):
|
|
return attributes
|
|
|
|
response_format = invocation_params.get("response_format")
|
|
if isinstance(response_format, type) and issubclass(response_format, pydantic.BaseModel):
|
|
try:
|
|
invocation_params["response_format"] = (
|
|
response_format.model_json_schema()
|
|
if IS_PYDANTIC_V2_OR_NEWER
|
|
else response_format.schema()
|
|
)
|
|
except Exception as e:
|
|
_logger.error(
|
|
"Failed to generate JSON schema for response_format: %s", e, exc_info=True
|
|
)
|
|
return attributes
|
|
|
|
def _start_span(
|
|
self,
|
|
span_name: str,
|
|
parent_run_id: Optional[UUID],
|
|
span_type: str,
|
|
run_id: UUID,
|
|
inputs: Optional[Union[str, dict[str, Any]]] = None,
|
|
attributes: Optional[dict[str, Any]] = None,
|
|
) -> LiveSpan:
|
|
"""Start MLflow Span (or Trace if it is root component)"""
|
|
serialized_attributes = self._serialize_invocation_params(attributes)
|
|
with maybe_set_prediction_context(self._prediction_context):
|
|
parent = self._get_parent_span(parent_run_id)
|
|
if parent:
|
|
span = self._mlflow_client.start_span(
|
|
name=span_name,
|
|
request_id=parent.request_id,
|
|
parent_id=parent.span_id,
|
|
span_type=span_type,
|
|
inputs=inputs,
|
|
attributes=serialized_attributes,
|
|
)
|
|
else:
|
|
# When parent_run_id is None, this is root component so start trace
|
|
dependencies_schemas = (
|
|
self._prediction_context.dependencies_schemas
|
|
if self._prediction_context
|
|
else None
|
|
)
|
|
span = self._mlflow_client.start_trace(
|
|
name=span_name,
|
|
span_type=span_type,
|
|
inputs=inputs,
|
|
attributes=serialized_attributes,
|
|
tags=dependencies_schemas,
|
|
)
|
|
|
|
# Attach the span to the current context to mark it "active"
|
|
token = set_span_in_context(span) if _should_attach_span_to_context.get() else None
|
|
self._run_span_mapping[str(run_id)] = SpanWithToken(span, token)
|
|
return span
|
|
|
|
def _get_parent_span(self, parent_run_id) -> Optional[LiveSpan]:
|
|
"""
|
|
Get parent span from multiple sources:
|
|
1. If there is an active span in current context, use it as parent span
|
|
2. If parent_run_id is provided, get the corresponding span from the run -> span mapping
|
|
3. If none of the above, return None
|
|
"""
|
|
if active_span := mlflow.get_current_active_span():
|
|
return active_span
|
|
elif parent_run_id:
|
|
return self._get_span_by_run_id(parent_run_id)
|
|
return None
|
|
|
|
def _end_span(
|
|
self,
|
|
run_id: UUID,
|
|
span: LiveSpan,
|
|
outputs=None,
|
|
attributes=None,
|
|
status=SpanStatus(SpanStatusCode.OK),
|
|
):
|
|
"""Close MLflow Span (or Trace if it is root component)"""
|
|
try:
|
|
with maybe_set_prediction_context(self._prediction_context):
|
|
self._mlflow_client.end_span(
|
|
request_id=span.request_id,
|
|
span_id=span.span_id,
|
|
outputs=outputs,
|
|
attributes=attributes,
|
|
status=status,
|
|
)
|
|
finally:
|
|
# Span should be detached from the context even when the client.end_span fails
|
|
st = self._run_span_mapping.pop(str(run_id), None)
|
|
if _should_attach_span_to_context.get():
|
|
if st.token is None:
|
|
raise MlflowException(
|
|
f"Token for span {st.span} is not found. "
|
|
"Cannot detach the span from context."
|
|
)
|
|
detach_span_from_context(st.token)
|
|
|
|
def flush(self):
|
|
"""Flush the state of the tracer."""
|
|
# Ideally, all spans should be popped and ended. However, LangChain sometimes
|
|
# does not trigger the end event properly and some spans may be left open.
|
|
# To avoid leaking tracing context, we remove all spans from the mapping.
|
|
for st in self._run_span_mapping.values():
|
|
if st.token:
|
|
_logger.debug(f"Found leaked span {st.span}. Force ending it.")
|
|
detach_span_from_context(st.token)
|
|
|
|
self._run_span_mapping = {}
|
|
|
|
def _assign_span_name(self, serialized: dict[str, Any], default_name="unknown") -> str:
|
|
return serialized.get("name", serialized.get("id", [default_name])[-1])
|
|
|
|
def on_chat_model_start(
|
|
self,
|
|
serialized: dict[str, Any],
|
|
messages: list[list[BaseMessage]],
|
|
*,
|
|
run_id: UUID,
|
|
tags: Optional[list[str]] = None,
|
|
parent_run_id: Optional[UUID] = None,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run when a chat model starts running."""
|
|
|
|
if metadata:
|
|
kwargs.update({"metadata": metadata})
|
|
|
|
span = self._start_span(
|
|
span_name=name or self._assign_span_name(serialized, "chat model"),
|
|
parent_run_id=parent_run_id,
|
|
span_type=SpanType.CHAT_MODEL,
|
|
run_id=run_id,
|
|
inputs=messages,
|
|
attributes=kwargs,
|
|
)
|
|
|
|
mlflow_messages = [
|
|
convert_lc_message_to_chat_message(msg)
|
|
for message_list in messages
|
|
for msg in message_list
|
|
]
|
|
set_span_chat_messages(span, mlflow_messages)
|
|
|
|
if tools := self._extract_tool_definitions(kwargs):
|
|
set_span_chat_tools(span, tools)
|
|
|
|
def on_llm_start(
|
|
self,
|
|
serialized: dict[str, Any],
|
|
prompts: list[str],
|
|
*,
|
|
run_id: UUID,
|
|
tags: Optional[list[str]] = None,
|
|
parent_run_id: Optional[UUID] = None,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Run when LLM (non-chat models) starts running."""
|
|
if metadata:
|
|
kwargs.update({"metadata": metadata})
|
|
|
|
span = self._start_span(
|
|
span_name=name or self._assign_span_name(serialized, "llm"),
|
|
parent_run_id=parent_run_id,
|
|
span_type=SpanType.LLM,
|
|
run_id=run_id,
|
|
inputs=prompts,
|
|
attributes=kwargs,
|
|
)
|
|
|
|
mlflow_messages = [ChatMessage(role="user", content=prompt) for prompt in prompts]
|
|
set_span_chat_messages(span, mlflow_messages)
|
|
|
|
if tools := self._extract_tool_definitions(kwargs):
|
|
set_span_chat_tools(span, tools)
|
|
|
|
def _extract_tool_definitions(self, kwargs: dict[str, Any]) -> list[ChatTool]:
|
|
raw_tools = kwargs.get("invocation_params", {}).get("tools", [])
|
|
tools = []
|
|
for raw_tool in raw_tools:
|
|
# First, try to parse the raw tool dictionary as OpenAI-style tool
|
|
try:
|
|
tool = ChatTool.validate_compat(raw_tool)
|
|
tools.append(tool)
|
|
except pydantic.ValidationError:
|
|
# If not OpenAI style, just try to extract the name and descriptions.
|
|
if name := raw_tool.get("name"):
|
|
tool = ChatTool(
|
|
type="function",
|
|
function=FunctionToolDefinition(
|
|
name=name, description=raw_tool.get("description")
|
|
),
|
|
)
|
|
tools.append(tool)
|
|
else:
|
|
_logger.warning(f"Failed to parse tool definition for tracing: {raw_tool}.")
|
|
|
|
return tools
|
|
|
|
def on_llm_new_token(
|
|
self,
|
|
token: str,
|
|
*,
|
|
chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
|
|
run_id: UUID,
|
|
parent_run_id: Optional[UUID] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run on new LLM token. Only available when streaming is enabled."""
|
|
llm_span = self._get_span_by_run_id(run_id)
|
|
event_kwargs = {"token": token}
|
|
if chunk:
|
|
event_kwargs["chunk"] = dumps(chunk)
|
|
llm_span.add_event(
|
|
SpanEvent(
|
|
name="new_token",
|
|
attributes=event_kwargs,
|
|
)
|
|
)
|
|
|
|
def on_retry(
|
|
self,
|
|
retry_state: RetryCallState,
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run on a retry event."""
|
|
span = self._get_span_by_run_id(run_id)
|
|
retry_d: dict[str, Any] = {
|
|
"slept": retry_state.idle_for,
|
|
"attempt": retry_state.attempt_number,
|
|
}
|
|
if retry_state.outcome is None:
|
|
retry_d["outcome"] = "N/A"
|
|
elif retry_state.outcome.failed:
|
|
retry_d["outcome"] = "failed"
|
|
exception = retry_state.outcome.exception()
|
|
retry_d["exception"] = str(exception)
|
|
retry_d["exception_type"] = exception.__class__.__name__
|
|
else:
|
|
retry_d["outcome"] = "success"
|
|
retry_d["result"] = str(retry_state.outcome.result())
|
|
span.add_event(
|
|
SpanEvent(
|
|
name="retry",
|
|
attributes=retry_d,
|
|
)
|
|
)
|
|
|
|
def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any):
|
|
"""End the span for an LLM run."""
|
|
llm_span = self._get_span_by_run_id(run_id)
|
|
|
|
# Record the chat messages attribute
|
|
input_messages = llm_span.get_attribute(SpanAttributeKey.CHAT_MESSAGES) or []
|
|
output_messages = [
|
|
convert_lc_generation_to_chat_message(gen)
|
|
for gen_list in response.generations
|
|
for gen in gen_list
|
|
]
|
|
set_span_chat_messages(llm_span, input_messages + output_messages)
|
|
self._end_span(run_id, llm_span, outputs=response)
|
|
|
|
def on_llm_error(
|
|
self,
|
|
error: BaseException,
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
):
|
|
"""Handle an error for an LLM run."""
|
|
llm_span = self._get_span_by_run_id(run_id)
|
|
llm_span.add_event(SpanEvent.from_exception(error))
|
|
self._end_span(run_id, llm_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
|
|
|
|
def on_chain_start(
|
|
self,
|
|
serialized: dict[str, Any],
|
|
inputs: Union[dict[str, Any], Any],
|
|
*,
|
|
run_id: UUID,
|
|
tags: Optional[list[str]] = None,
|
|
parent_run_id: Optional[UUID] = None,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
run_type: Optional[str] = None,
|
|
name: Optional[str] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Start span for a chain run."""
|
|
if metadata:
|
|
kwargs.update({"metadata": metadata})
|
|
# not considering streaming events for now
|
|
self._start_span(
|
|
span_name=name or self._assign_span_name(serialized, "chain"),
|
|
parent_run_id=parent_run_id,
|
|
span_type=SpanType.CHAIN,
|
|
run_id=run_id,
|
|
inputs=inputs,
|
|
attributes=kwargs,
|
|
)
|
|
|
|
def on_chain_end(
|
|
self,
|
|
outputs: dict[str, Any],
|
|
*,
|
|
run_id: UUID,
|
|
inputs: Optional[Union[dict[str, Any], Any]] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run when chain ends running."""
|
|
chain_span = self._get_span_by_run_id(run_id)
|
|
if inputs:
|
|
chain_span.set_inputs(inputs)
|
|
self._end_span(run_id, chain_span, outputs=outputs)
|
|
|
|
def on_chain_error(
|
|
self,
|
|
error: BaseException,
|
|
*,
|
|
inputs: Optional[Union[dict[str, Any], Any]] = None,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run when chain errors."""
|
|
chain_span = self._get_span_by_run_id(run_id)
|
|
if inputs:
|
|
chain_span.set_inputs(inputs)
|
|
chain_span.add_event(SpanEvent.from_exception(error))
|
|
self._end_span(run_id, chain_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
|
|
|
|
def on_tool_start(
|
|
self,
|
|
serialized: dict[str, Any],
|
|
input_str: str,
|
|
*,
|
|
run_id: UUID,
|
|
tags: Optional[list[str]] = None,
|
|
parent_run_id: Optional[UUID] = None,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
# We don't use inputs here because LangChain override the original inputs
|
|
# with None for some cases. In order to avoid losing the original inputs,
|
|
# we try to parse the input_str instead.
|
|
# https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/tools/base.py#L636-L640
|
|
inputs: Optional[dict[str, Any]] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Start span for a tool run."""
|
|
if metadata:
|
|
kwargs.update({"metadata": metadata})
|
|
|
|
# For function calling, input_str can be a stringified dictionary
|
|
# like "{'key': 'value'}". We try parsing it for better rendering,
|
|
# but conservatively fallback to original if it fails.
|
|
try:
|
|
inputs = ast.literal_eval(input_str)
|
|
except Exception:
|
|
inputs = input_str
|
|
|
|
self._start_span(
|
|
span_name=name or self._assign_span_name(serialized, "tool"),
|
|
parent_run_id=parent_run_id,
|
|
span_type=SpanType.TOOL,
|
|
run_id=run_id,
|
|
inputs=inputs,
|
|
attributes=kwargs,
|
|
)
|
|
|
|
def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any):
|
|
"""Run when tool ends running."""
|
|
tool_span = self._get_span_by_run_id(run_id)
|
|
self._end_span(run_id, tool_span, outputs=output)
|
|
|
|
def on_tool_error(
|
|
self,
|
|
error: BaseException,
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run when tool errors."""
|
|
tool_span = self._get_span_by_run_id(run_id)
|
|
tool_span.add_event(SpanEvent.from_exception(error))
|
|
self._end_span(run_id, tool_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
|
|
|
|
def on_retriever_start(
|
|
self,
|
|
serialized: dict[str, Any],
|
|
query: str,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: Optional[UUID] = None,
|
|
tags: Optional[list[str]] = None,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run when Retriever starts running."""
|
|
if metadata:
|
|
kwargs.update({"metadata": metadata})
|
|
self._start_span(
|
|
span_name=name or self._assign_span_name(serialized, "retriever"),
|
|
parent_run_id=parent_run_id,
|
|
span_type=SpanType.RETRIEVER,
|
|
run_id=run_id,
|
|
inputs=query,
|
|
attributes=kwargs,
|
|
)
|
|
|
|
def on_retriever_end(self, documents: Sequence[Document], *, run_id: UUID, **kwargs: Any):
|
|
"""Run when Retriever ends running."""
|
|
retriever_span = self._get_span_by_run_id(run_id)
|
|
try:
|
|
# attempt to convert documents to MlflowDocument
|
|
documents = [MlflowDocument.from_langchain_document(doc) for doc in documents]
|
|
except Exception as e:
|
|
_logger.debug(
|
|
f"Failed to convert LangChain Document to MLflow Document: {e}",
|
|
exc_info=True,
|
|
)
|
|
self._end_span(
|
|
run_id,
|
|
retriever_span,
|
|
outputs=documents,
|
|
)
|
|
|
|
def on_retriever_error(
|
|
self,
|
|
error: BaseException,
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
):
|
|
"""Run when Retriever errors."""
|
|
retriever_span = self._get_span_by_run_id(run_id)
|
|
retriever_span.add_event(SpanEvent.from_exception(error))
|
|
self._end_span(run_id, retriever_span, status=SpanStatus(SpanStatusCode.ERROR, str(error)))
|
|
|
|
def on_agent_action(
|
|
self,
|
|
action: AgentAction,
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""
|
|
Run on agent action.
|
|
|
|
NB: Agent action doesn't create a new LangChain Run, so instead of creating a new span,
|
|
an action will be recorded as an event of the existing span created by a parent chain.
|
|
"""
|
|
span = self._get_span_by_run_id(run_id)
|
|
span.add_event(
|
|
SpanEvent(
|
|
name="agent_action",
|
|
attributes={
|
|
"tool": action.tool,
|
|
"tool_input": dumps(action.tool_input),
|
|
"log": action.log,
|
|
},
|
|
)
|
|
)
|
|
|
|
def on_agent_finish(
|
|
self,
|
|
finish: AgentFinish,
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Run on agent end."""
|
|
span = self._get_span_by_run_id(run_id)
|
|
span.add_event(
|
|
SpanEvent(
|
|
name="agent_finish",
|
|
attributes={"return_values": dumps(finish.return_values), "log": finish.log},
|
|
)
|
|
)
|
|
|
|
def on_text(
|
|
self,
|
|
text: str,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: Optional[UUID] = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Run on arbitrary text."""
|
|
try:
|
|
span = self._get_span_by_run_id(run_id)
|
|
except MlflowException:
|
|
_logger.warning("Span not found for text event. Skipping text event logging.")
|
|
else:
|
|
span.add_event(
|
|
SpanEvent(
|
|
"text",
|
|
attributes={"text": text},
|
|
)
|
|
)
|