diff --git a/packages/traceloop-sdk/tests/test_workflow_context.py b/packages/traceloop-sdk/tests/test_workflow_context.py new file mode 100644 index 0000000000..f763099fc2 --- /dev/null +++ b/packages/traceloop-sdk/tests/test_workflow_context.py @@ -0,0 +1,237 @@ +"""Regression tests for workflow context scoping.""" + +import pytest +from opentelemetry import trace +from opentelemetry.semconv_ai import SpanAttributes + +from traceloop.sdk.decorators import task, workflow + + +def _finished_spans_by_name(exporter): + return {span.name: span for span in exporter.get_finished_spans()} + + +def _assert_without_workflow(span): + assert SpanAttributes.TRACELOOP_WORKFLOW_NAME not in span.attributes + assert span.parent is None + + +def test_completed_workflow_does_not_tag_following_task(exporter): + @task(name="inside") + def inside_task(): + pass + + @workflow(name="first") + def first_workflow(): + inside_task() + + @task(name="outside") + def outside_task(): + pass + + first_workflow() + outside_task() + + spans = _finished_spans_by_name(exporter) + assert spans["inside.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "first" + _assert_without_workflow(spans["outside.task"]) + + +def test_nested_workflow_restores_enclosing_workflow(exporter): + @task(name="inner_task") + def inner_task(): + pass + + @workflow(name="inner") + def inner_workflow(): + inner_task() + + @task(name="after_inner") + def task_after_inner(): + pass + + @workflow(name="outer") + def outer_workflow(): + inner_workflow() + task_after_inner() + + outer_workflow() + + spans = _finished_spans_by_name(exporter) + assert spans["inner_task.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "inner" + assert spans["after_inner.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "outer" + + +def test_failed_workflow_does_not_tag_following_task(exporter): + @workflow(name="failing") + def failing_workflow(): + raise RuntimeError("expected") + + @task(name="after_failure") + def task_after_failure(): + pass + + with pytest.raises(RuntimeError, match="expected"): + failing_workflow() + task_after_failure() + + spans = _finished_spans_by_name(exporter) + _assert_without_workflow(spans["after_failure.task"]) + + +@pytest.mark.asyncio +async def test_completed_async_workflow_does_not_tag_following_task(exporter): + @workflow(name="async_workflow") + async def async_workflow(): + pass + + @task(name="after_async") + async def task_after_async(): + pass + + await async_workflow() + await task_after_async() + + spans = _finished_spans_by_name(exporter) + _assert_without_workflow(spans["after_async.task"]) + + +def test_generator_workflow_scopes_context_to_iteration(exporter): + @task(name="inside_stream") + def task_inside_stream(): + pass + + @task(name="closing_stream") + def task_while_closing_stream(): + pass + + @workflow(name="stream") + def stream_workflow(): + try: + task_inside_stream() + yield 1 + finally: + task_while_closing_stream() + + @task(name="before_stream") + def task_before_stream(): + pass + + @task(name="after_stream") + def task_after_stream(): + pass + + @task(name="between_stream_items") + def task_between_stream_items(): + pass + + stream = stream_workflow() + task_before_stream() + assert "stream.workflow" not in _finished_spans_by_name(exporter) + assert next(stream) == 1 + task_between_stream_items() + stream.close() + task_after_stream() + + spans = _finished_spans_by_name(exporter) + _assert_without_workflow(spans["before_stream.task"]) + assert spans["inside_stream.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "stream" + assert spans["closing_stream.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "stream" + assert spans["closing_stream.task"].parent.span_id == spans["stream.workflow"].context.span_id + _assert_without_workflow(spans["between_stream_items.task"]) + _assert_without_workflow(spans["after_stream.task"]) + + +@pytest.mark.asyncio +async def test_async_generator_workflow_restores_context(exporter): + @task(name="inside_async_stream") + async def task_inside_async_stream(): + pass + + @task(name="closing_async_stream") + async def task_while_closing_async_stream(): + pass + + @workflow(name="async_stream") + async def async_stream_workflow(): + try: + await task_inside_async_stream() + yield 1 + finally: + await task_while_closing_async_stream() + + @task(name="after_async_stream") + async def task_after_async_stream(): + pass + + @task(name="between_async_stream_items") + async def task_between_async_stream_items(): + pass + + stream = async_stream_workflow() + result = await anext(stream) + await task_between_async_stream_items() + await stream.aclose() + await task_after_async_stream() + + spans = _finished_spans_by_name(exporter) + assert result == 1 + assert spans["inside_async_stream.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "async_stream" + assert spans["closing_async_stream.task"].attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "async_stream" + assert spans["closing_async_stream.task"].parent.span_id == spans["async_stream.workflow"].context.span_id + _assert_without_workflow(spans["between_async_stream_items.task"]) + _assert_without_workflow(spans["after_async_stream.task"]) + + +def test_generator_preserves_inner_span_across_yield(exporter): + tracer = trace.get_tracer(__name__) + + @workflow(name="stream") + def stream_workflow(): + with tracer.start_as_current_span("long_child"): + yield 1 + with tracer.start_as_current_span("after_yield"): + pass + + stream = stream_workflow() + assert next(stream) == 1 + with pytest.raises(StopIteration): + next(stream) + + spans = _finished_spans_by_name(exporter) + assert spans["after_yield"].parent.span_id == spans["long_child"].context.span_id + + +@pytest.mark.asyncio +async def test_async_generator_preserves_inner_span_across_yield(exporter): + tracer = trace.get_tracer(__name__) + + @workflow(name="async_stream") + async def async_stream_workflow(): + with tracer.start_as_current_span("async_long_child"): + yield 1 + with tracer.start_as_current_span("async_after_yield"): + pass + + stream = async_stream_workflow() + assert await anext(stream) == 1 + with pytest.raises(StopAsyncIteration): + await anext(stream) + + spans = _finished_spans_by_name(exporter) + assert spans["async_after_yield"].parent.span_id == spans["async_long_child"].context.span_id + + +def test_generator_task_propagates_entity_path(exporter): + tracer = trace.get_tracer(__name__) + + @task(name="stream_task") + def stream_task(): + with tracer.start_as_current_span("stream_child"): + pass + yield 1 + + assert list(stream_task()) == [1] + + spans = _finished_spans_by_name(exporter) + assert spans["stream_child"].attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] == "stream_task" diff --git a/packages/traceloop-sdk/traceloop/sdk/decorators/base.py b/packages/traceloop-sdk/traceloop/sdk/decorators/base.py index 04f00e0ad9..88f588a3a7 100644 --- a/packages/traceloop-sdk/traceloop/sdk/decorators/base.py +++ b/packages/traceloop-sdk/traceloop/sdk/decorators/base.py @@ -20,10 +20,9 @@ GEN_AI_TOOL_NAME, ) -from traceloop.sdk.tracing import get_tracer, set_workflow_name, set_agent_name +from traceloop.sdk.tracing import get_tracer, set_agent_name from traceloop.sdk.tracing.tracing import ( TracerWrapper, - set_entity_path, get_chained_entity_path, ) from traceloop.sdk.utils import camel_to_snake @@ -91,34 +90,56 @@ def aentity_class( ) -def _handle_generator(span, res): - # for some reason the SPAN_KEY is not being set in the context of the generator, so we re-set it - context_api.attach(trace.set_span_in_context(span)) +def _handle_generator(span, ctx, res): + generator_ctx = ctx try: - for item in res: + while True: + ctx_token = context_api.attach(generator_ctx) + try: + item = next(res) + except StopIteration: + return + finally: + generator_ctx = context_api.get_current() + context_api.detach(ctx_token) yield item except Exception as e: span.set_status(Status(StatusCode.ERROR, str(e))) span.record_exception(e) raise finally: - span.end() - # Note: we don't detach the context here as this fails in some situations - # https://github.com/open-telemetry/opentelemetry-python/issues/2606 - # This is not a problem since the context will be detached automatically during garbage collection + ctx_token = context_api.attach(generator_ctx) + try: + res.close() + finally: + context_api.detach(ctx_token) + span.end() -async def _ahandle_generator(span, ctx_token, res): +async def _ahandle_generator(span, ctx, res): + generator_ctx = ctx try: - async for part in res: + while True: + ctx_token = context_api.attach(generator_ctx) + try: + part = await anext(res) + except StopAsyncIteration: + return + finally: + generator_ctx = context_api.get_current() + context_api.detach(ctx_token) yield part except Exception as e: span.set_status(Status(StatusCode.ERROR, str(e))) span.record_exception(e) raise finally: - span.end() - context_api.detach(ctx_token) + ctx_token = context_api.attach(generator_ctx) + try: + await res.aclose() + finally: + context_api.detach(ctx_token) + span.end() def _should_send_prompts(): @@ -137,9 +158,7 @@ def _is_async_method(fn): def _setup_span(entity_name, tlp_span_kind, version): """Sets up the OpenTelemetry span and context""" - if tlp_span_kind == TraceloopSpanKindValues.WORKFLOW: - set_workflow_name(entity_name) - elif tlp_span_kind == TraceloopSpanKindValues.AGENT: + if tlp_span_kind == TraceloopSpanKindValues.AGENT: set_agent_name(entity_name) span_name = f"{entity_name}.{tlp_span_kind.value}" @@ -147,17 +166,20 @@ def _setup_span(entity_name, tlp_span_kind, version): with get_tracer() as tracer: span = tracer.start_span(span_name) ctx = trace.set_span_in_context(span) - ctx_token = context_api.attach(ctx) - + if tlp_span_kind == TraceloopSpanKindValues.WORKFLOW: + ctx = context_api.set_value("workflow_name", entity_name, ctx) if tlp_span_kind in [ TraceloopSpanKindValues.TASK, TraceloopSpanKindValues.TOOL, ]: entity_path = get_chained_entity_path(entity_name) - set_entity_path(entity_path) + ctx = context_api.set_value("entity_path", entity_path, ctx) + ctx_token = context_api.attach(ctx) span.set_attribute(SpanAttributes.TRACELOOP_SPAN_KIND, tlp_span_kind.value) span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name) + if tlp_span_kind == TraceloopSpanKindValues.WORKFLOW: + span.set_attribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, entity_name) if tlp_span_kind == TraceloopSpanKindValues.TOOL: span.set_attribute(GEN_AI_TOOL_NAME, entity_name) if version: @@ -224,10 +246,13 @@ async def async_gen_wrap(*args: Any, **kwargs: Any) -> Any: entity_name, tlp_span_kind, version ) _handle_span_input(span, args, kwargs, cls=JSONEncoder) - async for item in _ahandle_generator( - span, ctx_token, fn(*args, **kwargs) - ): - yield item + context_api.detach(ctx_token) + generator = _ahandle_generator(span, ctx, fn(*args, **kwargs)) + try: + async for item in generator: + yield item + finally: + await generator.aclose() return cast(F, async_gen_wrap) else: @@ -254,6 +279,22 @@ async def async_wrap(*args: Any, **kwargs: Any) -> Any: return cast(F, async_wrap) else: + if inspect.isgeneratorfunction(fn): + + @wraps(fn) + def sync_gen_wrap(*args: Any, **kwargs: Any) -> Any: + if not TracerWrapper.verify_initialized(): + yield from fn(*args, **kwargs) + return + + span, ctx, ctx_token = _setup_span( + entity_name, tlp_span_kind, version + ) + _handle_span_input(span, args, kwargs, cls=JSONEncoder) + context_api.detach(ctx_token) + yield from _handle_generator(span, ctx, fn(*args, **kwargs)) + + return cast(F, sync_gen_wrap) @wraps(fn) def sync_wrap(*args: Any, **kwargs: Any) -> Any: @@ -272,7 +313,8 @@ def sync_wrap(*args: Any, **kwargs: Any) -> Any: # span will be ended in the generator if isinstance(res, types.GeneratorType): - return _handle_generator(span, res) + context_api.detach(ctx_token) + return _handle_generator(span, ctx, res) _handle_span_output(span, res, cls=JSONEncoder) _cleanup_span(span, ctx_token)