wehub-resource-sync
523 lines
18 KiB
Python
523 lines
18 KiB
Python
import asyncio
|
|
from unittest import mock
|
|
|
|
import openai
|
|
import pytest
|
|
import pytest_asyncio
|
|
from semantic_kernel import Kernel
|
|
from semantic_kernel.agents import AgentResponseItem
|
|
from semantic_kernel.connectors.ai.open_ai import OpenAIChatCompletion
|
|
from semantic_kernel.contents import ChatMessageContent
|
|
from semantic_kernel.exceptions import FunctionExecutionException, KernelInvokeException
|
|
from semantic_kernel.functions.function_result import FunctionResult
|
|
from semantic_kernel.utils.telemetry.agent_diagnostics import (
|
|
gen_ai_attributes as agent_gen_ai_attributes,
|
|
)
|
|
from semantic_kernel.utils.telemetry.model_diagnostics import (
|
|
gen_ai_attributes as model_gen_ai_attributes,
|
|
)
|
|
|
|
import mlflow.semantic_kernel
|
|
from mlflow.entities import SpanType
|
|
from mlflow.entities.span_status import SpanStatusCode
|
|
from mlflow.environment_variables import MLFLOW_USE_DEFAULT_TRACER_PROVIDER
|
|
from mlflow.semantic_kernel.autolog import SemanticKernelSpanProcessor
|
|
from mlflow.tracing.constant import (
|
|
SpanAttributeKey,
|
|
TokenUsageKey,
|
|
)
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
from tests.semantic_kernel.resources import (
|
|
_create_and_invoke_chat_agent,
|
|
_create_and_invoke_chat_completion_direct,
|
|
_create_and_invoke_embeddings,
|
|
_create_and_invoke_kernel_complex,
|
|
_create_and_invoke_kernel_function_object,
|
|
_create_and_invoke_kernel_simple,
|
|
_create_and_invoke_text_completion,
|
|
)
|
|
from tests.tracing.helper import get_traces
|
|
|
|
lock = asyncio.Lock()
|
|
|
|
|
|
@pytest_asyncio.fixture(autouse=True)
|
|
async def lock_fixture():
|
|
async with lock:
|
|
yield
|
|
|
|
|
|
@pytest.fixture(params=[True, False])
|
|
def with_openai_autolog(request):
|
|
# Test with OpenAI autologging enabled and disabled
|
|
if request.param:
|
|
mlflow.openai.autolog()
|
|
else:
|
|
mlflow.openai.autolog(disable=True)
|
|
|
|
return request.param
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_invoke_simple(mock_openai, with_openai_autolog, mock_litellm_cost):
|
|
mlflow.semantic_kernel.autolog()
|
|
result = await _create_and_invoke_kernel_simple(mock_openai)
|
|
|
|
# The mock OpenAI endpoint echos the user message back
|
|
prompt = "Is sushi the best food ever?"
|
|
expected_content = '[{"role": "user", "content": "Is sushi the best food ever?"}]'
|
|
|
|
# Validate the result is not mutated by tracing logic
|
|
assert isinstance(result, FunctionResult)
|
|
assert isinstance(result.value[0], ChatMessageContent)
|
|
assert result.value[0].items[0].text == expected_content
|
|
|
|
# Trace
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
assert trace.info.request_id
|
|
assert trace.info.experiment_id == "0"
|
|
assert trace.info.timestamp_ms > 0
|
|
assert trace.info.status == "OK"
|
|
assert "Is sushi the best food ever?" in trace.info.request_preview
|
|
assert "Is sushi the best food ever?" in trace.info.response_preview
|
|
|
|
spans = trace.data.spans
|
|
assert len(spans) == (5 if with_openai_autolog else 4)
|
|
|
|
# Kernel.invoke_prompt
|
|
assert spans[0].name == "Kernel.invoke_prompt"
|
|
assert spans[0].span_type == SpanType.AGENT
|
|
assert spans[0].inputs == {"prompt": prompt}
|
|
assert spans[0].outputs == [{"role": "assistant", "content": expected_content}]
|
|
|
|
# Kernel.invoke_prompt
|
|
assert spans[1].name == "Kernel.invoke"
|
|
assert spans[1].span_type == SpanType.AGENT
|
|
assert spans[1].inputs["function"] is not None
|
|
assert spans[1].outputs == [{"role": "assistant", "content": expected_content}]
|
|
|
|
# Execute LLM as a tool
|
|
assert spans[2].name.startswith("execute_tool")
|
|
assert spans[2].span_type == SpanType.TOOL
|
|
|
|
# Actual LLM call
|
|
assert spans[3].name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini")
|
|
assert "gen_ai.operation.name" in spans[3].attributes
|
|
assert spans[3].inputs == {"messages": [{"role": "user", "content": prompt}]}
|
|
assert spans[3].outputs == {"messages": [{"role": "assistant", "content": expected_content}]}
|
|
|
|
chat_usage = spans[3].get_attribute(SpanAttributeKey.CHAT_USAGE)
|
|
assert chat_usage[TokenUsageKey.INPUT_TOKENS] == 9
|
|
assert chat_usage[TokenUsageKey.OUTPUT_TOKENS] == 12
|
|
assert chat_usage[TokenUsageKey.TOTAL_TOKENS] == 21
|
|
assert spans[3].get_attribute(SpanAttributeKey.SPAN_TYPE) == SpanType.CHAT_MODEL
|
|
assert spans[3].model_name == "gpt-4o-mini"
|
|
if not IS_TRACING_SDK_ONLY:
|
|
# Verify cost is calculated (9 input tokens * 1.0 + 12 output tokens * 2.0)
|
|
assert spans[3].llm_cost == {
|
|
"input_cost": 9.0,
|
|
"output_cost": 24.0,
|
|
"total_cost": 33.0,
|
|
}
|
|
|
|
# OpenAI autologging
|
|
if with_openai_autolog:
|
|
assert spans[4].name == "AsyncCompletions"
|
|
assert spans[4].span_type == SpanType.CHAT_MODEL
|
|
assert spans[4].parent_id == spans[3].span_id
|
|
assert spans[4].inputs == {
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"model": "gpt-4o-mini",
|
|
"stream": False,
|
|
}
|
|
assert spans[4].get_attribute(SpanAttributeKey.CHAT_USAGE) == {
|
|
"input_tokens": 9,
|
|
"output_tokens": 12,
|
|
"total_tokens": 21,
|
|
}
|
|
assert spans[4].model_name == "gpt-4o-mini"
|
|
if not IS_TRACING_SDK_ONLY:
|
|
assert spans[4].llm_cost == {
|
|
"input_cost": 9.0,
|
|
"output_cost": 24.0,
|
|
"total_cost": 33.0,
|
|
}
|
|
|
|
# Trace level token usage should not double-count
|
|
assert trace.info.token_usage == {
|
|
"input_tokens": 9,
|
|
"output_tokens": 12,
|
|
"total_tokens": 21,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_invoke_simple_with_sk_initialization_of_tracer(mock_openai):
|
|
from opentelemetry.sdk.resources import Resource
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
|
|
from opentelemetry.semconv.resource import ResourceAttributes
|
|
from opentelemetry.trace import get_tracer_provider, set_tracer_provider
|
|
|
|
resource = Resource.create({ResourceAttributes.SERVICE_NAME: "telemetry-console-quickstart"})
|
|
tracer_provider = TracerProvider(resource=resource)
|
|
tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
|
|
set_tracer_provider(tracer_provider)
|
|
|
|
mlflow.semantic_kernel.autolog()
|
|
_tracer_provider = get_tracer_provider()
|
|
assert isinstance(_tracer_provider, TracerProvider)
|
|
span_processors = _tracer_provider._active_span_processor._span_processors
|
|
assert len(span_processors) == 2
|
|
assert any(isinstance(p, SemanticKernelSpanProcessor) for p in span_processors)
|
|
|
|
_ = await _create_and_invoke_kernel_simple(mock_openai)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
assert trace.info.request_id
|
|
assert len(trace.data.spans) == 4
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_invoke_complex(mock_openai, mock_litellm_cost):
|
|
mlflow.semantic_kernel.autolog()
|
|
result = await _create_and_invoke_kernel_complex(mock_openai)
|
|
|
|
# Validate the result is not mutated by tracing logic
|
|
assert isinstance(result, FunctionResult)
|
|
assert isinstance(result.value[0], ChatMessageContent)
|
|
assert result.value[0].items[0].text.startswith('[{"role": "system",')
|
|
|
|
# Trace
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
spans = traces[0].data.spans
|
|
assert len(spans) == 3 # Kernel.invoke, execute_tool, chat.completions
|
|
|
|
kernel_span, tool_span, chat_span = spans
|
|
assert kernel_span.name == "Kernel.invoke"
|
|
assert kernel_span.span_type == SpanType.AGENT
|
|
function_metadata = kernel_span.inputs["function"]["metadata"]
|
|
assert function_metadata["name"] == "Chat"
|
|
assert function_metadata["plugin_name"] == "ChatBot"
|
|
prompt = kernel_span.inputs["function"]["prompt_template"]["prompt_template_config"]
|
|
assert prompt["template"] == "{{$chat_history}}{{$user_input}}"
|
|
arguments = kernel_span.inputs["arguments"]
|
|
assert arguments["user_input"] == "I want to find a hotel in Seattle with free wifi and a pool."
|
|
assert len(arguments["chat_history"]) == 2
|
|
|
|
assert tool_span.name == "execute_tool ChatBot-Chat"
|
|
assert tool_span.span_type == SpanType.TOOL
|
|
assert tool_span.parent_id == kernel_span.span_id
|
|
|
|
assert chat_span.name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini")
|
|
assert chat_span.parent_id == tool_span.span_id
|
|
assert chat_span.span_type == SpanType.CHAT_MODEL
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.OPERATION).startswith("chat")
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.SYSTEM) == "openai"
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.MODEL) == "gpt-4o-mini"
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.RESPONSE_ID) == "chatcmpl-123"
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.FINISH_REASON) == "FinishReason.STOP"
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.INPUT_TOKENS) == 9
|
|
assert chat_span.get_attribute(model_gen_ai_attributes.OUTPUT_TOKENS) == 12
|
|
assert chat_span.model_name == "gpt-4o-mini"
|
|
|
|
assert any(
|
|
"I want to find a hotel in Seattle with free wifi and a pool." in m.get("content", "")
|
|
for m in chat_span.inputs.get("messages", [])
|
|
)
|
|
assert isinstance(chat_span.outputs["messages"], list)
|
|
|
|
chat_usage = chat_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
|
|
assert chat_usage[TokenUsageKey.INPUT_TOKENS] == 9
|
|
assert chat_usage[TokenUsageKey.OUTPUT_TOKENS] == 12
|
|
assert chat_usage[TokenUsageKey.TOTAL_TOKENS] == 21
|
|
if not IS_TRACING_SDK_ONLY:
|
|
assert chat_span.llm_cost == {
|
|
"input_cost": 9.0,
|
|
"output_cost": 24.0,
|
|
"total_cost": 33.0,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_invoke_agent(mock_openai):
|
|
mlflow.semantic_kernel.autolog()
|
|
result = await _create_and_invoke_chat_agent(mock_openai)
|
|
assert isinstance(result, AgentResponseItem)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
spans = trace.data.spans
|
|
assert len(spans) == 3
|
|
|
|
root_span, child_span, grandchild_span = spans
|
|
|
|
assert root_span.name == "invoke_agent sushi_agent"
|
|
assert root_span.span_type == SpanType.AGENT
|
|
assert root_span.get_attribute(model_gen_ai_attributes.OPERATION) == "invoke_agent"
|
|
assert root_span.get_attribute(agent_gen_ai_attributes.AGENT_NAME) == "sushi_agent"
|
|
|
|
assert child_span.name == "AutoFunctionInvocationLoop"
|
|
assert child_span.span_type == SpanType.UNKNOWN
|
|
assert "sk.available_functions" in child_span.attributes
|
|
|
|
assert grandchild_span.name.startswith("chat")
|
|
assert grandchild_span.span_type == SpanType.CHAT_MODEL
|
|
assert grandchild_span.get_attribute(model_gen_ai_attributes.MODEL) == "gpt-4o-mini"
|
|
assert grandchild_span.model_name == "gpt-4o-mini"
|
|
assert isinstance(grandchild_span.inputs["messages"], list)
|
|
assert isinstance(grandchild_span.outputs["messages"], list)
|
|
assert (
|
|
grandchild_span.get_attribute(model_gen_ai_attributes.FINISH_REASON) == "FinishReason.STOP"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_autolog_trace_on_exception(mock_openai):
|
|
mlflow.semantic_kernel.autolog()
|
|
openai_client = openai.AsyncOpenAI(api_key="test", base_url=mock_openai)
|
|
|
|
kernel = Kernel()
|
|
kernel.add_service(
|
|
OpenAIChatCompletion(
|
|
service_id="chat-gpt",
|
|
ai_model_id="gpt-4o-mini",
|
|
async_client=openai_client,
|
|
)
|
|
)
|
|
|
|
error_message = "thiswillfail"
|
|
with mock.patch.object(
|
|
openai_client.chat.completions, "create", side_effect=RuntimeError(error_message)
|
|
):
|
|
with pytest.raises(
|
|
KernelInvokeException, match="Error occurred while invoking function"
|
|
) as exc_info:
|
|
await kernel.invoke_prompt("Hello?")
|
|
|
|
assert isinstance(exc_info.value.__cause__, FunctionExecutionException)
|
|
|
|
traces = get_traces()
|
|
assert traces, "No traces recorded"
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
assert len(trace.data.spans) == 4
|
|
assert trace.info.status == "ERROR"
|
|
|
|
_, _, _, llm_span = trace.data.spans
|
|
assert llm_span.status.status_code == SpanStatusCode.ERROR
|
|
assert llm_span.events[0].name == "exception"
|
|
assert error_message in llm_span.events[0].attributes["exception.message"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tracing_autolog_with_active_span(mock_openai, with_openai_autolog):
|
|
mlflow.semantic_kernel.autolog()
|
|
|
|
with mlflow.start_span("parent"):
|
|
response = await _create_and_invoke_kernel_simple(mock_openai)
|
|
|
|
assert isinstance(response, FunctionResult)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
|
|
trace = traces[0]
|
|
spans = trace.data.spans
|
|
assert len(spans) == (6 if with_openai_autolog else 5)
|
|
|
|
assert trace.info.request_id is not None
|
|
assert trace.info.status == "OK"
|
|
assert trace.info.tags["mlflow.traceName"] == "parent"
|
|
|
|
parent = trace.data.spans[0]
|
|
assert parent.name == "parent"
|
|
assert parent.parent_id is None
|
|
assert parent.span_type == SpanType.UNKNOWN
|
|
|
|
assert spans[1].name == "Kernel.invoke_prompt"
|
|
assert spans[1].parent_id == parent.span_id
|
|
assert spans[2].name == "Kernel.invoke"
|
|
assert spans[2].parent_id == spans[1].span_id
|
|
assert spans[3].name.startswith("execute_tool")
|
|
assert spans[3].parent_id == spans[2].span_id
|
|
assert spans[4].name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini")
|
|
assert spans[4].parent_id == spans[3].span_id
|
|
|
|
if with_openai_autolog:
|
|
assert spans[5].name == "AsyncCompletions"
|
|
assert spans[5].parent_id == spans[4].span_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tracing_attribution_with_threaded_calls(mock_openai):
|
|
mlflow.semantic_kernel.autolog()
|
|
|
|
n = 3
|
|
openai_client = openai.AsyncOpenAI(api_key="test", base_url=mock_openai)
|
|
|
|
kernel = Kernel()
|
|
kernel.add_service(
|
|
OpenAIChatCompletion(
|
|
service_id="chat-gpt",
|
|
ai_model_id="gpt-4o-mini",
|
|
async_client=openai_client,
|
|
)
|
|
)
|
|
|
|
async def call(prompt: str):
|
|
return await kernel.invoke_prompt(prompt)
|
|
|
|
prompts = [f"What is this number: {i}" for i in range(n)]
|
|
_ = await asyncio.gather(*(call(p) for p in prompts))
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == n
|
|
|
|
unique_messages = set()
|
|
for trace in traces:
|
|
spans = trace.data.spans
|
|
assert len(spans) == 4
|
|
|
|
assert spans[0].span_type == SpanType.AGENT
|
|
assert spans[1].span_type == SpanType.AGENT
|
|
assert spans[2].span_type == SpanType.TOOL
|
|
assert spans[3].span_type == SpanType.CHAT_MODEL
|
|
assert spans[3].model_name == "gpt-4o-mini"
|
|
|
|
message = spans[3].inputs["messages"][0]["content"]
|
|
assert message.startswith("What is this number: ")
|
|
unique_messages.add(message)
|
|
assert spans[3].outputs["messages"][0]["content"]
|
|
|
|
assert len(unique_messages) == n
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("create_and_invoke_func", "span_name_pattern", "expected_span_input_keys"),
|
|
[
|
|
(
|
|
_create_and_invoke_kernel_simple,
|
|
"chat",
|
|
["messages"],
|
|
),
|
|
(
|
|
_create_and_invoke_text_completion,
|
|
"text",
|
|
# Text completion input should be stored as a raw string
|
|
None,
|
|
),
|
|
(
|
|
_create_and_invoke_chat_completion_direct,
|
|
"chat",
|
|
["messages"],
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_sk_input_parsing(
|
|
mock_openai, create_and_invoke_func, span_name_pattern, expected_span_input_keys
|
|
):
|
|
mlflow.semantic_kernel.autolog()
|
|
|
|
_ = await create_and_invoke_func(mock_openai)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
trace = traces[0]
|
|
|
|
target_span = None
|
|
for span in trace.data.spans:
|
|
if span_name_pattern in span.name:
|
|
target_span = span
|
|
break
|
|
|
|
assert target_span is not None, f"No span found with pattern '{span_name_pattern}'"
|
|
if expected_span_input_keys:
|
|
for key in expected_span_input_keys:
|
|
assert key in target_span.inputs, (
|
|
f"Expected '{key}' in span inputs for {target_span.name}, got: {target_span.inputs}"
|
|
)
|
|
else:
|
|
assert isinstance(target_span.inputs, str)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_invoke_with_kernel_arguments(mock_openai):
|
|
mlflow.semantic_kernel.autolog()
|
|
_ = await _create_and_invoke_kernel_function_object(mock_openai)
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
|
|
# Check that kernel arguments were passed through to the prompt
|
|
child_span = next(s for s in traces[0].data.spans if "chat" in s.name)
|
|
assert child_span.inputs["messages"][0]["content"] == "Add 5 and 3"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_embeddings(mock_openai):
|
|
mlflow.semantic_kernel.autolog()
|
|
|
|
result = await _create_and_invoke_embeddings(mock_openai)
|
|
|
|
assert result is not None
|
|
assert len(result) == 3
|
|
|
|
# NOTE: Semantic Kernel currently does not instrument embeddings with OpenTelemetry
|
|
# spans, so no traces are generated for embedding operations
|
|
traces = get_traces()
|
|
assert len(traces) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_kernel_invoke_function_object(mock_openai):
|
|
mlflow.semantic_kernel.autolog()
|
|
|
|
await _create_and_invoke_kernel_function_object(mock_openai)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
|
|
# Verify trace structure
|
|
assert len(traces[0].data.spans) == 3
|
|
|
|
# Root span should be execute_tool
|
|
kernel_span, tool_span, chat_span = traces[0].data.spans
|
|
|
|
assert kernel_span.name == "Kernel.invoke"
|
|
assert kernel_span.span_type == SpanType.AGENT
|
|
assert kernel_span.inputs["function"] is not None
|
|
assert kernel_span.outputs is not None
|
|
|
|
assert tool_span.name == "execute_tool MathPlugin-Add"
|
|
assert tool_span.span_type == SpanType.TOOL
|
|
|
|
# Child span should be chat completion
|
|
assert chat_span.name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini")
|
|
assert chat_span.span_type == SpanType.CHAT_MODEL
|
|
assert chat_span.model_name == "gpt-4o-mini"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sk_shared_provider_no_recursion(monkeypatch, mock_openai):
|
|
# Verify semantic_kernel.autolog() works with shared tracer provider (no RecursionError)
|
|
monkeypatch.setenv(MLFLOW_USE_DEFAULT_TRACER_PROVIDER.name, "false")
|
|
|
|
mlflow.semantic_kernel.autolog()
|
|
result = await _create_and_invoke_kernel_simple(mock_openai)
|
|
|
|
assert isinstance(result, FunctionResult)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
spans = traces[0].data.spans
|
|
assert len(spans) >= 3
|
|
assert spans[0].name == "Kernel.invoke_prompt"
|