wehub-resource-sync
77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.entities.span import SpanType
|
|
from mlflow.entities.trace import Trace
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.genai.scorers.base import Scorer
|
|
|
|
|
|
class _TestScorer(Scorer):
|
|
"""Minimal Scorer subclass for use in discovery pipeline tests."""
|
|
|
|
def __call__(self, **kwargs):
|
|
return True
|
|
|
|
|
|
def _create_trace(
|
|
request="What is MLflow?",
|
|
response="MLflow is an ML platform.",
|
|
session_id=None,
|
|
span_type=SpanType.CHAIN,
|
|
error_span=False,
|
|
execution_duration_ms="default",
|
|
):
|
|
@mlflow.trace(name="agent", span_type=span_type)
|
|
def _run(question):
|
|
if session_id:
|
|
mlflow.update_current_trace(
|
|
metadata={"mlflow.trace.session": session_id},
|
|
)
|
|
with mlflow.start_span(name="llm_call", span_type=SpanType.LLM) as child:
|
|
child.set_inputs({"prompt": question})
|
|
child.set_outputs({"response": response})
|
|
if error_span:
|
|
with mlflow.start_span(name="tool_call", span_type=SpanType.TOOL) as tool:
|
|
tool.set_inputs({"action": "fetch"})
|
|
tool.record_exception("Connection failed")
|
|
return response
|
|
|
|
_run(request)
|
|
trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
|
|
|
|
# Override execution_duration if explicitly requested
|
|
if execution_duration_ms != "default":
|
|
trace_info = TraceInfo(
|
|
trace_id=trace.info.trace_id,
|
|
trace_location=trace.info.trace_location,
|
|
request_time=trace.info.timestamp_ms,
|
|
execution_duration=execution_duration_ms,
|
|
state=trace.info.state,
|
|
trace_metadata=trace.info.trace_metadata,
|
|
tags=trace.info.tags,
|
|
)
|
|
trace = Trace(trace_info, trace.data)
|
|
|
|
return trace
|
|
|
|
|
|
@pytest.fixture
|
|
def make_trace():
|
|
def _make(
|
|
request="What is MLflow?",
|
|
response="MLflow is an ML platform.",
|
|
session_id=None,
|
|
error_span=False,
|
|
execution_duration_ms="default",
|
|
):
|
|
return _create_trace(
|
|
request=request,
|
|
response=response,
|
|
session_id=session_id,
|
|
error_span=error_span,
|
|
execution_duration_ms=execution_duration_ms,
|
|
)
|
|
|
|
return _make
|