mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
90f1bee602
* feat: add prompt observability code * feat: improve logic for llm call * feat: add tests for traces * feat: remove unused improt * feat: add function to clear and add task traces * feat: fix import * feat: chagne time * feat: fix type checking issues * feat: add fixed time to fix test * feat: fix datetime test issue * feat: add add task traces function * feat: add same logic as entp * feat: add start_time as reference for duplication of tool call * feat: add max_depth * feat: add protocols file to properly import on LLM --------- Co-authored-by: Brandon Hancock (bhancock_ai) <109994880+bhancockio@users.noreply.github.com>
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from typing import Generator
|
|
|
|
|
|
class TraceContext:
|
|
"""Maintains the current trace context throughout the execution stack.
|
|
|
|
This class provides a context manager for tracking trace execution across
|
|
async and sync code paths using ContextVars.
|
|
"""
|
|
|
|
_context: ContextVar = ContextVar("trace_context", default=None)
|
|
|
|
@classmethod
|
|
def get_current(cls):
|
|
"""Get the current trace context.
|
|
|
|
Returns:
|
|
Optional[UnifiedTraceController]: The current trace controller or None if not set.
|
|
"""
|
|
return cls._context.get()
|
|
|
|
@classmethod
|
|
@contextmanager
|
|
def set_current(cls, trace):
|
|
"""Set the current trace context within a context manager.
|
|
|
|
Args:
|
|
trace: The trace controller to set as current.
|
|
|
|
Yields:
|
|
UnifiedTraceController: The current trace controller.
|
|
"""
|
|
token = cls._context.set(trace)
|
|
try:
|
|
yield trace
|
|
finally:
|
|
cls._context.reset(token)
|