Skip to content

Commit 0d6c87d

Browse files
authored
1 parent b3f292a commit 0d6c87d

2 files changed

Lines changed: 40 additions & 10 deletions

File tree

py/src/braintrust/integrations/livekit_agents/test_livekit_agents.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,7 @@ def _assert_openai_voice_turn_spans(logs, *, speech_text, outer_parent_name):
409409
_assert_eou_spans(logs)
410410
_assert_stt_spans(logs, speech_text)
411411
_assert_session_spans(logs)
412+
_assert_no_function_tool_spans_without_tools(logs)
412413
_assert_custom_metrics_stay_in_metadata(logs)
413414

414415
if outer_parent_name:
@@ -478,13 +479,17 @@ def _assert_session_spans(logs):
478479
_assert_no_metrics(session_logs, "prompt_tokens", "completion_tokens", "tokens")
479480

480481

482+
def _assert_no_function_tool_spans_without_tools(logs):
483+
assert not _spans_named(logs, "function_tool")
484+
485+
481486
def _assert_voice_turn_parenting(logs, outer_parent_name):
482487
outer_parent_log = _single_span(logs, outer_parent_name)
483488
session_log = _single_span(logs, "livekit_agent_session")
484489
assert session_log.get("span_parents") == [outer_parent_log.get("span_id")]
485490
assert "end" in session_log.get("metrics", {})
486491

487-
for child_name in ("function_tool", "llm_request_run", "user_speaking", "vad_endpointing"):
492+
for child_name in ("llm_request_run", "user_speaking", "vad_endpointing"):
488493
child_log = _single_span(logs, child_name)
489494
assert child_log.get("span_parents") == [session_log.get("span_id")], child_name
490495
assert _single_span(logs, "llm_request_run")["span_attributes"]["type"].value == "task"

py/src/braintrust/integrations/livekit_agents/tracing.py

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -218,18 +218,43 @@ async def traced_llm_stream_run(wrapped: Any, instance: Any, args: tuple[Any, ..
218218
async def traced_execute_tools_task(wrapped: Any, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
219219
session = kwargs.get("session")
220220
parent = getattr(session, _SESSION_PARENT_ATTR, None)
221-
result = await _traced_async_span(
222-
"function_tool",
223-
SpanTypeAttribute.TOOL,
224-
wrapped,
225-
args,
226-
kwargs,
227-
parent=parent if isinstance(parent, str) else None,
228-
event_builder=lambda: _tool_execution_event(kwargs.get("tool_output")),
229-
)
221+
parent_arg = parent if isinstance(parent, str) else None
222+
start_time = time.time()
223+
try:
224+
with _without_current_span(parent_arg):
225+
result = await wrapped(*args, **kwargs)
226+
except Exception as exc:
227+
end_time = time.time()
228+
event = _tool_execution_event(kwargs.get("tool_output"))
229+
if event:
230+
_log_function_tool_span(event, parent=parent_arg, start_time=start_time, end_time=end_time, error=exc)
231+
raise
232+
end_time = time.time()
233+
event = _tool_execution_event(kwargs.get("tool_output"))
234+
if event:
235+
_log_function_tool_span(event, parent=parent_arg, start_time=start_time, end_time=end_time)
230236
return result
231237

232238

239+
def _log_function_tool_span(
240+
event: dict[str, Any],
241+
*,
242+
parent: str | None,
243+
start_time: float,
244+
end_time: float,
245+
error: Exception | None = None,
246+
) -> None:
247+
span = start_span(
248+
name="function_tool",
249+
type=SpanTypeAttribute.TOOL,
250+
parent=parent,
251+
start_time=start_time,
252+
set_current=False,
253+
)
254+
span.log(error=error, **event)
255+
span.end(end_time=end_time)
256+
257+
233258
async def _traced_async_span(
234259
name: str,
235260
span_type: str,

0 commit comments

Comments
 (0)