Skip to content

Commit 1ea28d7

Browse files
viniciusdsmellogustavocidornelas
authored andcommitted
feat(tracer): enhance OpenlayerTracerProcessor with dynamic base class and type hinting
- Introduced a dynamic base class for `OpenlayerTracerProcessor` to handle the presence of the `agents` library, improving compatibility. - Added type hinting for tracing-related parameters and return types, enhancing code clarity and type safety. - Implemented an ImportError raise for better error handling when the `agents` library is not available. - Updated dictionary type annotations for improved type specificity.
1 parent c7232ce commit 1ea28d7

File tree

1 file changed

+32
-11
lines changed

1 file changed

+32
-11
lines changed

src/openlayer/lib/integrations/openai_agents.py

Lines changed: 32 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,24 @@
44
import logging
55
from pathlib import Path
66
import time
7-
from typing import Any, Dict, Optional, Union, List
7+
from datetime import datetime
8+
from typing import Any, Dict, Optional, Union, List, TYPE_CHECKING
89

910
from ..tracing import tracer, steps, enums
1011

12+
if TYPE_CHECKING:
13+
try:
14+
from agents import tracing # type: ignore[import]
15+
except ImportError:
16+
# When agents isn't available, we'll use string literals for type annotations
17+
pass
18+
1119
try:
1220
from agents import tracing # type: ignore[import]
13-
1421
HAVE_AGENTS = True
1522
except ImportError:
1623
HAVE_AGENTS = False
24+
tracing = None # type: ignore[assignment]
1725

1826
logger = logging.getLogger(__name__)
1927

@@ -582,7 +590,14 @@ def _configure_chat_completion_step(
582590
step.model_parameters = model_parameters or {}
583591

584592

585-
class OpenlayerTracerProcessor(tracing.TracingProcessor): # type: ignore[no-redef]
593+
# Dynamic base class to handle inheritance when agents is available
594+
if HAVE_AGENTS:
595+
_BaseProcessor = tracing.TracingProcessor # type: ignore[misc]
596+
else:
597+
_BaseProcessor = object # type: ignore[assignment,misc]
598+
599+
600+
class OpenlayerTracerProcessor(_BaseProcessor): # type: ignore[misc]
586601
"""Tracing processor for the `OpenAI Agents SDK
587602
<https://openai.github.io/openai-agents-python/>`_.
588603
@@ -649,6 +664,12 @@ def __init__(self, **kwargs: Any) -> None:
649664
Args:
650665
**kwargs: Additional metadata to associate with all traces.
651666
"""
667+
if not HAVE_AGENTS:
668+
raise ImportError(
669+
"The 'agents' library is required to use OpenlayerTracerProcessor. "
670+
"Please install it with: pip install openai-agents"
671+
)
672+
652673
self.metadata: Dict[str, Any] = kwargs or {}
653674
self._active_traces: Dict[str, Dict[str, Any]] = {}
654675
self._active_steps: Dict[str, steps.Step] = {}
@@ -676,7 +697,7 @@ def __init__(self, **kwargs: Any) -> None:
676697
global _active_openlayer_processor
677698
_active_openlayer_processor = self
678699

679-
def on_trace_start(self, trace: tracing.Trace) -> None:
700+
def on_trace_start(self, trace: "tracing.Trace") -> None:
680701
"""Handle the start of a trace (root agent workflow)."""
681702
try:
682703
# Get trace information
@@ -693,7 +714,7 @@ def on_trace_start(self, trace: tracing.Trace) -> None:
693714
except Exception as e:
694715
logger.error(f"Failed to handle trace start: {e}")
695716

696-
def on_trace_end(self, trace: tracing.Trace) -> None:
717+
def on_trace_end(self, trace: "tracing.Trace") -> None:
697718
"""Handle the end of a trace (root agent workflow)."""
698719
try:
699720
trace_data = self._active_traces.pop(trace.trace_id, None)
@@ -786,7 +807,7 @@ def on_trace_end(self, trace: tracing.Trace) -> None:
786807
except Exception as e:
787808
logger.error(f"Failed to handle trace end: {e}")
788809

789-
def on_span_start(self, span: tracing.Span) -> None:
810+
def on_span_start(self, span: "tracing.Span") -> None:
790811
"""Handle the start of a span (individual agent step)."""
791812
try:
792813
# Extract span attributes using helper function
@@ -840,7 +861,7 @@ def on_span_start(self, span: tracing.Span) -> None:
840861
except Exception as e:
841862
logger.error(f"Failed to handle span start: {e}")
842863

843-
def on_span_end(self, span: tracing.Span) -> None:
864+
def on_span_end(self, span: "tracing.Span") -> None:
844865
"""Handle the end of a span (individual agent step)."""
845866
try:
846867
# Extract span attributes using helper function
@@ -912,7 +933,7 @@ def on_span_end(self, span: tracing.Span) -> None:
912933
logger.error(f"Failed to handle span end: {e}")
913934

914935
def _create_step_for_span(
915-
self, span: tracing.Span, span_data: Any
936+
self, span: "tracing.Span", span_data: Any
916937
) -> Optional[steps.Step]:
917938
"""Create the appropriate Openlayer step for a span."""
918939
try:
@@ -1315,7 +1336,7 @@ def _create_generic_step(
13151336
step.start_time = start_time
13161337
return step
13171338

1318-
def _extract_usage_from_response(self, response: Any, field: str = None) -> int:
1339+
def _extract_usage_from_response(self, response: Any, field: Optional[str] = None) -> Union[int, Dict[str, int]]:
13191340
"""Extract usage information from response object."""
13201341
if not response:
13211342
return 0
@@ -1339,7 +1360,7 @@ def _extract_usage_from_response(self, response: Any, field: str = None) -> int:
13391360
}
13401361

13411362
def _update_step_with_span_data(
1342-
self, step: steps.Step, span: tracing.Span, span_data: Any
1363+
self, step: steps.Step, span: "tracing.Span", span_data: Any
13431364
) -> None:
13441365
"""Update step with final span data."""
13451366
try:
@@ -1650,7 +1671,7 @@ def _extract_actual_llm_output(self, span_data: Any) -> Optional[str]:
16501671
except Exception:
16511672
return None
16521673

1653-
def _cleanup_dict_with_warning(self, dict_obj: Dict, name: str) -> None:
1674+
def _cleanup_dict_with_warning(self, dict_obj: Dict[str, Any], name: str) -> None:
16541675
"""Helper to clean up dictionaries with warning logging."""
16551676
if dict_obj:
16561677
dict_obj.clear()

0 commit comments

Comments
 (0)