mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
feat: add same logic as entp
This commit is contained in:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user