Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: extend component run info #8436

Closed
wants to merge 6 commits into from
Closed

feat: extend component run info #8436

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e3680d0
feat: extend component run info
tstadel Oct 2, 2024
2dd43bf
fix tests
tstadel Oct 4, 2024
05a2f59
add tests for lengths
tstadel Oct 4, 2024
0526ae6
extend logging test
tstadel Oct 4, 2024
5c3e11e
Merge branch 'main' into feat/extend_component_run_info
tstadel Oct 4, 2024
7fdc3ec
fix merge
tstadel Oct 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 47 additions & 17 deletions haystack/core/pipeline/pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,31 +39,46 @@ def _run_component(self, name: str, inputs: Dict[str, Any]) -> Dict[str, Any]:
:return: The output of the Component.
"""
instance: Component = self.graph.nodes[name]["instance"]
component_type = instance.__class__.__name__
input_types = {k: type(v).__name__ for k, v in inputs.items()}
input_lengths = {k: len(v) for k, v in inputs.items() if isinstance(v, (list, tuple, set))}
input_spec = {
key: {
"type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
"senders": value.senders,
}
for key, value in instance.__haystack_input__._sockets_dict.items() # type: ignore
}
output_spec = {
key: {
"type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
"receivers": value.receivers,
}
for key, value in instance.__haystack_output__._sockets_dict.items() # type: ignore
}

with tracing.tracer.trace(
"haystack.component.run",
tags={
"haystack.component.name": name,
"haystack.component.type": instance.__class__.__name__,
"haystack.component.input_types": {k: type(v).__name__ for k, v in inputs.items()},
"haystack.component.input_spec": {
key: {
"type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
"senders": value.senders,
}
for key, value in instance.__haystack_input__._sockets_dict.items() # type: ignore
},
"haystack.component.output_spec": {
key: {
"type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
"receivers": value.receivers,
}
for key, value in instance.__haystack_output__._sockets_dict.items() # type: ignore
},
"haystack.component.type": component_type,
"haystack.component.input_types": input_types,
"haystack.component.input_lengths": input_lengths,
"haystack.component.input_spec": input_spec,
"haystack.component.output_spec": output_spec,
},
) as span:
span.set_content_tag("haystack.component.input", inputs)
logger.info("Running component {component_name}", component_name=name)
logger.info(
"Running component {component_name}",
component_name=name,
extra={
"component_name": name,
"component_type": component_type,
"input_types": input_types,
"input_lengths": input_lengths,
},
)
res: Dict[str, Any] = instance.run(**inputs)
self.graph.nodes[name]["visits"] += 1

Expand All @@ -79,8 +94,23 @@ def _run_component(self, name: str, inputs: Dict[str, Any]) -> Dict[str, Any]:
f"Component '{name}' didn't return a dictionary. "
"Components must always return dictionaries: check the documentation."
)
output_types = {k: type(v).__name__ for k, v in res.items()}
output_lengths = {k: len(v) for k, v in res.items() if isinstance(v, (list, tuple, set))}

span.set_tag("haystack.component.visits", self.graph.nodes[name]["visits"])
span.set_tag("haystack.component.output_types", output_types)
span.set_tag("haystack.component.output_lengths", output_lengths)
span.set_content_tag("haystack.component.output", res)
logger.info(
"Component {component_name} finished",
component_name=name,
extra={
"component_name": name,
"component_type": component_type,
"output_types": output_types,
"output_lengths": output_lengths,
},
)

return res

Expand Down
13 changes: 12 additions & 1 deletion test/core/pipeline/test_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1094,13 +1094,24 @@ def test__run_component(self, spying_tracer, caplog):
assert span.tags == {
"haystack.component.name": "document_builder",
"haystack.component.type": "DocumentBuilder",
"haystack.component.input_lengths": {},
"haystack.component.input_types": {"text": "str"},
"haystack.component.input_spec": {"text": {"type": "str", "senders": ["sentence_builder"]}},
"haystack.component.output_spec": {"doc": {"type": "Document", "receivers": ["document_cleaner"]}},
"haystack.component.output_lengths": {},
"haystack.component.output_types": {"doc": "Document"},
"haystack.component.visits": 1,
}

assert caplog.messages == ["Running component document_builder"]
assert caplog.messages == ["Running component document_builder", "Component document_builder finished"]
assert caplog.records[0].component_name == "document_builder"
assert caplog.records[0].component_type == "DocumentBuilder"
assert caplog.records[0].input_lengths == {}
assert caplog.records[0].input_types == {"text": "str"}
assert caplog.records[1].component_name == "document_builder"
assert caplog.records[1].component_type == "DocumentBuilder"
assert caplog.records[1].output_lengths == {}
assert caplog.records[1].output_types == {"doc": "Document"}

def test__run_component_with_variadic_input(self):
document_joiner = component_class("DocumentJoiner", input_types={"docs": Variadic[Document]})()
Expand Down
77 changes: 76 additions & 1 deletion test/core/pipeline/test_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
#
# SPDX-License-Identifier: Apache-2.0
from test.tracing.utils import SpyingSpan, SpyingTracer
from typing import Optional
from typing import List, Optional
from unittest.mock import ANY

import pytest
Expand All @@ -23,6 +23,18 @@ def run(self, word: Optional[str]): # use optional to spice up the typing tags
return {"output": f"Hello, {word}!"}


@component
class HelloList:
@component.output_types(output=List[str])
def run(self, words: Optional[List[str]]): # use optional to spice up the typing tags
"""
Takes a string in input and returns "Hello, <string>!"
in output.
"""
outputs = [f"Hello, {word}!" for word in words]
return {"output": outputs}


@pytest.fixture()
def pipeline() -> Pipeline:
pipeline = Pipeline()
Expand All @@ -32,6 +44,13 @@ def pipeline() -> Pipeline:
return pipeline


@pytest.fixture()
def pipeline_list_input() -> Pipeline:
pipeline = Pipeline()
pipeline.add_component("hello", HelloList())
return pipeline


class TestTracing:
def test_with_enabled_tracing(self, pipeline: Pipeline, spying_tracer: SpyingTracer) -> None:
pipeline.run(data={"word": "world"})
Expand All @@ -55,8 +74,11 @@ def test_with_enabled_tracing(self, pipeline: Pipeline, spying_tracer: SpyingTra
tags={
"haystack.component.name": "hello",
"haystack.component.type": "Hello",
"haystack.component.input_lengths": {},
"haystack.component.input_types": {"word": "str"},
"haystack.component.input_spec": {"word": {"type": ANY, "senders": []}},
"haystack.component.output_lengths": {},
"haystack.component.output_types": {"output": "str"},
"haystack.component.output_spec": {"output": {"type": "str", "receivers": ["hello2"]}},
"haystack.component.visits": 1,
},
Expand All @@ -68,8 +90,11 @@ def test_with_enabled_tracing(self, pipeline: Pipeline, spying_tracer: SpyingTra
tags={
"haystack.component.name": "hello2",
"haystack.component.type": "Hello",
"haystack.component.input_lengths": {},
"haystack.component.input_types": {"word": "str"},
"haystack.component.input_spec": {"word": {"type": ANY, "senders": ["hello"]}},
"haystack.component.output_lengths": {},
"haystack.component.output_types": {"output": "str"},
"haystack.component.output_spec": {"output": {"type": "str", "receivers": []}},
"haystack.component.visits": 1,
},
Expand Down Expand Up @@ -111,8 +136,11 @@ def test_with_enabled_content_tracing(
tags={
"haystack.component.name": "hello",
"haystack.component.type": "Hello",
"haystack.component.input_lengths": {},
"haystack.component.input_types": {"word": "str"},
"haystack.component.input_spec": {"word": {"type": ANY, "senders": []}},
"haystack.component.output_lengths": {},
"haystack.component.output_types": {"output": "str"},
"haystack.component.output_spec": {"output": {"type": "str", "receivers": ["hello2"]}},
"haystack.component.input": {"word": "world"},
"haystack.component.visits": 1,
Expand All @@ -126,8 +154,11 @@ def test_with_enabled_content_tracing(
tags={
"haystack.component.name": "hello2",
"haystack.component.type": "Hello",
"haystack.component.input_lengths": {},
"haystack.component.input_types": {"word": "str"},
"haystack.component.input_spec": {"word": {"type": ANY, "senders": ["hello"]}},
"haystack.component.output_lengths": {},
"haystack.component.output_types": {"output": "str"},
"haystack.component.output_spec": {"output": {"type": "str", "receivers": []}},
"haystack.component.input": {"word": "Hello, world!"},
"haystack.component.visits": 1,
Expand All @@ -137,3 +168,47 @@ def test_with_enabled_content_tracing(
span_id=ANY,
),
]

def test_with_enabled_tracing_input_lengths(
self, pipeline_list_input: Pipeline, spying_tracer: SpyingTracer
) -> None:
pipeline_list_input.run(data={"words": ["happy", "world"]})

assert len(spying_tracer.spans) == 2

assert spying_tracer.spans == [
SpyingSpan(
operation_name="haystack.pipeline.run",
tags={
"haystack.pipeline.input_data": {"hello": {"words": ["happy", "world"]}},
"haystack.pipeline.output_data": {"hello": {"output": ["Hello, happy!", "Hello, world!"]}},
"haystack.pipeline.metadata": {},
"haystack.pipeline.max_runs_per_component": 100,
},
trace_id=ANY,
span_id=ANY,
),
SpyingSpan(
operation_name="haystack.component.run",
tags={
"haystack.component.name": "hello",
"haystack.component.type": "HelloList",
"haystack.component.input_lengths": {"words": 2},
"haystack.component.input_types": {"words": "list"},
"haystack.component.input_spec": {"words": {"type": ANY, "senders": []}},
"haystack.component.output_lengths": {"output": 2},
"haystack.component.output_types": {"output": "list"},
"haystack.component.output_spec": {"output": {"type": "typing.List[str]", "receivers": []}},
"haystack.component.visits": 1,
},
trace_id=ANY,
span_id=ANY,
),
]

# We need to check the type of the input_spec because it can be rendered differently
# depending on the Python version 🫠
assert spying_tracer.spans[1].tags["haystack.component.input_spec"]["words"]["type"] in [
"typing.Union[typing.List[str], NoneType]",
"typing.Optional[typing.List[str]]",
]
Loading