wehub-resource-sync
132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
import random
|
|
|
|
from _data import generate_trace_data
|
|
from pytest_benchmark.fixture import BenchmarkFixture
|
|
|
|
import mlflow
|
|
from mlflow.entities.span import SpanType
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
|
|
DEFAULT_SPANS = 100
|
|
INGEST_ROUNDS = 20
|
|
INGEST_WARMUP = 3
|
|
|
|
|
|
def test_ingest(benchmark: BenchmarkFixture, store: SqlAlchemyStore, experiment_id: str) -> None:
|
|
rng = random.Random(42)
|
|
|
|
def setup():
|
|
ti, sp = generate_trace_data(experiment_id, DEFAULT_SPANS, rng)
|
|
return (ti, sp), {}
|
|
|
|
def do(ti, sp):
|
|
store.start_trace(ti)
|
|
store.log_spans(experiment_id, sp)
|
|
|
|
benchmark.pedantic(
|
|
do, setup=setup, iterations=1, rounds=INGEST_ROUNDS, warmup_rounds=INGEST_WARMUP
|
|
)
|
|
|
|
|
|
def test_search_by_tag(
|
|
benchmark: BenchmarkFixture,
|
|
store: SqlAlchemyStore,
|
|
experiment_id: str,
|
|
seeded: list[str],
|
|
) -> None:
|
|
benchmark(
|
|
store.search_traces,
|
|
locations=[experiment_id],
|
|
max_results=100,
|
|
filter_string="tag.env = 'prod'",
|
|
)
|
|
|
|
|
|
def test_search_by_state(
|
|
benchmark: BenchmarkFixture,
|
|
store: SqlAlchemyStore,
|
|
experiment_id: str,
|
|
seeded: list[str],
|
|
) -> None:
|
|
benchmark(
|
|
store.search_traces,
|
|
locations=[experiment_id],
|
|
max_results=100,
|
|
filter_string="status = 'ERROR'",
|
|
)
|
|
|
|
|
|
def test_search_by_name_like(
|
|
benchmark: BenchmarkFixture,
|
|
store: SqlAlchemyStore,
|
|
experiment_id: str,
|
|
seeded: list[str],
|
|
) -> None:
|
|
benchmark(
|
|
store.search_traces,
|
|
locations=[experiment_id],
|
|
max_results=100,
|
|
filter_string="name LIKE 'rag_pipeline%'",
|
|
)
|
|
|
|
|
|
def test_search_by_timestamp(
|
|
benchmark: BenchmarkFixture,
|
|
store: SqlAlchemyStore,
|
|
experiment_id: str,
|
|
seeded: list[str],
|
|
) -> None:
|
|
benchmark(
|
|
store.search_traces,
|
|
locations=[experiment_id],
|
|
max_results=100,
|
|
filter_string="timestamp > 0",
|
|
order_by=["timestamp DESC"],
|
|
)
|
|
|
|
|
|
def _run_agent_workflow(num_tools: int, num_docs: int, query: str) -> None:
|
|
with mlflow.start_span(name="agent_run", span_type=SpanType.AGENT) as root:
|
|
root.set_inputs({"query": query})
|
|
|
|
with mlflow.start_span(name="retrieve", span_type=SpanType.RETRIEVER) as retr:
|
|
retr.set_inputs({"query": query})
|
|
docs = [
|
|
{"id": f"doc_{i}", "score": 0.9 - i * 0.01, "text": f"doc text {i} " * 10}
|
|
for i in range(num_docs)
|
|
]
|
|
retr.set_outputs({"documents": docs})
|
|
|
|
with mlflow.start_span(name="plan", span_type=SpanType.CHAIN) as planner:
|
|
planner.set_inputs({"query": query, "num_docs": len(docs)})
|
|
steps = [f"step_{i}" for i in range(num_tools)]
|
|
planner.set_outputs({"steps": steps})
|
|
|
|
tool_results = []
|
|
for step in steps:
|
|
with mlflow.start_span(name=f"tool:{step}", span_type=SpanType.TOOL) as tool:
|
|
tool.set_inputs({"step": step})
|
|
result = {"step": step, "status": "ok", "value": len(step)}
|
|
tool.set_outputs(result)
|
|
tool_results.append(result)
|
|
|
|
with mlflow.start_span(name="summarize", span_type=SpanType.LLM) as summ:
|
|
summ.set_inputs({"query": query, "tool_results": tool_results})
|
|
response = f"Answer to {query!r} using {num_docs} docs and {num_tools} tool calls."
|
|
summ.set_outputs({"response": response})
|
|
summ.set_attribute("model", "gpt-test")
|
|
summ.set_attribute("usage.input_tokens", 1234)
|
|
summ.set_attribute("usage.output_tokens", 567)
|
|
|
|
root.set_outputs({"response": response})
|
|
|
|
|
|
def test_e2e_agent(benchmark: BenchmarkFixture, e2e_setup: None) -> None:
|
|
counter = [0]
|
|
|
|
def do():
|
|
_run_agent_workflow(num_tools=20, num_docs=20, query=f"q-{counter[0]}")
|
|
counter[0] += 1
|
|
|
|
benchmark(do)
|