mirrors / mlflow--mlflow

项目文件夹

mlflow--mlflow
文件
main
T
wehub-resource-sync 4b22cfda96 chore: import upstream snapshot with attribution
2026-07-13 13:22:34 +08:00

1249 行
42 KiB
Python
原始文件 永久链接 Blame 文件历史

import json
import re
import sys
from unittest import mock
import httpx
import openai
import pytest
from openai.resources.chat.completions import Completions as ChatCompletions
from openai.resources.completions import Completions
from openai.resources.embeddings import Embeddings
from packaging.version import Version
from pydantic import BaseModel
import mlflow
from mlflow.entities import SpanLogLevel
from mlflow.entities.span import SpanType
from mlflow.exceptions import MlflowException
from mlflow.openai.autolog import _get_span_type
from mlflow.openai.utils.chat_schema import _parse_tools
from mlflow.tracing.constant import (
STREAM_CHUNK_EVENT_VALUE_KEY,
CostKey,
SpanAttributeKey,
TokenUsageKey,
TraceMetadataKey,
)
from mlflow.version import IS_TRACING_SDK_ONLY
from tests.openai.mock_openai import AZURE_ANNOTATIONS, EMPTY_CHOICES, LIST_CONTENT
from tests.tracing.helper import get_traces, skip_when_testing_trace_sdk
MOCK_TOOLS = [
{
"type": "function",
"function": {
"name": "add",
"description": "Add two numbers",
"parameters": {
"type": "object",
"properties": {
"a": {"type": "number"},
"b": {"type": "number"},
},
"required": ["a", "b"],
},
},
}
]
@pytest.fixture(params=[True, False], ids=["sync", "async"])
def client(request, monkeypatch, mock_openai):
monkeypatch.setenv("OPENAI_API_KEY", "test")
monkeypatch.setenv("OPENAI_API_BASE", mock_openai)
if request.param:
client = openai.OpenAI(api_key="test", base_url=mock_openai)
client._is_async = False
return client
else:
client = openai.AsyncOpenAI(api_key="test", base_url=mock_openai)
client._is_async = True
return client
@pytest.fixture
def completion_models():
return [
mlflow.openai.log_model(
"gpt-4o-mini",
"completions",
name="model",
temperature=temp,
prompt="Say {text}",
pip_requirements=["mlflow"], # Hard code for speed up
)
for temp in [0.1, 0.2, 0.3]
]
@pytest.fixture
def embedding_models():
float_model = mlflow.openai.log_model(
"text-embedding-ada-002",
"embeddings",
name="model",
encoding_format="float",
pip_requirements=["mlflow"], # Hard code for speed up
)
base64_model = mlflow.openai.log_model(
"text-embedding-ada-002",
"embeddings",
name="model",
encoding_format="base64",
pip_requirements=["mlflow"], # Hard code for speed up
)
return [float_model, base64_model]
@pytest.mark.asyncio
@pytest.mark.skipif(
Version(openai.__version__) < Version("1.66"), reason="Cost tracking does not work before 1.66"
)
async def test_chat_completions_autolog(client, mock_litellm_cost):
mlflow.openai.autolog()
messages = [{"role": "user", "content": "test"}]
response = client.chat.completions.create(
messages=messages,
model="gpt-4o-mini",
temperature=0,
)
if client._is_async:
await response
traces = get_traces()
assert len(traces) == 1
trace = traces[0]
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.CHAT_MODEL
assert span.log_level == SpanLogLevel.INFO
assert span.inputs == {"messages": messages, "model": "gpt-4o-mini", "temperature": 0}
assert span.outputs["id"] == "chatcmpl-123"
assert span.attributes["model"] == "gpt-4o-mini"
assert span.attributes["temperature"] == 0
assert span.get_attribute(SpanAttributeKey.CHAT_USAGE) == {
TokenUsageKey.INPUT_TOKENS: 9,
TokenUsageKey.OUTPUT_TOKENS: 12,
TokenUsageKey.TOTAL_TOKENS: 21,
}
assert span.model_name == "gpt-4o-mini"
assert span.get_attribute(SpanAttributeKey.MESSAGE_FORMAT) == "openai"
if not IS_TRACING_SDK_ONLY:
# Verify cost is calculated (9 input tokens * 1.0 + 12 output tokens * 2.0)
assert span.llm_cost == {
"input_cost": 9.0,
"output_cost": 24.0,
"total_cost": 33.0,
}
assert TraceMetadataKey.SOURCE_RUN not in trace.info.request_metadata
assert trace.info.token_usage == {
TokenUsageKey.INPUT_TOKENS: 9,
TokenUsageKey.OUTPUT_TOKENS: 12,
TokenUsageKey.TOTAL_TOKENS: 21,
}
if not IS_TRACING_SDK_ONLY:
assert trace.info.cost == {
CostKey.INPUT_COST: 9.0,
CostKey.OUTPUT_COST: 24.0,
CostKey.TOTAL_COST: 33.0,
}
def test_get_span_type_resolves_subclasses():
# Regression test for https://github.com/mlflow/mlflow/issues/23754:
# third-party clients (e.g. DatabricksOpenAI) subclass the OpenAI resource
# classes, so exact dict lookup misses and a previous fallback returned an
# unhashable tuple, breaking span log-level resolution downstream.
class CustomChatCompletions(ChatCompletions):
pass
assert _get_span_type(CustomChatCompletions) == SpanType.CHAT_MODEL
assert _get_span_type(ChatCompletions) == SpanType.CHAT_MODEL
assert _get_span_type(Completions) == SpanType.LLM
assert _get_span_type(Embeddings) == SpanType.EMBEDDING
assert _get_span_type(object) == SpanType.UNKNOWN
@pytest.mark.asyncio
@pytest.mark.skipif(
Version(openai.__version__) < Version("1.66"), reason="Cost tracking does not work before 1.66"
)
async def test_chat_completions_autolog_with_cached_tokens(client, mock_litellm_cost):
mlflow.openai.autolog()
mock_response = {
"id": "chatcmpl-cached",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hello"},
"logprobs": None,
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 50,
"completion_tokens": 20,
"total_tokens": 70,
"prompt_tokens_details": {"cached_tokens": 30, "audio_tokens": 0},
"completion_tokens_details": {"reasoning_tokens": 0},
},
}
if client._is_async:
patch_target = "httpx.AsyncClient.send"
async def send_patch(self, request, *args, **kwargs):
return httpx.Response(status_code=200, request=request, json=mock_response)
else:
patch_target = "httpx.Client.send"
def send_patch(self, request, *args, **kwargs):
return httpx.Response(status_code=200, request=request, json=mock_response)
with mock.patch(patch_target, send_patch):
response = client.chat.completions.create(
messages=[{"role": "user", "content": "test"}],
model="gpt-4o-mini",
temperature=0,
)
if client._is_async:
response = await response
traces = get_traces()
assert len(traces) == 1
span = traces[0].data.spans[0]
assert span.get_attribute(SpanAttributeKey.CHAT_USAGE) == {
TokenUsageKey.INPUT_TOKENS: 50,
TokenUsageKey.OUTPUT_TOKENS: 20,
TokenUsageKey.TOTAL_TOKENS: 70,
TokenUsageKey.CACHE_READ_INPUT_TOKENS: 30,
}
assert traces[0].info.token_usage == {
TokenUsageKey.INPUT_TOKENS: 50,
TokenUsageKey.OUTPUT_TOKENS: 20,
TokenUsageKey.TOTAL_TOKENS: 70,
TokenUsageKey.CACHE_READ_INPUT_TOKENS: 30,
}
@pytest.mark.asyncio
@pytest.mark.skipif(
Version(openai.__version__) < Version("1.66"), reason="Cost tracking does not work before 1.66"
)
async def test_chat_completions_autolog_under_current_active_span(client):
# If a user have an active span, the autologging should create a child span under it.
mlflow.openai.autolog()
messages = [{"role": "user", "content": "test"}]
with mlflow.start_span(name="parent"):
for _ in range(3):
response = client.chat.completions.create(
messages=messages,
model="gpt-4o-mini",
temperature=0,
)
if client._is_async:
await response
traces = get_traces()
assert len(traces) == 1
trace = traces[0]
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 4
parent_span = trace.data.spans[0]
assert parent_span.name == "parent"
child_span = trace.data.spans[1]
assert child_span.name == "AsyncCompletions" if client._is_async else "Completions"
assert child_span.inputs == {"messages": messages, "model": "gpt-4o-mini", "temperature": 0}
assert child_span.outputs["id"] == "chatcmpl-123"
assert child_span.parent_id == parent_span.span_id
# Token usage should be aggregated correctly
assert trace.info.token_usage == {
TokenUsageKey.INPUT_TOKENS: 27,
TokenUsageKey.OUTPUT_TOKENS: 36,
TokenUsageKey.TOTAL_TOKENS: 63,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("include_usage", [True, False])
async def test_chat_completions_autolog_streaming(client, include_usage):
mlflow.openai.autolog()
stream_options_supported = Version(openai.__version__) >= Version("1.26")
if not stream_options_supported and include_usage:
pytest.skip("OpenAI SDK version does not support usage tracking in streaming")
messages = [{"role": "user", "content": "test"}]
input_params = {
"messages": messages,
"model": "gpt-4o-mini",
"temperature": 0,
"stream": True,
}
if stream_options_supported:
input_params["stream_options"] = {"include_usage": include_usage}
stream = client.chat.completions.create(**input_params)
if client._is_async:
async for _ in await stream:
pass
else:
for _ in stream:
pass
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.CHAT_MODEL
assert span.inputs == input_params
# Reconstructed response from streaming chunks
assert isinstance(span.outputs, dict)
assert span.outputs["id"] == "chatcmpl-123"
assert span.outputs["object"] == "chat.completion"
assert span.outputs["model"] == "gpt-4o-mini"
assert span.outputs["system_fingerprint"] == "fp_44709d6fcb"
assert "choices" in span.outputs
assert span.outputs["choices"][0]["message"]["role"] == "assistant"
assert span.outputs["choices"][0]["message"]["content"] == "Hello world"
# Usage should be preserved when include_usage=True
if include_usage:
assert "usage" in span.outputs
assert span.outputs["usage"]["prompt_tokens"] == 9
assert span.outputs["usage"]["completion_tokens"] == 12
assert span.outputs["usage"]["total_tokens"] == 21
stream_event_data = trace.data.spans[0].events
assert stream_event_data[0].name == "mlflow.chunk.item.0"
chunk_1 = json.loads(stream_event_data[0].attributes[STREAM_CHUNK_EVENT_VALUE_KEY])
assert chunk_1["id"] == "chatcmpl-123"
assert chunk_1["choices"][0]["delta"]["content"] == "Hello"
assert stream_event_data[1].name == "mlflow.chunk.item.1"
chunk_2 = json.loads(stream_event_data[1].attributes[STREAM_CHUNK_EVENT_VALUE_KEY])
assert chunk_2["id"] == "chatcmpl-123"
assert chunk_2["choices"][0]["delta"]["content"] == " world"
if include_usage:
assert trace.info.token_usage == {
TokenUsageKey.INPUT_TOKENS: 9,
TokenUsageKey.OUTPUT_TOKENS: 12,
TokenUsageKey.TOTAL_TOKENS: 21,
}
@pytest.mark.asyncio
async def test_chat_completions_autolog_tracing_error(client):
mlflow.openai.autolog()
messages = [{"role": "user", "content": "test"}]
with pytest.raises(openai.UnprocessableEntityError, match="Input should be less"): # noqa: PT012
response = client.chat.completions.create(
messages=messages,
model="gpt-4o-mini",
temperature=5.0,
)
if client._is_async:
await response
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace.info.status == "ERROR"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.name == "AsyncCompletions" if client._is_async else "Completions"
assert span.inputs["messages"][0]["content"] == "test"
assert span.outputs is None
assert span.events[0].name == "exception"
assert span.events[0].attributes["exception.type"] == "UnprocessableEntityError"
@pytest.mark.asyncio
async def test_chat_completions_autolog_tracing_error_with_parent_span(client):
mlflow.openai.autolog()
if client._is_async:
@mlflow.trace
async def create_completions(text: str) -> str:
try:
response = await client.chat.completions.create(
messages=[{"role": "user", "content": text}],
model="gpt-4o-mini",
temperature=5.0,
)
return response.choices[0].delta.content
except openai.OpenAIError as e:
raise MlflowException("Failed to create completions") from e
with pytest.raises(MlflowException, match="Failed to create completions"):
await create_completions("test")
else:
@mlflow.trace
def create_completions(text: str) -> str:
try:
response = client.chat.completions.create(
messages=[{"role": "user", "content": text}],
model="gpt-4o-mini",
temperature=5.0,
)
return response.choices[0].delta.content
except openai.OpenAIError as e:
raise MlflowException("Failed to create completions") from e
with pytest.raises(MlflowException, match="Failed to create completions"):
create_completions("test")
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace.info.status == "ERROR"
assert len(trace.data.spans) == 2
parent_span = trace.data.spans[0]
assert parent_span.name == "create_completions"
assert parent_span.inputs == {"text": "test"}
assert parent_span.outputs is None
assert parent_span.status.status_code == "ERROR"
assert parent_span.events[0].name == "exception"
assert parent_span.events[0].attributes["exception.type"] == "MlflowException"
assert parent_span.events[0].attributes["exception.message"] == "Failed to create completions"
child_span = trace.data.spans[1]
assert child_span.name == "AsyncCompletions" if client._is_async else "Completions"
assert child_span.inputs["messages"][0]["content"] == "test"
assert child_span.outputs is None
assert child_span.status.status_code == "ERROR"
assert child_span.events[0].name == "exception"
assert child_span.events[0].attributes["exception.type"] == "UnprocessableEntityError"
@pytest.mark.asyncio
async def test_chat_completions_streaming_empty_choices(client):
mlflow.openai.autolog()
stream = client.chat.completions.create(
messages=[{"role": "user", "content": EMPTY_CHOICES}],
model="gpt-4o-mini",
stream=True,
)
chunks = [chunk async for chunk in await stream] if client._is_async else list(stream)
# Ensure the stream has a chunk with empty choices
assert chunks[0].choices == []
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace.info.status == "OK"
@pytest.mark.asyncio
async def test_chat_completions_streaming_ignores_azure_annotation_chunks(client):
mlflow.openai.autolog()
stream = client.chat.completions.create(
messages=[{"role": "user", "content": AZURE_ANNOTATIONS}],
model="gpt-4o-mini",
stream=True,
stream_options={"include_usage": True},
)
chunks = [chunk async for chunk in await stream] if client._is_async else list(stream)
assert chunks[0].object == ""
assert chunks[-1].object == ""
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace is not None
assert trace.info.status == "OK"
assert trace.info.token_usage == {
TokenUsageKey.INPUT_TOKENS: 9,
TokenUsageKey.OUTPUT_TOKENS: 12,
TokenUsageKey.TOTAL_TOKENS: 21,
}
span = trace.data.spans[0]
assert span.outputs["id"] == "chatcmpl-123"
assert span.outputs["model"] == "gpt-4o-mini"
assert span.outputs["choices"][0]["message"]["content"] == "Hello world"
assert span.outputs["usage"]["prompt_tokens"] == 9
assert span.get_attribute(SpanAttributeKey.CHAT_USAGE) == {
TokenUsageKey.INPUT_TOKENS: 9,
TokenUsageKey.OUTPUT_TOKENS: 12,
TokenUsageKey.TOTAL_TOKENS: 21,
}
@pytest.mark.asyncio
async def test_chat_completions_streaming_with_list_content(client):
# Test streaming with Databricks-style list content in chunks.
mlflow.openai.autolog()
stream = client.chat.completions.create(
messages=[{"role": "user", "content": LIST_CONTENT}],
model="gpt-4o-mini",
stream=True,
)
chunks = [chunk async for chunk in await stream] if client._is_async else list(stream)
assert len(chunks) == 2
assert chunks[0].choices[0].delta.content == [{"type": "text", "text": "Hello"}]
assert chunks[1].choices[0].delta.content == [{"type": "text", "text": " world"}]
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.CHAT_MODEL
# Verify the reconstructed message content is correct (text extracted from list)
assert isinstance(span.outputs, dict)
assert span.outputs["choices"][0]["message"]["content"] == "Hello world"
@pytest.mark.asyncio
async def test_completions_autolog(client):
mlflow.openai.autolog()
response = client.completions.create(
prompt="test",
model="gpt-4o-mini",
temperature=0,
)
if client._is_async:
await response
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.LLM
assert span.inputs == {"prompt": "test", "model": "gpt-4o-mini", "temperature": 0}
assert span.outputs["id"] == "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7"
assert span.model_name == "gpt-4o-mini"
assert span.get_attribute(SpanAttributeKey.MESSAGE_FORMAT) == "openai"
assert TraceMetadataKey.SOURCE_RUN not in trace.info.request_metadata
@pytest.mark.asyncio
async def test_completions_autolog_streaming_empty_choices(client):
mlflow.openai.autolog()
stream = client.completions.create(
prompt=EMPTY_CHOICES,
model="gpt-4o-mini",
stream=True,
)
chunks = [chunk async for chunk in await stream] if client._is_async else list(stream)
# Ensure the stream has a chunk with empty choices
assert chunks[0].choices == []
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace.info.status == "OK"
@pytest.mark.asyncio
async def test_completions_autolog_streaming(client):
mlflow.openai.autolog()
stream = client.completions.create(
prompt="test",
model="gpt-4o-mini",
temperature=0,
stream=True,
)
if client._is_async:
async for _ in await stream:
pass
else:
for _ in stream:
pass
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.LLM
assert span.inputs == {
"prompt": "test",
"model": "gpt-4o-mini",
"temperature": 0,
"stream": True,
}
assert span.outputs == "Hello world" # aggregated string of streaming response
stream_event_data = trace.data.spans[0].events
assert stream_event_data[0].name == "mlflow.chunk.item.0"
chunk_1 = json.loads(stream_event_data[0].attributes[STREAM_CHUNK_EVENT_VALUE_KEY])
assert chunk_1["id"] == "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7"
assert chunk_1["choices"][0]["text"] == "Hello"
assert stream_event_data[1].name == "mlflow.chunk.item.1"
chunk_2 = json.loads(stream_event_data[1].attributes[STREAM_CHUNK_EVENT_VALUE_KEY])
assert chunk_2["id"] == "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7"
assert chunk_2["choices"][0]["text"] == " world"
@pytest.mark.asyncio
async def test_embeddings_autolog(client):
mlflow.openai.autolog()
response = client.embeddings.create(
input="test",
model="text-embedding-ada-002",
)
if client._is_async:
await response
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert trace is not None
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.EMBEDDING
assert span.log_level == SpanLogLevel.INFO
assert span.inputs == {"input": "test", "model": "text-embedding-ada-002"}
assert span.outputs["data"][0]["embedding"] == list(range(1536))
assert span.model_name == "text-embedding-ada-002"
assert TraceMetadataKey.SOURCE_RUN not in trace.info.request_metadata
@skip_when_testing_trace_sdk
@pytest.mark.asyncio
async def test_autolog_use_active_run_id(client):
mlflow.openai.autolog()
messages = [{"role": "user", "content": "test"}]
async def _call_create():
response = client.chat.completions.create(messages=messages, model="gpt-4o-mini")
if client._is_async:
await response
return response
with mlflow.start_run() as run_1:
await _call_create()
with mlflow.start_run() as run_2:
await _call_create()
await _call_create()
with mlflow.start_run() as run_3:
mlflow.openai.autolog()
await _call_create()
traces = get_traces()[::-1] # reverse order to sort by timestamp in ascending order
assert len(traces) == 4
assert traces[0].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_1.info.run_id
assert traces[1].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_2.info.run_id
assert traces[2].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_2.info.run_id
assert traces[3].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_3.info.run_id
@pytest.mark.asyncio
async def test_autolog_raw_response(client):
mlflow.openai.autolog()
messages = [{"role": "user", "content": "test"}]
resp = client.chat.completions.with_raw_response.create(
model="gpt-4o-mini",
messages=messages,
tools=MOCK_TOOLS,
)
if client._is_async:
resp = await resp
resp = resp.parse() # ensure the raw response is returned
assert resp.choices[0].message.content == '[{"role": "user", "content": "test"}]'
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.CHAT_MODEL
assert isinstance(span.outputs, dict)
assert (
span.outputs["choices"][0]["message"]["content"] == '[{"role": "user", "content": "test"}]'
)
assert span.attributes[SpanAttributeKey.CHAT_TOOLS] == MOCK_TOOLS
assert span.model_name == "gpt-4o-mini"
assert trace.info.token_usage == {
TokenUsageKey.INPUT_TOKENS: 9,
TokenUsageKey.OUTPUT_TOKENS: 12,
TokenUsageKey.TOTAL_TOKENS: 21,
}
@pytest.mark.asyncio
async def test_autolog_raw_response_stream(client):
mlflow.openai.autolog()
messages = [{"role": "user", "content": "test"}]
resp = client.chat.completions.with_raw_response.create(
model="gpt-4o-mini",
messages=messages,
tools=MOCK_TOOLS,
stream=True,
)
if client._is_async:
resp = await resp
resp = resp.parse() # ensure the raw response is returned
if client._is_async:
chunks = [c.choices[0].delta.content async for c in resp]
else:
chunks = [c.choices[0].delta.content for c in resp]
assert chunks == ["Hello", " world"]
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.CHAT_MODEL
assert span.model_name == "gpt-4o-mini"
# Reconstructed response from streaming chunks
assert isinstance(span.outputs, dict)
assert span.outputs["id"] == "chatcmpl-123"
assert span.outputs["object"] == "chat.completion"
assert span.outputs["model"] == "gpt-4o-mini"
assert span.outputs["choices"][0]["message"]["content"] == "Hello world"
assert span.attributes[SpanAttributeKey.CHAT_TOOLS] == MOCK_TOOLS
@pytest.mark.skipif(
Version(openai.__version__) < Version("1.40"), reason="Requires OpenAI SDK >= 1.40"
)
@pytest.mark.asyncio
async def test_response_format(client):
mlflow.openai.autolog()
class Person(BaseModel):
name: str
age: int
mock_response = {
"id": "chatcmpl-Ax4UAd5xf32KjgLkS1SEEY9oorI9m",
"object": "chat.completion",
"created": 1738641958,
"model": "gpt-4o-2024-08-06",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": '{"name":"Angelo","age":42}',
"refusal": None,
},
"logprobs": None,
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 68,
"completion_tokens": 11,
"total_tokens": 79,
"prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0},
"completion_tokens_details": {
"reasoning_tokens": 0,
"audio_tokens": 0,
"accepted_prediction_tokens": 0,
"rejected_prediction_tokens": 0,
},
},
"service_tier": "default",
"system_fingerprint": "fp_50cad350e4",
}
if client._is_async:
patch_target = "httpx.AsyncClient.send"
async def send_patch(self, request, *args, **kwargs):
return httpx.Response(
status_code=200,
request=request,
json=mock_response,
)
else:
patch_target = "httpx.Client.send"
def send_patch(self, request, *args, **kwargs):
return httpx.Response(
status_code=200,
request=request,
json=mock_response,
)
with mock.patch(patch_target, send_patch):
response = client.beta.chat.completions.parse(
messages=[
{"role": "system", "content": "Extract info from text"},
{"role": "user", "content": "I am Angelo and I am 42."},
],
model="gpt-4o",
temperature=0,
response_format=Person,
)
if client._is_async:
response = await response
assert response.choices[0].message.parsed == Person(name="Angelo", age=42)
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.outputs["choices"][0]["message"]["content"] == '{"name":"Angelo","age":42}'
assert span.span_type == SpanType.CHAT_MODEL
assert span.model_name == "gpt-4o"
assert trace.info.trace_metadata.get(TraceMetadataKey.TOKEN_USAGE) == json.dumps({
TokenUsageKey.INPUT_TOKENS: 68,
TokenUsageKey.OUTPUT_TOKENS: 11,
TokenUsageKey.TOTAL_TOKENS: 79,
TokenUsageKey.CACHE_READ_INPUT_TOKENS: 0,
})
@skip_when_testing_trace_sdk
@pytest.mark.asyncio
async def test_autolog_link_traces_to_loaded_model_chat_completions(client, completion_models):
mlflow.openai.autolog()
for model_info in completion_models:
model_dict = mlflow.openai.load_model(model_info.model_uri)
resp = client.chat.completions.create(
messages=[{"role": "user", "content": f"test {model_info.model_id}"}],
model=model_dict["model"],
temperature=model_dict["temperature"],
)
if client._is_async:
await resp
traces = get_traces()
assert len(traces) == len(completion_models)
for trace in traces:
span = trace.data.spans[0]
model_id = trace.info.request_metadata[TraceMetadataKey.MODEL_ID]
assert model_id is not None
assert span.inputs["messages"][0]["content"] == f"test {model_id}"
assert span.model_name == model_dict["model"]
@skip_when_testing_trace_sdk
@pytest.mark.asyncio
async def test_autolog_link_traces_to_loaded_model_completions(client, completion_models):
mlflow.openai.autolog()
for model_info in completion_models:
model_dict = mlflow.openai.load_model(model_info.model_uri)
resp = client.completions.create(
prompt=f"test {model_info.model_id}",
model=model_dict["model"],
temperature=model_dict["temperature"],
)
if client._is_async:
await resp
traces = get_traces()
assert len(traces) == len(completion_models)
for trace in traces:
span = trace.data.spans[0]
model_id = trace.info.request_metadata[TraceMetadataKey.MODEL_ID]
assert model_id is not None
assert span.inputs["prompt"] == f"test {model_id}"
assert span.model_name == model_dict["model"]
@skip_when_testing_trace_sdk
@pytest.mark.asyncio
async def test_autolog_link_traces_to_loaded_model_embeddings(client, embedding_models):
mlflow.openai.autolog()
for model_info in embedding_models:
model_dict = mlflow.openai.load_model(model_info.model_uri)
resp = client.embeddings.create(
input=f"test {model_info.model_id}",
model=model_dict["model"],
encoding_format=model_dict["encoding_format"],
)
if client._is_async:
await resp
traces = get_traces()
assert len(traces) == len(embedding_models)
for trace in traces:
span = trace.data.spans[0]
model_id = trace.info.request_metadata[TraceMetadataKey.MODEL_ID]
assert model_id is not None
assert span.inputs["input"] == f"test {model_id}"
assert span.model_name == model_dict["model"]
@skip_when_testing_trace_sdk
def test_autolog_link_traces_to_loaded_model_embeddings_pyfunc(
monkeypatch, mock_openai, embedding_models
):
monkeypatch.setenv("OPENAI_API_KEY", "test")
monkeypatch.setenv("OPENAI_API_BASE", mock_openai)
mlflow.openai.autolog()
for model_info in embedding_models:
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
assert mlflow.get_active_model_id() == model_info.model_id
pyfunc_model.predict(f"test {model_info.model_id}")
traces = get_traces()
assert len(traces) == len(embedding_models)
for trace in traces:
span = trace.data.spans[0]
model_id = trace.info.request_metadata[TraceMetadataKey.MODEL_ID]
assert model_id is not None
assert span.inputs["input"] == [f"test {model_id}"]
assert span.model_name == "text-embedding-ada-002"
@skip_when_testing_trace_sdk
def test_autolog_link_traces_to_active_model(monkeypatch, mock_openai, embedding_models):
monkeypatch.setenv("OPENAI_API_KEY", "test")
monkeypatch.setenv("OPENAI_API_BASE", mock_openai)
model = mlflow.create_external_model(name="test_model")
mlflow.set_active_model(model_id=model.model_id)
mlflow.openai.autolog()
for model_info in embedding_models:
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
pyfunc_model.predict(model_info.model_id)
traces = get_traces()
assert len(traces) == len(embedding_models)
for trace in traces:
span = trace.data.spans[0]
assert trace.info.request_metadata[TraceMetadataKey.MODEL_ID] == model.model_id
logged_model_id = span.inputs["input"][0]
assert logged_model_id != model.model_id
assert span.model_name == "text-embedding-ada-002"
@pytest.mark.asyncio
async def test_images_generate_autolog(client):
mlflow.openai.autolog()
# Disable tracing header injection — safe_patch rejects the extra_headers
# dict as a "new input" because it's not an ExceptionSafe-wrapped object.
# This is a known limitation shared with other non-chat endpoints.
openai_autolog_module = sys.modules["mlflow.openai.autolog"]
with mock.patch.object(openai_autolog_module, "_inject_tracing_headers"):
response = client.images.generate(
model="dall-e-3",
prompt="a white siamese cat",
n=1,
response_format="b64_json",
)
if client._is_async:
await response
traces = get_traces()
assert len(traces) == 1
trace = traces[0]
assert trace.info.status == "OK"
assert len(trace.data.spans) == 1
span = trace.data.spans[0]
assert span.span_type == SpanType.TOOL
assert span.inputs["prompt"] == "a white siamese cat"
assert span.outputs["data"][0]["revised_prompt"] == "a test image"
@pytest.mark.parametrize(
"sentinel",
[None, 42, object()],
)
def test_parse_tools_handles_openai_not_given_sentinel(sentinel):
assert _parse_tools({"tools": sentinel}) == []
@skip_when_testing_trace_sdk
@pytest.mark.asyncio
async def test_model_loading_set_active_model_id_without_fetching_logged_model(
client, completion_models
):
mlflow.openai.autolog()
model_info = completion_models[0]
with mock.patch("mlflow.get_logged_model", side_effect=Exception("get_logged_model failed")):
model_dict = mlflow.openai.load_model(model_info.model_uri)
resp = client.chat.completions.create(
messages=[{"role": "user", "content": f"test {model_info.model_id}"}],
model=model_dict["model"],
temperature=model_dict["temperature"],
)
if client._is_async:
await resp
traces = get_traces()
assert len(traces) == 1
span = traces[0].data.spans[0]
model_id = traces[0].info.request_metadata[TraceMetadataKey.MODEL_ID]
assert model_id is not None
assert span.inputs["messages"][0]["content"] == f"test {model_id}"
assert span.model_name == model_dict["model"]
@pytest.mark.skipif(
Version(openai.__version__) < Version("1.66"), reason="Requires OpenAI SDK >= 1.66"
)
@skip_when_testing_trace_sdk
def test_reconstruct_response_from_stream():
from openai.types.responses import (
ResponseOutputItemDoneEvent,
ResponseOutputMessage,
ResponseOutputText,
)
from mlflow.openai.autolog import _reconstruct_response_from_stream
from mlflow.types.responses_helpers import OutputItem
content1 = ResponseOutputText(annotations=[], text="Hello", type="output_text")
content2 = ResponseOutputText(annotations=[], text=" world", type="output_text")
message1 = ResponseOutputMessage(
id="test-1", content=[content1], role="assistant", status="completed", type="message"
)
message2 = ResponseOutputMessage(
id="test-2", content=[content2], role="assistant", status="completed", type="message"
)
chunk1 = ResponseOutputItemDoneEvent(
item=message1, output_index=0, sequence_number=1, type="response.output_item.done"
)
chunk2 = ResponseOutputItemDoneEvent(
item=message2, output_index=1, sequence_number=2, type="response.output_item.done"
)
chunks = [chunk1, chunk2]
result = _reconstruct_response_from_stream(chunks)
assert result.output == [
OutputItem(**chunk1.item.to_dict()),
OutputItem(**chunk2.item.to_dict()),
]
@pytest.mark.asyncio
async def test_tracing_headers_injected(client):
mlflow.openai.autolog()
captured_request = {}
mock_response = {
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21},
}
if client._is_async:
patch_target = "httpx.AsyncClient.send"
original_send = httpx.AsyncClient.send
async def send_patch(self, request, *args, **kwargs):
if "chat/completions" in str(request.url):
captured_request["headers"] = dict(request.headers)
return httpx.Response(status_code=200, request=request, json=mock_response)
return await original_send(self, request, *args, **kwargs)
else:
patch_target = "httpx.Client.send"
original_send = httpx.Client.send
def send_patch(self, request, *args, **kwargs):
if "chat/completions" in str(request.url):
captured_request["headers"] = dict(request.headers)
return httpx.Response(status_code=200, request=request, json=mock_response)
return original_send(self, request, *args, **kwargs)
with mock.patch(patch_target, send_patch):
response = client.chat.completions.create(
messages=[{"role": "user", "content": "test"}],
model="gpt-4o-mini",
)
if client._is_async:
response = await response
# Verify traceparent header was injected
assert "traceparent" in captured_request["headers"]
traceparent = captured_request["headers"]["traceparent"]
assert re.fullmatch(r"00-[0-9a-f]{32}-[0-9a-f]{16}-[0-9a-f]{2}", traceparent)
# Verify the traceparent points to the LLM span
traces = get_traces()
assert len(traces) == 1
span = traces[0].data.spans[0]
span_ctx = span._span.get_span_context()
trace_id_hex = format(span_ctx.trace_id, "032x")
span_id_hex = format(span_ctx.span_id, "016x")
assert traceparent.startswith(f"00-{trace_id_hex}-{span_id_hex}-")
@pytest.mark.asyncio
async def test_tracing_headers_preserve_user_headers(client):
mlflow.openai.autolog()
captured_request = {}
mock_response = {
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21},
}
if client._is_async:
patch_target = "httpx.AsyncClient.send"
original_send = httpx.AsyncClient.send
async def send_patch(self, request, *args, **kwargs):
if "chat/completions" in str(request.url):
captured_request["headers"] = dict(request.headers)
return httpx.Response(status_code=200, request=request, json=mock_response)
return await original_send(self, request, *args, **kwargs)
else:
patch_target = "httpx.Client.send"
original_send = httpx.Client.send
def send_patch(self, request, *args, **kwargs):
if "chat/completions" in str(request.url):
captured_request["headers"] = dict(request.headers)
return httpx.Response(status_code=200, request=request, json=mock_response)
return original_send(self, request, *args, **kwargs)
with mock.patch(patch_target, send_patch):
response = client.chat.completions.create(
messages=[{"role": "user", "content": "test"}],
model="gpt-4o-mini",
extra_headers={"X-Custom": "my-value"},
)
if client._is_async:
response = await response
# User-provided headers should be preserved alongside traceparent
assert "traceparent" in captured_request["headers"]
assert captured_request["headers"].get("x-custom") == "my-value"
@pytest.mark.asyncio
@pytest.mark.skipif(
Version(openai.__version__) < Version("1.66"), reason="Cost tracking does not work before 1.66"
)
async def test_chat_completions_autolog_streaming_with_cached_tokens(client, mock_litellm_cost):
mlflow.openai.autolog()
mock_chunk = {
"id": "chatcmpl-stream-cached",
"object": "chat.completion.chunk",
"created": 1677652288,
"model": "gpt-4o-mini",
"choices": [],
"usage": {
"prompt_tokens": 50,
"completion_tokens": 20,
"total_tokens": 70,
"prompt_tokens_details": {"cached_tokens": 30, "audio_tokens": 0},
"completion_tokens_details": {"reasoning_tokens": 0},
},
}
if client._is_async:
patch_target = "httpx.AsyncClient.send"
async def send_patch(self, request, *args, **kwargs):
content = f"data: {json.dumps(mock_chunk)}\n\ndata: [DONE]\n\n".encode()
return httpx.Response(status_code=200, request=request, content=content)
else:
patch_target = "httpx.Client.send"
def send_patch(self, request, *args, **kwargs):
content = f"data: {json.dumps(mock_chunk)}\n\ndata: [DONE]\n\n".encode()
return httpx.Response(status_code=200, request=request, content=content)
with mock.patch(patch_target, send_patch):
stream = client.chat.completions.create(
messages=[{"role": "user", "content": "test"}],
model="gpt-4o-mini",
stream=True,
stream_options={"include_usage": True},
)
if client._is_async:
async for _ in await stream:
pass
else:
for _ in stream:
pass
traces = get_traces()
assert len(traces) == 1
span = traces[0].data.spans[0]
assert span.get_attribute(SpanAttributeKey.CHAT_USAGE) == {
TokenUsageKey.INPUT_TOKENS: 50,
TokenUsageKey.OUTPUT_TOKENS: 20,
TokenUsageKey.TOTAL_TOKENS: 70,
TokenUsageKey.CACHE_READ_INPUT_TOKENS: 30,
}
在新工单中引用 查看 Git Blame 复制永久链接