Skip to content

Fixing OpenAI tracing issues #974

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

Merged
merged 9 commits into from
Jul 18, 2025
Merged

Fixing OpenAI tracing issues #974

Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
252254e
Fixing tracing issues. Propagate span parentage, remove startActivity
tconley1428 Jul 16, 2025
40602bb
Remove traceback
tconley1428 Jul 16, 2025
38385e2
Merge branch 'main' into openai/tracing_fixes
tconley1428 Jul 16, 2025
22c0245
Merge branch 'main' into openai/tracing_fixes
tconley1428 Jul 16, 2025
8476649
Add start traces with correct durations
tconley1428 Jul 17, 2025
5fb122f
Add tracing test
tconley1428 Jul 17, 2025
c930698
Remove cast
tconley1428 Jul 17, 2025
d2c8784
Remove flaky test check
tconley1428 Jul 18, 2025
54c4555
Merge branch 'main' into openai/tracing_fixes
tconley1428 Jul 18, 2025
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
88 changes: 61 additions & 27 deletions temporalio/contrib/openai_agents/_trace_interceptor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,15 @@
from __future__ import annotations

from contextlib import contextmanager
from typing import Any, Mapping, Protocol, Type
from typing import Any, Mapping, Protocol, Type, cast

from agents import custom_span, get_current_span, trace
from agents import CustomSpanData, custom_span, get_current_span, trace
from agents.tracing import (
get_trace_provider,
)
from agents.tracing.spans import NoOpSpan
from agents.tracing.provider import DefaultTraceProvider
from agents.tracing.scope import Scope
from agents.tracing.spans import NoOpSpan, SpanImpl

import temporalio.activity
import temporalio.api.common.v1
Expand Down Expand Up @@ -65,11 +67,15 @@ def context_from_header(
else workflow.info().workflow_type
)
data = (
{"activityId": activity.info().activity_id}
{
"activityId": activity.info().activity_id,
"activity": activity.info().activity_type,
}
if activity.in_activity()
else None
)
if get_trace_provider().get_current_trace() is None:
current_trace = get_trace_provider().get_current_trace()
if current_trace is None:
metadata = {
"temporal:workflowId": activity.info().workflow_id
if activity.in_activity()
Expand All @@ -79,16 +85,21 @@ def context_from_header(
else workflow.info().run_id,
"temporal:workflowType": workflow_type,
}
with trace(
current_trace = trace(
span_info["traceName"],
trace_id=span_info["traceId"],
metadata=metadata,
) as t:
with custom_span(name=span_name, parent=t, data=data):
yield
else:
with custom_span(name=span_name, parent=None, data=data):
yield
)
Scope.set_current_trace(current_trace)
current_span = get_trace_provider().get_current_span()
if current_span is None:
current_span = get_trace_provider().create_span(
span_data=CustomSpanData(name="", data={}), span_id=span_info["spanId"]
)
Scope.set_current_span(current_span)

with custom_span(name=span_name, parent=current_span, data=data):
yield


class OpenAIAgentsTracingInterceptor(
Expand All @@ -115,7 +126,7 @@ class OpenAIAgentsTracingInterceptor(
worker = Worker(client, task_queue="my-task-queue", interceptors=[interceptor])
"""

def __init__( # type: ignore[reportMissingSuperCall]
def __init__(
self,
payload_converter: temporalio.converter.PayloadConverter = temporalio.converter.default().payload_converter,
) -> None:
Expand Down Expand Up @@ -188,7 +199,7 @@ async def start_workflow(
**({"temporal:workflowId": input.id} if input.id else {}),
}
data = {"workflowId": input.id} if input.id else None
span_name = "temporal:startWorkflow"
span_name = f"temporal:startWorkflow"
if get_trace_provider().get_current_trace() is None:
with trace(
span_name + ":" + input.workflow, metadata=metadata, group_id=input.id
Expand All @@ -207,7 +218,7 @@ async def query_workflow(self, input: temporalio.client.QueryWorkflowInput) -> A
**({"temporal:workflowId": input.id} if input.id else {}),
}
data = {"workflowId": input.id, "query": input.query}
span_name = "temporal:queryWorkflow"
span_name = f"temporal:queryWorkflow"
if get_trace_provider().get_current_trace() is None:
with trace(span_name, metadata=metadata, group_id=input.id):
with custom_span(name=span_name, data=data):
Expand All @@ -226,7 +237,7 @@ async def signal_workflow(
**({"temporal:workflowId": input.id} if input.id else {}),
}
data = {"workflowId": input.id, "signal": input.signal}
span_name = "temporal:signalWorkflow"
span_name = f"temporal:signalWorkflow"
if get_trace_provider().get_current_trace() is None:
with trace(span_name, metadata=metadata, group_id=input.id):
with custom_span(name=span_name, data=data):
Expand Down Expand Up @@ -325,32 +336,55 @@ class _ContextPropagationWorkflowOutboundInterceptor(
async def signal_child_workflow(
self, input: temporalio.worker.SignalChildWorkflowInput
) -> None:
set_header_from_context(input, temporalio.workflow.payload_converter())
return await self.next.signal_child_workflow(input)
with custom_span(
name=f"temporal:signalChildWorkflow",
data={"workflowId": input.child_workflow_id},
):
set_header_from_context(input, temporalio.workflow.payload_converter())
await self.next.signal_child_workflow(input)

async def signal_external_workflow(
self, input: temporalio.worker.SignalExternalWorkflowInput
) -> None:
set_header_from_context(input, temporalio.workflow.payload_converter())
return await self.next.signal_external_workflow(input)
with custom_span(
name=f"temporal:signalExternalWorkflow",
data={"workflowId": input.workflow_id},
):
set_header_from_context(input, temporalio.workflow.payload_converter())
await self.next.signal_external_workflow(input)

def start_activity(
self, input: temporalio.worker.StartActivityInput
) -> temporalio.workflow.ActivityHandle:
with custom_span(
name=f"temporal:startActivity:{input.activity}",
):
set_header_from_context(input, temporalio.workflow.payload_converter())
return self.next.start_activity(input)
span = custom_span(
name=f"temporal:startActivity", data={"activity": input.activity}
)
span.start(mark_as_current=True)
set_header_from_context(input, temporalio.workflow.payload_converter())
handle = self.next.start_activity(input)
handle.add_done_callback(lambda _: span.finish())
return handle

async def start_child_workflow(
self, input: temporalio.worker.StartChildWorkflowInput
) -> temporalio.workflow.ChildWorkflowHandle:
span = custom_span(
name=f"temporal:startChildWorkflow", data={"workflow": input.workflow}
)
span.start(mark_as_current=True)
set_header_from_context(input, temporalio.workflow.payload_converter())
return await self.next.start_child_workflow(input)
handle = await self.next.start_child_workflow(input)
handle.add_done_callback(lambda _: span.finish())
return handle

def start_local_activity(
self, input: temporalio.worker.StartLocalActivityInput
) -> temporalio.workflow.ActivityHandle:
span = custom_span(
name=f"temporal:startLocalActivity", data={"activity": input.activity}
)
span.start(mark_as_current=True)
set_header_from_context(input, temporalio.workflow.payload_converter())
return self.next.start_local_activity(input)
handle = self.next.start_local_activity(input)
handle.add_done_callback(lambda _: span.finish())
return handle
Loading