feat: add same logic as entp

This commit is contained in:
Eduardo Chiarotti
2025-02-10 12:31:56 -03:00
parent e3026ebd56
commit 8f3bf31339

View File

@@ -80,6 +80,10 @@ class UnifiedTraceController:
self.flow_step = flow_step self.flow_step = flow_step
self.status: str = "running" self.status: str = "running"
# Add trace to task's trace collection if task_id is present
if task_id:
self._add_to_task_traces()
def _add_to_task_traces(self) -> None: def _add_to_task_traces(self) -> None:
"""Add this trace to the task's trace collection.""" """Add this trace to the task's trace collection."""
if not hasattr(UnifiedTraceController, "_task_traces"): if not hasattr(UnifiedTraceController, "_task_traces"):
@@ -213,15 +217,20 @@ class UnifiedTraceController:
) )
def _record_trace(self, result: Any = None) -> None: def _record_trace(self, result: Any = None) -> None:
"""Record the trace using PlusClient. """Record the trace.
This method is called when a trace is completed. It ensures the trace
is properly recorded and associated with its task if applicable.
Args: Args:
result: Optional result to include in the trace. result: Optional result to include in the trace
""" """
if result: if result:
self.context["response"] = result self.context["response"] = result
# TODO: Add trace to record_task_finished # Add to task traces if this trace belongs to a task
if self.task_id:
self._add_to_task_traces()
def should_trace() -> bool: def should_trace() -> bool:
@@ -311,8 +320,7 @@ def init_flow_main_trace(
with TraceContext.set_current(trace): with TraceContext.set_current(trace):
try: try:
return await func(self, *args, **kwargs) return await func(self, *args, **kwargs)
except Exception as e: except Exception:
trace.end_trace(error=str(e))
raise raise
return wrapper return wrapper