wehub-resource-sync
64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
from semantic_kernel.kernel import Kernel
|
|
from semantic_kernel.utils.telemetry.model_diagnostics import decorators
|
|
|
|
from mlflow.semantic_kernel.autolog import setup_semantic_kernel_tracing
|
|
from mlflow.semantic_kernel.tracing_utils import (
|
|
patched_kernel_entry_point,
|
|
semantic_kernel_diagnostics_wrapper,
|
|
)
|
|
from mlflow.telemetry.events import AutologgingEvent
|
|
from mlflow.telemetry.track import _record_event
|
|
from mlflow.utils.annotations import experimental as experimental
|
|
from mlflow.utils.autologging_utils import autologging_integration, safe_patch
|
|
|
|
FLAVOR_NAME = "semantic_kernel"
|
|
|
|
|
|
@autologging_integration(FLAVOR_NAME)
|
|
def autolog(
|
|
log_traces: bool = True,
|
|
disable: bool = False,
|
|
silent: bool = False,
|
|
):
|
|
"""
|
|
Enables (or disables) and configures autologging from Semantic Kernel to MLflow.
|
|
Only synchronous calls are supported. Asynchnorous APIs and streaming are not recorded.
|
|
|
|
Args:
|
|
log_traces: If ``True``, traces are logged for Semantic Kernel.
|
|
If ``False``, no traces are collected during inference. Default to ``True``.
|
|
disable: If ``True``, disables the Semantic Kernel autologging. Default to ``False``.
|
|
silent: If ``True``, suppress all event logs and warnings from MLflow during Anthropic
|
|
autologging. If ``False``, show all events and warnings.
|
|
"""
|
|
|
|
setup_semantic_kernel_tracing()
|
|
|
|
# Create root spans for the kernel entry points.
|
|
for method in ["invoke", "invoke_prompt"]:
|
|
safe_patch(
|
|
FLAVOR_NAME,
|
|
Kernel,
|
|
method,
|
|
patched_kernel_entry_point,
|
|
)
|
|
|
|
# NB: Semantic Kernel uses logging to serialize inputs/outputs. These parsers are used by their
|
|
# tracing decorators to log the inputs/outputs. These patches give coverage for many additional
|
|
# methods that are not covered by above entry point patches.
|
|
for method_name in [
|
|
"_set_completion_input",
|
|
"_set_completion_response",
|
|
"_set_completion_error",
|
|
]:
|
|
safe_patch(
|
|
FLAVOR_NAME,
|
|
decorators,
|
|
method_name,
|
|
semantic_kernel_diagnostics_wrapper,
|
|
)
|
|
|
|
_record_event(
|
|
AutologgingEvent, {"flavor": FLAVOR_NAME, "log_traces": log_traces, "disable": disable}
|
|
)
|