import pytest import mlflow from mlflow import get_experiment_by_name, set_experiment from mlflow.demo.base import DEMO_EXPERIMENT_NAME, DemoFeature, DemoResult from mlflow.demo.generators.traces import ( _PROVIDER_TO_LLM_SPAN_NAME, DEMO_SESSION_TURN_TAG, DEMO_TRACE_TYPE_TAG, DEMO_VERSION_TAG, TracesDemoGenerator, ) from mlflow.entities import SpanType from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey from mlflow.tracking._tracking_service.utils import _use_tracking_uri @pytest.fixture def traces_generator(): generator = TracesDemoGenerator() original_version = generator.version yield generator TracesDemoGenerator.version = original_version @pytest.fixture(scope="module") def generated_traces(tmp_path_factory): """Generate the demo once per module and return the materialized list of traces. Used by read-only structural tests so we don't re-generate the full demo per test. The fixture controls its own tracking URI (the autouse function-scoped tracking_uri_mock in the parent conftest doesn't apply at module setup time), and the returned traces are in-memory Python objects so consuming tests don't need the URI still set when they run. `flush=True` ensures the async trace export queue is drained before we read. Uses `_use_tracking_uri` (not get/set around `get_tracking_uri()`) so that the pre-fixture state — typically `_tracking_uri = None` falling back to the default — is restored exactly. Calling `set_tracking_uri(get_tracking_uri())` would materialise the default's absolute path and leave it stuck in `_tracking_uri`, poisoning later tests that opt out of the autouse fixture. """ db_path = tmp_path_factory.mktemp("demo_shared") / "mlflow.db" with _use_tracking_uri(f"sqlite:///{db_path}"): TracesDemoGenerator().generate() experiment = get_experiment_by_name(DEMO_EXPERIMENT_NAME) return mlflow.search_traces( locations=[experiment.experiment_id], max_results=100, return_type="list", flush=True, ) def test_generator_attributes(): generator = TracesDemoGenerator() assert generator.name == DemoFeature.TRACES assert generator.version == 3 def test_data_exists_false_when_no_experiment(): generator = TracesDemoGenerator() assert generator._data_exists() is False def test_data_exists_false_when_experiment_empty(): set_experiment(DEMO_EXPERIMENT_NAME) generator = TracesDemoGenerator() assert generator._data_exists() is False def test_generate_creates_traces(): generator = TracesDemoGenerator() result = generator.generate() assert isinstance(result, DemoResult) assert result.feature == DemoFeature.TRACES assert len(result.entity_ids) > 0 assert "experiments" in result.navigation_url def test_generate_creates_experiment(): generator = TracesDemoGenerator() generator.generate() experiment = get_experiment_by_name(DEMO_EXPERIMENT_NAME) assert experiment is not None assert experiment.lifecycle_stage == "active" def test_data_exists_true_after_generate(): generator = TracesDemoGenerator() assert generator._data_exists() is False generator.generate() assert generator._data_exists() is True def test_delete_demo_removes_traces(): generator = TracesDemoGenerator() generator.generate() assert generator._data_exists() is True generator.delete_demo() assert generator._data_exists() is False def test_traces_have_expected_structure(generated_traces): assert len(generated_traces) > 0 all_span_names = set() for trace in generated_traces: all_span_names.update(span.name for span in trace.data.spans) assert "rag_pipeline" in all_span_names assert "agent" in all_span_names assert "chat_agent" in all_span_names assert "prompt_chain" in all_span_names assert "render_prompt" in all_span_names assert "embed_query" in all_span_names assert "retrieve_docs" in all_span_names assert "chat.completions.create" in all_span_names # OpenAI assert "messages.create" in all_span_names # Anthropic assert "generate_content" in all_span_names # Google def test_traces_have_version_metadata(generated_traces): v1_traces = [t for t in generated_traces if t.info.trace_metadata.get(DEMO_VERSION_TAG) == "v1"] v2_traces = [t for t in generated_traces if t.info.trace_metadata.get(DEMO_VERSION_TAG) == "v2"] # 2 RAG + 2 agent + 6 prompt + 4 multimodal + 7 session = 21 per version assert len(v1_traces) == 21 assert len(v2_traces) == 21 assert len(generated_traces) == 42 def test_traces_have_type_metadata(generated_traces): rag_traces = [ t for t in generated_traces if t.info.trace_metadata.get(DEMO_TRACE_TYPE_TAG) == "rag" ] agent_traces = [ t for t in generated_traces if t.info.trace_metadata.get(DEMO_TRACE_TYPE_TAG) == "agent" ] prompt_traces = [ t for t in generated_traces if t.info.trace_metadata.get(DEMO_TRACE_TYPE_TAG) == "prompt" ] session_traces = [ t for t in generated_traces if t.info.trace_metadata.get(DEMO_TRACE_TYPE_TAG) == "session" ] # 2 RAG per version = 4 total # 2 agent per version = 4 total # 6 prompt per version = 12 total # 7 session per version = 14 total assert len(rag_traces) == 4 assert len(agent_traces) == 4 assert len(prompt_traces) == 12 assert len(session_traces) == 14 def test_is_generated_checks_version(traces_generator): traces_generator.generate() traces_generator.store_version() assert traces_generator.is_generated() is True TracesDemoGenerator.version = 99 assert traces_generator.is_generated() is False def _is_chat_message(obj): """ chat-utils/openai.ts has a normalizeOpenAIChatInput function that asserts a chat-renderable message for inputs """ return isinstance(obj, dict) and "role" in obj and "content" in obj def _has_openai_choices_shape(outputs): """ ModalTraceExplorer.utils.tsx has a fallback which tries to normalise responses if they are OpenAI-shaped """ if not isinstance(outputs, dict): return False choices = outputs.get("choices") if not isinstance(choices, list) or not choices: return False return all(_is_chat_message(c.get("message")) for c in choices) def test_root_span_inputs_are_chat_renderable(generated_traces): for trace in generated_traces: root = next(s for s in trace.data.spans if s.parent_id is None) inputs = root.inputs assert isinstance(inputs, dict), f"Root span {root.name} inputs is not a dict" messages = inputs.get("messages") assert isinstance(messages, list), f"Root span {root.name} inputs missing 'messages' list" assert messages, f"Root span {root.name} inputs has empty 'messages' list" assert all(_is_chat_message(m) for m in messages), ( f"Root span {root.name} has malformed message in inputs" ) def test_llm_span_outputs_are_chat_renderable(generated_traces): for trace in generated_traces: for span in trace.data.spans: if span.span_type != SpanType.LLM: continue assert _has_openai_choices_shape(span.outputs), ( f"LLM span {span.name} in trace {trace.info.trace_id} " f"does not have OpenAI choices output shape" ) def test_root_span_outputs_are_chat_renderable(generated_traces): for trace in generated_traces: # Multimodal traces use the OpenAI Images / Audio API response shapes, # not ChatCompletions; both render in the UI but via different normalizers. if trace.info.trace_metadata.get(DEMO_TRACE_TYPE_TAG) == "multimodal": continue root = next(s for s in trace.data.spans if s.parent_id is None) assert _has_openai_choices_shape(root.outputs), ( f"Root span {root.name} does not have OpenAI choices output shape" ) def test_trace_with_tools_has_react_shape(generated_traces): for trace in generated_traces: root = next(s for s in trace.data.spans if s.parent_id is None) children = [s for s in trace.data.spans if s.parent_id == root.span_id] tool_count = sum(1 for s in children if s.span_type == SpanType.TOOL) if tool_count == 0: continue assert len(children) == 2 * tool_count + 1 ordered = sorted(children, key=lambda s: s.start_time_ns) # Assert ordering now expected = [SpanType.LLM, SpanType.TOOL] * tool_count + [SpanType.LLM] assert [s.span_type for s in ordered] == expected def test_final_llm_span_emits_content(generated_traces): for trace in generated_traces: llm_spans = [s for s in trace.data.spans if s.span_type == SpanType.LLM] if not llm_spans: continue last_llm = max(llm_spans, key=lambda s: s.start_time_ns) choices = last_llm.outputs.get("choices", []) content = choices[0].get("message", {}).get("content") if choices else None assert isinstance(content, str) assert content def test_span_name_matches_provider(generated_traces): for trace in generated_traces: for span in trace.data.spans: # Multimodal traces use SpanType.CHAT_MODEL and are skipped by the LLM check below. if span.span_type != SpanType.LLM: continue provider = span.attributes.get(SpanAttributeKey.MODEL_PROVIDER) assert span.name == _PROVIDER_TO_LLM_SPAN_NAME[provider] def test_session_turns_thread_prior_history(generated_traces): session_traces = [ t for t in generated_traces if t.info.trace_metadata.get(DEMO_TRACE_TYPE_TAG) == "session" ] by_session: dict[str, list[object]] = {} for trace in session_traces: sid = trace.info.trace_metadata.get(TraceMetadataKey.TRACE_SESSION) by_session.setdefault(sid, []).append(trace) multi_turn_sessions = [ts for ts in by_session.values() if len(ts) > 1] assert multi_turn_sessions, "expected at least one session with multiple turns" for turns in multi_turn_sessions: turns.sort(key=lambda t: int(t.info.trace_metadata[DEMO_SESSION_TURN_TAG])) prior_queries: list[str] = [] for turn in turns: root = next(s for s in turn.data.spans if s.parent_id is None) first_llm = min( (s for s in turn.data.spans if s.parent_id == root.span_id), key=lambda s: s.start_time_ns, ) user_messages = [m for m in first_llm.inputs["messages"] if m.get("role") == "user"] assert [m["content"] for m in user_messages] == [ *prior_queries, root.inputs["messages"][0]["content"], ] prior_queries.append(root.inputs["messages"][0]["content"])