From 2811db2e07cc88f1bde8016d4ded1f1b60e95335 Mon Sep 17 00:00:00 2001 From: Josh Lebedinsky Date: Fri, 15 Aug 2025 13:57:31 -0400 Subject: [PATCH 1/3] fix: add gentrace.in_experiment to span attributes --- src/gentrace/__init__.py | 4 +++- src/gentrace/lib/constants.py | 2 ++ src/gentrace/lib/eval.py | 4 ++++ src/gentrace/lib/eval_dataset.py | 4 ++++ src/gentrace/lib/span_processor.py | 7 ++++++- 5 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/gentrace/__init__.py b/src/gentrace/__init__.py index b7da1e46..04001b86 100644 --- a/src/gentrace/__init__.py +++ b/src/gentrace/__init__.py @@ -107,6 +107,7 @@ def __getattr__(self, name: str) -> _t.Any: ATTR_GENTRACE_PIPELINE_ID, ATTR_GENTRACE_TEST_CASE_ID, ATTR_GENTRACE_EXPERIMENT_ID, + ATTR_GENTRACE_IN_EXPERIMENT, ATTR_GENTRACE_TEST_CASE_NAME, ATTR_GENTRACE_FN_ARGS_EVENT_NAME, ATTR_GENTRACE_FN_OUTPUT_EVENT_NAME, @@ -160,7 +161,6 @@ def __getattr__(self, name: str) -> _t.Any: "DefaultHttpxClient", "DefaultAsyncHttpxClient", "DefaultAioHttpClient", - # Start custom Gentrace exports "init", "setup", @@ -170,6 +170,8 @@ def __getattr__(self, name: str) -> _t.Any: "eval", "eval_dataset", "TestInput", + "ATTR_GENTRACE_SAMPLE_KEY", + "ATTR_GENTRACE_IN_EXPERIMENT", "ATTR_GENTRACE_PIPELINE_ID", "ATTR_GENTRACE_TEST_CASE_ID", "ATTR_GENTRACE_EXPERIMENT_ID", diff --git a/src/gentrace/lib/constants.py b/src/gentrace/lib/constants.py index 6f96ea2b..7ecff51a 100644 --- a/src/gentrace/lib/constants.py +++ b/src/gentrace/lib/constants.py @@ -12,6 +12,7 @@ ATTR_GENTRACE_TEST_CASE_ID = "gentrace.test_case_id" ATTR_GENTRACE_SAMPLE_KEY = "gentrace.sample" +ATTR_GENTRACE_IN_EXPERIMENT = "gentrace.in_experiment" # Maximum allowed concurrency for eval_dataset MAX_EVAL_DATASET_CONCURRENCY = 100 @@ -25,5 +26,6 @@ "ATTR_GENTRACE_TEST_CASE_ID", "ATTR_GENTRACE_PIPELINE_ID", "ATTR_GENTRACE_SAMPLE_KEY", + "ATTR_GENTRACE_IN_EXPERIMENT", "MAX_EVAL_DATASET_CONCURRENCY", ] diff --git a/src/gentrace/lib/eval.py b/src/gentrace/lib/eval.py index a5ca3e7a..c437b062 100644 --- a/src/gentrace/lib/eval.py +++ b/src/gentrace/lib/eval.py @@ -13,6 +13,7 @@ ANONYMOUS_SPAN_NAME, ATTR_GENTRACE_SAMPLE_KEY, ATTR_GENTRACE_EXPERIMENT_ID, + ATTR_GENTRACE_IN_EXPERIMENT, ATTR_GENTRACE_TEST_CASE_NAME, ATTR_GENTRACE_FN_ARGS_EVENT_NAME, ATTR_GENTRACE_FN_OUTPUT_EVENT_NAME, @@ -90,6 +91,9 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> Any: context_with_modified_baggage = otel_baggage.set_baggage( ATTR_GENTRACE_SAMPLE_KEY, "true", context=current_context ) + context_with_modified_baggage = otel_baggage.set_baggage( + ATTR_GENTRACE_IN_EXPERIMENT, "true", context=current_context + ) token = otel_context.attach(context_with_modified_baggage) try: diff --git a/src/gentrace/lib/eval_dataset.py b/src/gentrace/lib/eval_dataset.py index dff60f2a..80588707 100644 --- a/src/gentrace/lib/eval_dataset.py +++ b/src/gentrace/lib/eval_dataset.py @@ -43,6 +43,7 @@ ATTR_GENTRACE_SAMPLE_KEY, ATTR_GENTRACE_TEST_CASE_ID, ATTR_GENTRACE_EXPERIMENT_ID, + ATTR_GENTRACE_IN_EXPERIMENT, ATTR_GENTRACE_TEST_CASE_NAME, MAX_EVAL_DATASET_CONCURRENCY, ATTR_GENTRACE_FN_ARGS_EVENT_NAME, @@ -215,6 +216,9 @@ async def _run_single_test_case_for_dataset( # Set up baggage context similar to @interaction() current_context = otel_context.get_current() context_with_modified_baggage = otel_baggage.set_baggage(ATTR_GENTRACE_SAMPLE_KEY, "true", context=current_context) + context_with_modified_baggage = otel_baggage.set_baggage( + ATTR_GENTRACE_IN_EXPERIMENT, "true", context=current_context + ) token = otel_context.attach(context_with_modified_baggage) try: diff --git a/src/gentrace/lib/span_processor.py b/src/gentrace/lib/span_processor.py index b3e7ae70..ef7ef426 100644 --- a/src/gentrace/lib/span_processor.py +++ b/src/gentrace/lib/span_processor.py @@ -6,7 +6,7 @@ from opentelemetry.context import Context from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor -from gentrace.lib.constants import ATTR_GENTRACE_SAMPLE_KEY +from gentrace.lib.constants import ATTR_GENTRACE_SAMPLE_KEY, ATTR_GENTRACE_IN_EXPERIMENT class GentraceSpanProcessor(SpanProcessor): @@ -35,6 +35,11 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None if isinstance(sample_value, str): span.set_attribute(ATTR_GENTRACE_SAMPLE_KEY, sample_value) + in_experiment_value = baggage.get_baggage(ATTR_GENTRACE_IN_EXPERIMENT, context=parent_context) + if in_experiment_value is not None: + if isinstance(in_experiment_value, str): + span.set_attribute(ATTR_GENTRACE_IN_EXPERIMENT, in_experiment_value) + @override def on_end(self, span: ReadableSpan) -> None: """ From 56e6656a45adc08d77e2a6b3fa444f2a9ba9b525 Mon Sep 17 00:00:00 2001 From: Josh Lebedinsky Date: Fri, 15 Aug 2025 14:46:21 -0400 Subject: [PATCH 2/3] nit --- src/gentrace/lib/eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gentrace/lib/eval.py b/src/gentrace/lib/eval.py index c437b062..7e375d00 100644 --- a/src/gentrace/lib/eval.py +++ b/src/gentrace/lib/eval.py @@ -92,7 +92,7 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> Any: ATTR_GENTRACE_SAMPLE_KEY, "true", context=current_context ) context_with_modified_baggage = otel_baggage.set_baggage( - ATTR_GENTRACE_IN_EXPERIMENT, "true", context=current_context + ATTR_GENTRACE_IN_EXPERIMENT, "true", context=context_with_modified_baggage ) token = otel_context.attach(context_with_modified_baggage) From 744b01b150033575cee1de7c3255e991fae7cb2a Mon Sep 17 00:00:00 2001 From: Josh Lebedinsky Date: Fri, 15 Aug 2025 15:12:36 -0400 Subject: [PATCH 3/3] nit --- src/gentrace/lib/eval_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gentrace/lib/eval_dataset.py b/src/gentrace/lib/eval_dataset.py index 80588707..0e7f310f 100644 --- a/src/gentrace/lib/eval_dataset.py +++ b/src/gentrace/lib/eval_dataset.py @@ -217,7 +217,7 @@ async def _run_single_test_case_for_dataset( current_context = otel_context.get_current() context_with_modified_baggage = otel_baggage.set_baggage(ATTR_GENTRACE_SAMPLE_KEY, "true", context=current_context) context_with_modified_baggage = otel_baggage.set_baggage( - ATTR_GENTRACE_IN_EXPERIMENT, "true", context=current_context + ATTR_GENTRACE_IN_EXPERIMENT, "true", context=context_with_modified_baggage ) token = otel_context.attach(context_with_modified_baggage)