wehub-resource-sync
214 lines
6.5 KiB
Python
214 lines
6.5 KiB
Python
from unittest.mock import patch
|
|
|
|
from haystack import Document, Pipeline, component
|
|
from haystack.components.rankers import LostInTheMiddleRanker
|
|
from haystack.components.retrievers import InMemoryBM25Retriever
|
|
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
|
|
|
import mlflow
|
|
from mlflow.entities import SpanType
|
|
from mlflow.environment_variables import MLFLOW_USE_DEFAULT_TRACER_PROVIDER
|
|
from mlflow.tracing.constant import SpanAttributeKey
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
from tests.tracing.helper import get_traces
|
|
|
|
|
|
@component
|
|
class Add:
|
|
def run(self, a: int, b: int):
|
|
return {"sum": a + b}
|
|
|
|
|
|
@component
|
|
class Multiply:
|
|
def run(self, value: int, factor: int):
|
|
return {"product": value * factor}
|
|
|
|
|
|
def test_haystack_autolog_single_trace():
|
|
mlflow.haystack.autolog()
|
|
|
|
pipe = Pipeline()
|
|
pipe.add_component("adder", Add())
|
|
pipe.run({"adder": {"a": 1, "b": 2}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
spans = traces[0].data.spans
|
|
assert spans[0].span_type == SpanType.CHAIN
|
|
assert spans[0].name == "haystack.pipeline.run"
|
|
assert spans[0].inputs == {"adder": {"a": 1, "b": 2}}
|
|
assert spans[0].outputs == {"adder": {"sum": 3}}
|
|
assert spans[1].span_type == SpanType.TOOL
|
|
assert spans[1].name == "Add"
|
|
assert spans[1].inputs == {"a": 1, "b": 2}
|
|
assert spans[1].outputs == {"sum": 3}
|
|
|
|
mlflow.haystack.autolog(disable=True)
|
|
pipe.run({"adder": {"a": 3, "b": 4}})
|
|
assert len(get_traces()) == 1
|
|
|
|
|
|
def test_pipeline_with_multiple_components_single_trace():
|
|
mlflow.haystack.autolog()
|
|
|
|
pipe = Pipeline()
|
|
pipe.add_component("adder", Add())
|
|
pipe.add_component("multiplier", Multiply())
|
|
|
|
pipe.run({"adder": {"a": 1, "b": 2}, "multiplier": {"value": 3, "factor": 4}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
spans = traces[0].data.spans
|
|
assert spans[0].span_type == SpanType.CHAIN
|
|
assert spans[0].name == "haystack.pipeline.run"
|
|
assert spans[1].span_type == SpanType.TOOL
|
|
assert spans[2].span_type == SpanType.TOOL
|
|
assert spans[1].name == "Add"
|
|
assert spans[2].name == "Multiply"
|
|
assert spans[1].inputs == {"a": 1, "b": 2}
|
|
assert spans[1].outputs == {"sum": 3}
|
|
assert spans[2].inputs == {"value": 3, "factor": 4}
|
|
assert spans[2].outputs == {"product": 12}
|
|
|
|
mlflow.haystack.autolog(disable=True)
|
|
pipe.run({"adder": {"a": 1, "b": 2}, "multiplier": {"value": 3, "factor": 4}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
|
|
|
|
def test_token_usage_parsed_for_llm_component(mock_litellm_cost):
|
|
mlflow.haystack.autolog()
|
|
|
|
@component
|
|
class MyLLM:
|
|
def run(self, prompt: str, model: str):
|
|
return {}
|
|
|
|
pipe = Pipeline()
|
|
pipe.add_component("my_llm", MyLLM())
|
|
|
|
output = {
|
|
"replies": [
|
|
{
|
|
"content": [{"text": "hi"}],
|
|
"meta": {"usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}},
|
|
}
|
|
]
|
|
}
|
|
|
|
with patch.object(MyLLM, "run", return_value=output):
|
|
pipe.run({"my_llm": {"prompt": "hello", "model": "gpt-4"}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
span = traces[0].data.spans[1]
|
|
assert span.span_type == SpanType.LLM
|
|
assert span.name == "MyLLM"
|
|
assert span.attributes[SpanAttributeKey.CHAT_USAGE] == {
|
|
"input_tokens": 1,
|
|
"output_tokens": 2,
|
|
"total_tokens": 3,
|
|
}
|
|
assert span.model_name == "gpt-4"
|
|
if not IS_TRACING_SDK_ONLY:
|
|
# Verify cost is calculated (1 input token * 1.0 + 2 output tokens * 2.0)
|
|
assert span.llm_cost == {
|
|
"input_cost": 1.0,
|
|
"output_cost": 4.0,
|
|
"total_cost": 5.0,
|
|
}
|
|
|
|
mlflow.haystack.autolog(disable=True)
|
|
|
|
traces = get_traces()
|
|
with patch.object(MyLLM, "run", return_value=output):
|
|
pipe.run({"my_llm": {"prompt": "hello", "model": "gpt-4"}})
|
|
|
|
assert len(traces) == 1
|
|
|
|
|
|
def test_autolog_disable():
|
|
mlflow.haystack.autolog()
|
|
|
|
pipe1 = Pipeline()
|
|
pipe1.add_component("adder", Add())
|
|
pipe1.run({"adder": {"a": 1, "b": 2}})
|
|
assert len(get_traces()) == 1
|
|
|
|
mlflow.haystack.autolog(disable=True)
|
|
pipe2 = Pipeline()
|
|
pipe2.add_component("adder", Add())
|
|
pipe2.run({"adder": {"a": 2, "b": 3}})
|
|
assert len(get_traces()) == 1
|
|
|
|
|
|
def test_in_memory_retriever_component_traced():
|
|
mlflow.set_experiment("haystack_retriever")
|
|
mlflow.haystack.autolog()
|
|
|
|
store = InMemoryDocumentStore()
|
|
store.write_documents([Document(content="foo")])
|
|
pipe = Pipeline()
|
|
pipe.add_component("retriever", InMemoryBM25Retriever(document_store=store))
|
|
pipe.run({"retriever": {"query": "foo"}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
span = traces[0].data.spans[1]
|
|
assert span.span_type == SpanType.RETRIEVER
|
|
assert span.name == "InMemoryBM25Retriever"
|
|
assert span.outputs["documents"][0]["content"] == "foo"
|
|
|
|
|
|
def test_multiple_components_in_pipeline_reranker():
|
|
mlflow.haystack.autolog()
|
|
|
|
pipe = Pipeline()
|
|
store = InMemoryDocumentStore()
|
|
store.write_documents([Document(content="foo")])
|
|
|
|
pipe.add_component("retriever", InMemoryBM25Retriever(document_store=store))
|
|
pipe.add_component("reranker", LostInTheMiddleRanker())
|
|
|
|
pipe.connect("retriever.documents", "reranker.documents")
|
|
|
|
pipe.run({"retriever": {"query": "foo"}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
spans = traces[0].data.spans
|
|
assert spans[0].span_type == SpanType.CHAIN
|
|
assert spans[0].name == "haystack.pipeline.run"
|
|
assert spans[1].name == "InMemoryBM25Retriever"
|
|
assert spans[2].name == "LostInTheMiddleRanker"
|
|
assert spans[1].span_type == SpanType.RETRIEVER
|
|
assert spans[2].span_type == SpanType.RERANKER
|
|
assert spans[1].inputs["query"] == "foo"
|
|
assert spans[2].inputs["documents"][0]["content"] == "foo"
|
|
|
|
mlflow.haystack.autolog(disable=True)
|
|
pipe.run({"retriever": {"query": "foo"}})
|
|
assert len(get_traces()) == 1
|
|
|
|
|
|
def test_haystack_autolog_shared_provider_no_recursion(monkeypatch):
|
|
# Verify haystack.autolog() works with shared tracer provider (no RecursionError)
|
|
monkeypatch.setenv(MLFLOW_USE_DEFAULT_TRACER_PROVIDER.name, "false")
|
|
|
|
mlflow.haystack.autolog()
|
|
|
|
pipe = Pipeline()
|
|
pipe.add_component("adder", Add())
|
|
pipe.run({"adder": {"a": 1, "b": 2}})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
spans = traces[0].data.spans
|
|
assert spans[0].span_type == SpanType.CHAIN
|
|
assert spans[0].inputs == {"adder": {"a": 1, "b": 2}}
|
|
assert spans[0].outputs == {"adder": {"sum": 3}}
|