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..7e375d00 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=context_with_modified_baggage + ) 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..0e7f310f 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=context_with_modified_baggage + ) 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: """