diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 138eedbcc..602bd718d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,7 +1,9 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.2.1" + rev: "v0.9.1" hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] + exclude: ^examples/ - id: ruff-format + exclude: ^examples/ diff --git a/agentops/__init__.py b/agentops/__init__.py index 25e28dcf3..4d66e7b57 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -7,21 +7,22 @@ # Client global instance; one per process runtime _client = Client() + def record(event): """ Legacy function to record an event. This is kept for backward compatibility. - + In the current version, this simply sets the end_timestamp on the event. - + Args: event: The event to record """ from agentops.helpers.time import get_ISO_time - + # TODO: Manual timestamp assignment is a temporary fix; should use proper event lifecycle - if event and hasattr(event, 'end_timestamp'): + if event and hasattr(event, "end_timestamp"): event.end_timestamp = get_ISO_time() - + return event @@ -141,6 +142,7 @@ def configure(**kwargs): _client.configure(**kwargs) + # For backwards compatibility and testing @@ -149,8 +151,7 @@ def get_client() -> Client: return _client - -from agentops.legacy import * # type: ignore +from agentops.legacy import * # noqa: E402, F403 __all__ = [ "init", @@ -159,6 +160,4 @@ def get_client() -> Client: "record", "start_session", "end_session", - "track_agent", - "track_tool", ] diff --git a/agentops/client/api/base.py b/agentops/client/api/base.py index c7654154e..92a162293 100644 --- a/agentops/client/api/base.py +++ b/agentops/client/api/base.py @@ -15,8 +15,7 @@ class TokenFetcher(Protocol): """Protocol for token fetching functions""" - def __call__(self, api_key: str) -> str: - ... + def __call__(self, api_key: str) -> str: ... class BaseApiClient: diff --git a/agentops/client/client.py b/agentops/client/client.py index bce0423ff..c9c4d5808 100644 --- a/agentops/client/client.py +++ b/agentops/client/client.py @@ -68,7 +68,7 @@ def init(self, **kwargs): session = start_session(tags=list(self.config.default_tags)) else: session = start_session() - + return session def configure(self, **kwargs): diff --git a/agentops/config.py b/agentops/config.py index 8ee08db22..cf5d1cf1d 100644 --- a/agentops/config.py +++ b/agentops/config.py @@ -49,7 +49,7 @@ class Config: default_factory=lambda: get_env_int("AGENTOPS_MAX_WAIT_TIME", 5000), metadata={"description": "Maximum time in milliseconds to wait for API responses"}, ) - + export_flush_interval: int = field( default_factory=lambda: get_env_int("AGENTOPS_EXPORT_FLUSH_INTERVAL", 1000), metadata={"description": "Time interval in milliseconds between automatic exports of telemetry data"}, @@ -154,7 +154,7 @@ def configure( if max_wait_time is not None: self.max_wait_time = max_wait_time - + if export_flush_interval is not None: self.export_flush_interval = export_flush_interval diff --git a/agentops/helpers/validation.py b/agentops/helpers/validation.py index 2a0c219cf..78c5d2008 100644 --- a/agentops/helpers/validation.py +++ b/agentops/helpers/validation.py @@ -4,4 +4,5 @@ def is_coroutine_or_generator(fn: Any) -> bool: """Check if a function is asynchronous (coroutine or async generator)""" import inspect + return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn) diff --git a/agentops/legacy/__init__.py b/agentops/legacy/__init__.py index ea097dec1..057b5ea44 100644 --- a/agentops/legacy/__init__.py +++ b/agentops/legacy/__init__.py @@ -23,7 +23,7 @@ class Session: """ This class provides compatibility with CrewAI >= 0.105.0, which uses an event-based integration pattern where it calls methods directly on the Session object: - + - create_agent(): Called when a CrewAI agent is created - record(): Called when a CrewAI tool is used - end_session(): Called when a CrewAI run completes @@ -42,7 +42,7 @@ def __del__(self): def create_agent(self, name: Optional[str] = None, agent_id: Optional[str] = None, **kwargs): """ Method to create an agent for CrewAI >= 0.105.0 compatibility. - + CrewAI >= 0.105.0 calls this with: - name=agent.role - agent_id=str(agent.id) @@ -52,7 +52,7 @@ def create_agent(self, name: Optional[str] = None, agent_id: Optional[str] = Non def record(self, event=None): """ Method to record events for CrewAI >= 0.105.0 compatibility. - + CrewAI >= 0.105.0 calls this with a tool event when a tool is used. """ pass @@ -60,11 +60,11 @@ def record(self, event=None): def end_session(self, **kwargs): """ Method to end the session for CrewAI >= 0.105.0 compatibility. - + CrewAI >= 0.105.0 calls this with: - end_state="Success" - end_state_reason="Finished Execution" - + forces a flush to ensure the span is exported immediately. """ _set_span_attributes(self.span, kwargs) @@ -75,7 +75,7 @@ def end_session(self, **kwargs): def _create_session_span(tags: Union[Dict[str, Any], List[str], None] = None) -> tuple: """ Helper function to create a session span with tags. - + This is an internal function used by start_session() to create the from the SDK to create a span with kind=SpanKind.SESSION. @@ -106,12 +106,12 @@ def start_session( This function creates and starts a new session span, which can be used to group related operations together. The session will remain active until end_session is called either with the Session object or with kwargs. - + Usage patterns: 1. Standard pattern: session = start_session(); end_session(session) 2. CrewAI < 0.105.0: start_session(); end_session(end_state="Success", ...) 3. CrewAI >= 0.105.0: session = start_session(); session.end_session(end_state="Success", ...) - + This function stores the session in a global variable to support the CrewAI < 0.105.0 pattern where end_session is called without the session object. @@ -127,11 +127,12 @@ def start_session( AgentOpsClientNotInitializedException: If the client is not initialized """ global _current_session - + if not TracingCore.get_instance().initialized: from agentops import Client + Client().init() - + span, context, token = _create_session_span(tags) session = Session(span, token) _current_session = session @@ -141,14 +142,14 @@ def start_session( def _set_span_attributes(span: Any, attributes: Dict[str, Any]) -> None: """ Helper to set attributes on a span. - + Args: span: The span to set attributes on attributes: The attributes to set as a dictionary """ if not attributes or not hasattr(span, "set_attribute"): return - + for key, value in attributes.items(): span.set_attribute(f"agentops.status.{key}", str(value)) @@ -159,11 +160,12 @@ def _flush_span_processors() -> None: """ try: from opentelemetry.trace import get_tracer_provider + tracer_provider = get_tracer_provider() tracer_provider.force_flush() # type: ignore except Exception as e: logger.warning(f"Failed to force flush span processor: {e}") - + def end_session(session_or_status: Any = None, **kwargs) -> None: """ @@ -181,24 +183,25 @@ def end_session(session_or_status: Any = None, **kwargs) -> None: Args: session_or_status: The session object returned by start_session, or a string representing the status (for backwards compatibility) - **kwargs: Additional arguments for CrewAI < 0.105.0 compatibility. + **kwargs: Additional arguments for CrewAI < 0.105.0 compatibility. CrewAI < 0.105.0 passes these named arguments: - end_state="Success" - end_state_reason="Finished Execution" - is_auto_end=True - + When called this way, the function will use the most recently created session via start_session(). """ from agentops.sdk.decorators.utility import _finalize_span - + from agentops.sdk.core import TracingCore + if not TracingCore.get_instance().initialized: logger.debug("Ignoring end_session call - TracingCore not initialized") return - # In some old implementations, and in crew < 0.10.5 `end_session` will be - # called with a single string as a positional argument like: "Success" + # In some old implementations, and in crew < 0.10.5 `end_session` will be + # called with a single string as a positional argument like: "Success" # Handle the CrewAI < 0.105.0 integration pattern where end_session is called # with only named parameters. In this pattern, CrewAI does not keep a reference @@ -211,18 +214,18 @@ def end_session(session_or_status: Any = None, **kwargs) -> None: # ) if session_or_status is None and kwargs: global _current_session - + if _current_session is not None: _set_span_attributes(_current_session.span, kwargs) _finalize_span(_current_session.span, _current_session.token) _flush_span_processors() _current_session = None return - + # Handle the standard pattern and CrewAI >= 0.105.0 pattern where a Session object is passed. # In both cases, we call _finalize_span with the span and token from the Session. # This is the most direct and precise way to end a specific session. - if hasattr(session_or_status, 'span') and hasattr(session_or_status, 'token'): + if hasattr(session_or_status, "span") and hasattr(session_or_status, "token"): _set_span_attributes(session_or_status.span, kwargs) _finalize_span(session_or_status.span, session_or_status.token) _flush_span_processors() @@ -231,8 +234,8 @@ def end_session(session_or_status: Any = None, **kwargs) -> None: def end_all_sessions(): """ @deprecated - We don't automatically track more than one session, so just end the session - that we are tracking. + We don't automatically track more than one session, so just end the session + that we are tracking. """ end_session() @@ -249,17 +252,17 @@ def ErrorEvent(*args, **kwargs): """ @deprecated Use tracing instead. - + For backward compatibility with tests, this returns a minimal object with the required attributes. """ from agentops.helpers.time import get_ISO_time - + class LegacyErrorEvent: def __init__(self): self.init_timestamp = get_ISO_time() self.end_timestamp = None - + return LegacyErrorEvent() @@ -267,17 +270,17 @@ def ActionEvent(*args, **kwargs): """ @deprecated Use tracing instead. - + For backward compatibility with tests, this returns a minimal object with the required attributes. """ from agentops.helpers.time import get_ISO_time - + class LegacyActionEvent: def __init__(self): self.init_timestamp = get_ISO_time() self.end_timestamp = None - + return LegacyActionEvent() @@ -294,28 +297,32 @@ def track_agent(*args, **kwargs): @deprecated Decorator for marking agents in legacy projects. """ + def noop(f): return f + return noop def track_tool(*args, **kwargs): """ @deprecated - Decorator for marking tools and legacy projects. + Decorator for marking tools and legacy projects. """ + def noop(f): return f + return noop __all__ = [ - "start_session", - "end_session", - "ToolEvent", - "ErrorEvent", - "ActionEvent", - "track_agent", + "start_session", + "end_session", + "ToolEvent", + "ErrorEvent", + "ActionEvent", + "track_agent", "track_tool", - "end_all_sessions" + "end_all_sessions", ] diff --git a/agentops/sdk/__init__.py b/agentops/sdk/__init__.py index 1b0779dd5..f1be1f718 100644 --- a/agentops/sdk/__init__.py +++ b/agentops/sdk/__init__.py @@ -7,8 +7,10 @@ # Import core components from agentops.sdk.core import TracingCore + # Import decorators from agentops.sdk.decorators import agent, operation, session, task, workflow + # from agentops.sdk.traced import TracedObject # Merged into TracedObject from agentops.sdk.types import TracingConfig diff --git a/agentops/sdk/core.py b/agentops/sdk/core.py index 540089e8c..887e17d35 100644 --- a/agentops/sdk/core.py +++ b/agentops/sdk/core.py @@ -5,10 +5,8 @@ from typing import List, Optional from opentelemetry import metrics, trace -from opentelemetry.exporter.otlp.proto.http.metric_exporter import \ - OTLPMetricExporter -from opentelemetry.exporter.otlp.proto.http.trace_exporter import \ - OTLPSpanExporter +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader from opentelemetry.sdk.resources import Resource @@ -66,10 +64,7 @@ def setup_telemetry( trace.set_tracer_provider(provider) # Create exporter with authentication - exporter = OTLPSpanExporter( - endpoint=exporter_endpoint, - headers={"Authorization": f"Bearer {jwt}"} if jwt else {} - ) + exporter = OTLPSpanExporter(endpoint=exporter_endpoint, headers={"Authorization": f"Bearer {jwt}"} if jwt else {}) # Regular processor for normal spans and immediate export processor = BatchSpanProcessor( @@ -82,10 +77,7 @@ def setup_telemetry( # Setup metrics metric_reader = PeriodicExportingMetricReader( - OTLPMetricExporter( - endpoint=metrics_endpoint, - headers={"Authorization": f"Bearer {jwt}"} if jwt else {} - ) + OTLPMetricExporter(endpoint=metrics_endpoint, headers={"Authorization": f"Bearer {jwt}"} if jwt else {}) ) meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader]) metrics.set_meter_provider(meter_provider) @@ -201,7 +193,7 @@ def shutdown(self) -> None: # Perform a single flush on the SynchronousSpanProcessor (which takes care of all processors' shutdown) if not self._initialized: return - self._provider._active_span_processor.force_flush(self.config['max_wait_time']) # type: ignore + self._provider._active_span_processor.force_flush(self.config["max_wait_time"]) # type: ignore # Shutdown provider if self._provider: diff --git a/agentops/sdk/decorators/__init__.py b/agentops/sdk/decorators/__init__.py index a7ffad8a9..706bd4624 100644 --- a/agentops/sdk/decorators/__init__.py +++ b/agentops/sdk/decorators/__init__.py @@ -16,12 +16,6 @@ session = create_entity_decorator(SpanKind.SESSION) operation = task -__all__ = [ - 'agent', - 'task', - 'workflow', - 'session', - 'operation' -] +__all__ = ["agent", "task", "workflow", "session", "operation"] # Create decorators task, workflow, session, agent diff --git a/agentops/sdk/decorators/factory.py b/agentops/sdk/decorators/factory.py index b29ade4d6..1852ed314 100644 --- a/agentops/sdk/decorators/factory.py +++ b/agentops/sdk/decorators/factory.py @@ -8,26 +8,33 @@ from agentops.logging import logger from agentops.sdk.core import TracingCore -from .utility import (_create_as_current_span, _finalize_span, _make_span, - _process_async_generator, _process_sync_generator, - _record_entity_input, _record_entity_output) +from .utility import ( + _create_as_current_span, + _finalize_span, + _make_span, + _process_async_generator, + _process_sync_generator, + _record_entity_input, + _record_entity_output, +) def create_entity_decorator(entity_kind: str): """ Factory function that creates decorators for specific entity kinds. - + Args: entity_kind: The type of operation being performed (SpanKind.*) - + Returns: A decorator with optional arguments for name and version """ + def decorator(wrapped=None, *, name=None, version=None): # Handle case where decorator is called with parameters if wrapped is None: return functools.partial(decorator, name=name, version=version) - + # Handle class decoration if inspect.isclass(wrapped): # Create a proxy class that wraps the original class @@ -37,33 +44,33 @@ def __init__(self, *args, **kwargs): operation_name = name or wrapped.__name__ self._agentops_span_context_manager = _create_as_current_span(operation_name, entity_kind, version) self._agentops_active_span = self._agentops_span_context_manager.__enter__() - + try: _record_entity_input(self._agentops_active_span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") - + # Call the original __init__ super().__init__(*args, **kwargs) - + def __del__(self): # End span when instance is destroyed - if hasattr(self, '_agentops_active_span') and hasattr(self, '_agentops_span_context_manager'): + if hasattr(self, "_agentops_active_span") and hasattr(self, "_agentops_span_context_manager"): try: _record_entity_output(self._agentops_active_span, self) except Exception as e: logger.warning(f"Failed to record entity output: {e}") - + self._agentops_span_context_manager.__exit__(None, None, None) - + # Preserve metadata of the original class WrappedClass.__name__ = wrapped.__name__ WrappedClass.__qualname__ = wrapped.__qualname__ WrappedClass.__module__ = wrapped.__module__ WrappedClass.__doc__ = wrapped.__doc__ - + return WrappedClass - + # Create the actual decorator wrapper function for functions @wrapt.decorator def wrapper(wrapped, instance, args, kwargs): @@ -87,10 +94,10 @@ def wrapper(wrapped, instance, args, kwargs): _record_entity_input(span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") - + result = wrapped(*args, **kwargs) return _process_sync_generator(span, result) - + # Handle async generator functions elif is_async_generator: # Use the old approach for async generators @@ -99,19 +106,20 @@ def wrapper(wrapped, instance, args, kwargs): _record_entity_input(span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") - + result = wrapped(*args, **kwargs) return _process_async_generator(span, token, result) - + # Handle async functions elif is_async: + async def _wrapped_async(): with _create_as_current_span(operation_name, entity_kind, version) as span: try: _record_entity_input(span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") - + try: result = await wrapped(*args, **kwargs) try: @@ -122,9 +130,9 @@ async def _wrapped_async(): except Exception as e: span.record_exception(e) raise - + return _wrapped_async() - + # Handle sync functions else: with _create_as_current_span(operation_name, entity_kind, version) as span: @@ -132,7 +140,7 @@ async def _wrapped_async(): _record_entity_input(span, args, kwargs) except Exception as e: logger.warning(f"Failed to record entity input: {e}") - + try: result = wrapped(*args, **kwargs) try: @@ -145,8 +153,6 @@ async def _wrapped_async(): raise # Return the wrapper for functions, we already returned WrappedClass for classes - return wrapper(wrapped) # type: ignore - - return decorator - + return wrapper(wrapped) # type: ignore + return decorator diff --git a/agentops/sdk/decorators/utility.py b/agentops/sdk/decorators/utility.py index 0e3953dfd..6426a1fd8 100644 --- a/agentops/sdk/decorators/utility.py +++ b/agentops/sdk/decorators/utility.py @@ -31,6 +31,7 @@ def set_workflow_name(workflow_name: str) -> None: def set_entity_path(entity_path: str) -> None: attach(set_value("entity_path", entity_path)) + # Helper functions for content management @@ -73,17 +74,14 @@ def _get_current_span_info(): "span_id": f"{ctx.span_id:x}" if hasattr(ctx, "span_id") else "None", "trace_id": f"{ctx.trace_id:x}" if hasattr(ctx, "trace_id") else "None", "name": getattr(current_span, "name", "Unknown"), - "is_recording": getattr(current_span, "is_recording", False) + "is_recording": getattr(current_span, "is_recording", False), } return {"name": "No current span"} @contextmanager def _create_as_current_span( - operation_name: str, - span_kind: str, - version: Optional[int] = None, - attributes: Optional[Dict[str, Any]] = None + operation_name: str, span_kind: str, version: Optional[int] = None, attributes: Optional[Dict[str, Any]] = None ) -> Generator[Span, None, None]: """ Create and yield an instrumentation span as the current span using proper context management. @@ -104,7 +102,7 @@ def _create_as_current_span( # Log before we do anything before_span = _get_current_span_info() logger.debug(f"[DEBUG] BEFORE {operation_name}.{span_kind} - Current context: {before_span}") - + # Create span with proper naming convention span_name = f"{operation_name}.{span_kind}" @@ -125,26 +123,25 @@ def _create_as_current_span( # Get current context explicitly to debug it current_context = context_api.get_current() - + # Use OpenTelemetry's context manager to properly handle span lifecycle with tracer.start_as_current_span(span_name, attributes=attributes, context=current_context) as span: # Log after span creation if hasattr(span, "get_span_context"): span_ctx = span.get_span_context() - logger.debug(f"[DEBUG] CREATED {span_name} - span_id: {span_ctx.span_id:x}, parent: {before_span.get('span_id', 'None')}") - + logger.debug( + f"[DEBUG] CREATED {span_name} - span_id: {span_ctx.span_id:x}, parent: {before_span.get('span_id', 'None')}" + ) + yield span - + # Log after we're done after_span = _get_current_span_info() logger.debug(f"[DEBUG] AFTER {operation_name}.{span_kind} - Returned to context: {after_span}") def _make_span( - operation_name: str, - span_kind: str, - version: Optional[int] = None, - attributes: Optional[Dict[str, Any]] = None + operation_name: str, span_kind: str, version: Optional[int] = None, attributes: Optional[Dict[str, Any]] = None ) -> tuple: """ Create a span without context management for manual span lifecycle control. @@ -167,7 +164,7 @@ def _make_span( # Log before we do anything before_span = _get_current_span_info() logger.debug(f"[DEBUG] BEFORE _make_span {operation_name}.{span_kind} - Current context: {before_span}") - + # Create span with proper naming convention span_name = f"{operation_name}.{span_kind}" @@ -188,18 +185,20 @@ def _make_span( # Get current context explicitly current_context = context_api.get_current() - + # Create the span with explicit context span = tracer.start_span(span_name, context=current_context, attributes=attributes) # Set as current context and get token for later detachment ctx = trace.set_span_in_context(span) token = context_api.attach(ctx) - + # Log after span creation if hasattr(span, "get_span_context"): span_ctx = span.get_span_context() - logger.debug(f"[DEBUG] CREATED _make_span {span_name} - span_id: {span_ctx.span_id:x}, parent: {before_span.get('span_id', 'None')}") + logger.debug( + f"[DEBUG] CREATED _make_span {span_name} - span_id: {span_ctx.span_id:x}, parent: {before_span.get('span_id', 'None')}" + ) return span, ctx, token @@ -236,15 +235,15 @@ def _finalize_span(span: trace.Span, token: Any) -> None: if hasattr(span, "get_span_context") and hasattr(span.get_span_context(), "span_id"): span_id = f"{span.get_span_context().span_id:x}" logger.debug(f"[DEBUG] ENDING span {getattr(span, 'name', 'unknown')} - span_id: {span_id}") - + span.end() - + # Debug info before detaching current_after_end = _get_current_span_info() logger.debug(f"[DEBUG] AFTER span.end() - Current context: {current_after_end}") - + context_api.detach(token) - + # Debug info after detaching final_context = _get_current_span_info() logger.debug(f"[DEBUG] AFTER detach - Final context: {final_context}") diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index 324c7b443..d6f23bc98 100644 --- a/agentops/semconv/span_attributes.py +++ b/agentops/semconv/span_attributes.py @@ -58,4 +58,4 @@ class SpanAttributes: # Operation attributes OPERATION_NAME = "operation.name" - OPERATION_VERSION = "operation.version" \ No newline at end of file + OPERATION_VERSION = "operation.version" diff --git a/agentops/semconv/span_kinds.py b/agentops/semconv/span_kinds.py index 190afacbe..3315e04b4 100644 --- a/agentops/semconv/span_kinds.py +++ b/agentops/semconv/span_kinds.py @@ -16,7 +16,7 @@ class SpanKind: # Workflow kinds WORKFLOW_STEP = "workflow.step" # Step in a workflow - WORKFLOW = 'workflow' + WORKFLOW = "workflow" SESSION = "session" TASK = "task" OPERATION = "operation" diff --git a/tests/benchmark/benchmark_init.py b/tests/benchmark/benchmark_init.py index 674062a0f..93432d2b5 100644 --- a/tests/benchmark/benchmark_init.py +++ b/tests/benchmark/benchmark_init.py @@ -7,10 +7,11 @@ Benchmark script for measuring TracingCore initialization time. """ + def run_benchmark(): """ Run a benchmark of TracingCore initialization. - + Returns: Dictionary with timing results """ @@ -24,19 +25,19 @@ def run_benchmark(): return { "init": init_time, - "total": init_time # Total time is just init time now + "total": init_time, # Total time is just init time now } def print_results(results): """ Print benchmark results in a formatted way. - + Args: results: Dictionary with timing results """ print("\n=== BENCHMARK RESULTS ===") - + print(f"\nINIT TIME: {results['init']:.6f}s") print(f"TOTAL TIME: {results['total']:.6f}s") @@ -44,4 +45,4 @@ def print_results(results): if __name__ == "__main__": print("Running TracingCore benchmark...") results = run_benchmark() - print_results(results) + print_results(results) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 742c1c89f..90ebecd64 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -32,8 +32,10 @@ def mock_req(endpoint, api_key): """ with requests_mock.Mocker(real_http=False) as m: # Map session IDs to their JWTs - m.post(endpoint + "/v3/auth/token", json={"token": str(uuid.uuid4()), - "project_id": "test-project-id", "api_key": api_key}) + m.post( + endpoint + "/v3/auth/token", + json={"token": str(uuid.uuid4()), "project_id": "test-project-id", "api_key": api_key}, + ) yield m diff --git a/tests/unit/sdk/instrumentation_tester.py b/tests/unit/sdk/instrumentation_tester.py index 9e5dc80d5..a3f868805 100644 --- a/tests/unit/sdk/instrumentation_tester.py +++ b/tests/unit/sdk/instrumentation_tester.py @@ -5,8 +5,7 @@ from opentelemetry import trace as trace_api from opentelemetry.sdk.trace import ReadableSpan, Span, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import \ - InMemorySpanExporter +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.util.types import Attributes from agentops.sdk.core import TracingCore, setup_telemetry @@ -46,8 +45,7 @@ def reset_trace_globals(): class HasAttributesViaProperty(Protocol): @property - def attributes(self) -> Attributes: - ... + def attributes(self) -> Attributes: ... class HasAttributesViaAttr(Protocol): @@ -92,8 +90,7 @@ def __init__(self): # Patch the setup_telemetry function to return our test providers self.setup_telemetry_patcher = mock.patch( - 'agentops.sdk.core.setup_telemetry', - return_value=(self.tracer_provider, self.mock_meter_provider) + "agentops.sdk.core.setup_telemetry", return_value=(self.tracer_provider, self.mock_meter_provider) ) self.mock_setup_telemetry = self.setup_telemetry_patcher.start() diff --git a/tests/unit/sdk/test_decorators.py b/tests/unit/sdk/test_decorators.py index e67c85c6c..8227d9427 100644 --- a/tests/unit/sdk/test_decorators.py +++ b/tests/unit/sdk/test_decorators.py @@ -8,7 +8,6 @@ from agentops.sdk.decorators import agent, operation, session, workflow, task from agentops.semconv import SpanKind from agentops.semconv.span_attributes import SpanAttributes -from agentops.semconv import SpanAttributes from tests.unit.sdk.instrumentation_tester import InstrumentationTester @@ -62,10 +61,15 @@ def test_session(): assert len(spans) == 4 # Verify span kinds - session_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION] - agent_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT] - operation_spans = [s for s in spans if s.attributes and s.attributes.get( - SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK] + session_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION + ] + agent_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT + ] + operation_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK + ] assert len(session_spans) == 1 assert len(agent_spans) == 1 @@ -74,31 +78,31 @@ def test_session(): # Find the main_operation and nested_operation spans main_operation = None nested_operation = None - + for span in operation_spans: - if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'main_operation': + if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_operation": main_operation = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'nested_operation': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_operation": nested_operation = span - + assert main_operation is not None, "main_operation span not found" assert nested_operation is not None, "nested_operation span not found" - + # Verify the session span is the root session_span = session_spans[0] assert session_span.parent is None - + # Verify the agent span is a child of the session span agent_span = agent_spans[0] assert agent_span.parent is not None assert session_span.context is not None assert agent_span.parent.span_id == session_span.context.span_id - + # Verify main_operation is a child of the agent span assert main_operation.parent is not None assert agent_span.context is not None assert main_operation.parent.span_id == agent_span.context.span_id - + # Verify nested_operation is a child of main_operation assert nested_operation.parent is not None assert main_operation.context is not None @@ -151,10 +155,15 @@ async def test_async_session(): assert len(spans) == 4 # Verify span kinds - session_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION] - agent_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT] - operation_spans = [s for s in spans if s.attributes and s.attributes.get( - SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK] + session_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION + ] + agent_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT + ] + operation_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK + ] assert len(session_spans) == 1 assert len(agent_spans) == 1 @@ -163,31 +172,31 @@ async def test_async_session(): # Find the main_operation and nested_operation spans main_operation = None nested_operation = None - + for span in operation_spans: - if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'main_async_operation': + if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_async_operation": main_operation = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'nested_async_operation': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_async_operation": nested_operation = span - + assert main_operation is not None, "main_async_operation span not found" assert nested_operation is not None, "nested_async_operation span not found" - + # Verify the session span is the root session_span = session_spans[0] assert session_span.parent is None - + # Verify the agent span is a child of the session span agent_span = agent_spans[0] assert agent_span.parent is not None assert session_span.context is not None assert agent_span.parent.span_id == session_span.context.span_id - + # Verify main_operation is a child of the agent span assert main_operation.parent is not None assert agent_span.context is not None assert main_operation.parent.span_id == agent_span.context.span_id - + # Verify nested_operation is a child of main_operation assert nested_operation.parent is not None assert main_operation.context is not None @@ -242,10 +251,15 @@ def test_generator_session(): assert len(spans) == 4 # Verify span kinds - session_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION] - agent_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT] - operation_spans = [s for s in spans if s.attributes and s.attributes.get( - SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK] + session_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION + ] + agent_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT + ] + operation_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK + ] assert len(session_spans) == 1 assert len(agent_spans) == 1 @@ -254,31 +268,31 @@ def test_generator_session(): # Find the main_operation and nested_operation spans main_operation = None nested_operation = None - + for span in operation_spans: - if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'main_generator_operation': + if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_generator_operation": main_operation = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'nested_generator': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_generator": nested_operation = span - + assert main_operation is not None, "main_generator_operation span not found" assert nested_operation is not None, "nested_generator span not found" - + # Verify the session span is the root session_span = session_spans[0] assert session_span.parent is None - + # Verify the agent span is a child of the session span agent_span = agent_spans[0] assert agent_span.parent is not None assert session_span.context is not None assert agent_span.parent.span_id == session_span.context.span_id - + # Verify main_operation is a child of the agent span assert main_operation.parent is not None assert agent_span.context is not None assert main_operation.parent.span_id == agent_span.context.span_id - + # Verify nested_operation is a child of main_operation assert nested_operation.parent is not None assert main_operation.context is not None @@ -334,10 +348,15 @@ async def test_async_generator_session(): assert len(spans) == 4 # Verify span kinds - session_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION] - agent_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT] - operation_spans = [s for s in spans if s.attributes and s.attributes.get( - SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK] + session_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION + ] + agent_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT + ] + operation_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK + ] assert len(session_spans) == 1 assert len(agent_spans) == 1 @@ -346,31 +365,34 @@ async def test_async_generator_session(): # Find the main_operation and nested_operation spans main_operation = None nested_operation = None - + for span in operation_spans: - if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'main_async_generator_operation': + if ( + span.attributes + and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_async_generator_operation" + ): main_operation = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'nested_async_generator': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_async_generator": nested_operation = span - + assert main_operation is not None, "main_async_generator_operation span not found" assert nested_operation is not None, "nested_async_generator span not found" - + # Verify the session span is the root session_span = session_spans[0] assert session_span.parent is None - + # Verify the agent span is a child of the session span agent_span = agent_spans[0] assert agent_span.parent is not None assert session_span.context is not None assert agent_span.parent.span_id == session_span.context.span_id - + # Verify main_operation is a child of the agent span assert main_operation.parent is not None assert agent_span.context is not None assert main_operation.parent.span_id == agent_span.context.span_id - + # Verify nested_operation is a child of main_operation assert nested_operation.parent is not None assert main_operation.context is not None @@ -428,10 +450,15 @@ def test_complex_session(): assert len(spans) == 5 # Verify span kinds - session_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION] - agent_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT] - operation_spans = [s for s in spans if s.attributes and s.attributes.get( - SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK] + session_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION + ] + agent_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT + ] + operation_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK + ] assert len(session_spans) == 1 assert len(agent_spans) == 1 @@ -441,39 +468,39 @@ def test_complex_session(): level1_operation = None level2_operation = None level3_operation = None - + for span in operation_spans: - if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'level1_operation': + if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "level1_operation": level1_operation = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'level2_operation': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "level2_operation": level2_operation = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'level3_operation': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "level3_operation": level3_operation = span - + assert level1_operation is not None, "level1_operation span not found" assert level2_operation is not None, "level2_operation span not found" assert level3_operation is not None, "level3_operation span not found" - + # Verify the session span is the root session_span = session_spans[0] assert session_span.parent is None - + # Verify the agent span is a child of the session span agent_span = agent_spans[0] assert agent_span.parent is not None assert session_span.context is not None assert agent_span.parent.span_id == session_span.context.span_id - + # Verify level1_operation is a child of the agent span assert level1_operation.parent is not None assert agent_span.context is not None assert level1_operation.parent.span_id == agent_span.context.span_id - + # Verify level2_operation is a child of level1_operation assert level2_operation.parent is not None assert level1_operation.context is not None assert level2_operation.parent.span_id == level1_operation.context.span_id - + # Verify level3_operation is a child of level2_operation assert level3_operation.parent is not None assert level2_operation.context is not None @@ -525,10 +552,17 @@ def test_workflow_session(): assert len(spans) == 4 # Verify span kinds - session_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION] - workflow_spans = [s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.WORKFLOW] - task_spans = [s for s in spans if s.attributes and s.attributes.get( - SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK] + session_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION + ] + workflow_spans = [ + s + for s in spans + if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.WORKFLOW + ] + task_spans = [ + s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK + ] assert len(session_spans) == 1 assert len(workflow_spans) == 1 @@ -538,34 +572,34 @@ def test_workflow_session(): workflow_span = None process_task = None transform_task = None - + for span in spans: - if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'data_processing_workflow': + if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "data_processing_workflow": workflow_span = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'process_input': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "process_input": process_task = span - elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == 'transform_data': + elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "transform_data": transform_task = span - + assert workflow_span is not None, "workflow span not found" assert process_task is not None, "process_input task span not found" assert transform_task is not None, "transform_data task span not found" - + # Verify the session span is the root session_span = session_spans[0] assert session_span.parent is None - + # Verify the workflow span is a child of the session span assert workflow_span.parent is not None assert session_span.context is not None assert workflow_span.parent.span_id == session_span.context.span_id - + # Verify process_task is a child of the workflow span assert process_task.parent is not None assert workflow_span.context is not None assert process_task.parent.span_id == workflow_span.context.span_id - + # Verify transform_task is a child of the workflow span assert transform_task.parent is not None assert workflow_span.context is not None - assert transform_task.parent.span_id == workflow_span.context.span_id \ No newline at end of file + assert transform_task.parent.span_id == workflow_span.context.span_id diff --git a/tests/unit/test_session_legacy.py b/tests/unit/test_session_legacy.py index e63557787..455874c2e 100644 --- a/tests/unit/test_session_legacy.py +++ b/tests/unit/test_session_legacy.py @@ -1,12 +1,10 @@ - - def test_session_auto_start(instrumentation): import agentops from agentops.legacy import Session # Pass a dummy API key for the test session = agentops.init(api_key="test-api-key", auto_start_session=True) - + assert isinstance(session, Session) @@ -42,12 +40,12 @@ def __init__(self): self.role = "Test Agent" self.goal = "Testing" self.id = "test-agent-id" - + agent = MockAgent() agentops.track_agent(agent) except Exception as e: assert False, f"track_agent raised an exception: {e}" - + # Test track_tool function exists and doesn't raise errors try: # Mock a tool object similar to what CrewAI would provide @@ -55,7 +53,7 @@ class MockTool: def __init__(self): self.name = "Test Tool" self.description = "A test tool" - + tool = MockTool() agentops.track_tool(tool, "Test Agent") except Exception as e: @@ -64,97 +62,85 @@ def __init__(self): # Test events that CrewAI might use tool_event = agentops.ToolEvent(name="test_tool") action_event = agentops.ActionEvent(action_type="test_action") - + # Verify that record function works with these events agentops.record(tool_event) agentops.record(action_event) - - + + def test_crewai_kwargs_pattern(instrumentation): """ Test the CrewAI < 0.105.0 pattern where end_session is called with only kwargs. - + In versions < 0.105.0, CrewAI directly calls: agentops.end_session( end_state="Success", - end_state_reason="Finished Execution", + end_state_reason="Finished Execution", is_auto_end=True ) """ import agentops from agentops.legacy import Session - + # Initialize with test API key agentops.init(api_key="test-api-key") - + # Create a session session = agentops.start_session(tags=["test", "crewai-kwargs"]) assert isinstance(session, Session) - + # Test the CrewAI < 0.105.0 pattern - calling end_session with only kwargs - agentops.end_session( - end_state="Success", - end_state_reason="Finished Execution", - is_auto_end=True - ) - + agentops.end_session(end_state="Success", end_state_reason="Finished Execution", is_auto_end=True) + # After calling end_session, creating a new session should work correctly # (this implicitly tests that the internal state is reset properly) new_session = agentops.start_session(tags=["test", "post-end"]) assert isinstance(new_session, Session) - - + + def test_crewai_kwargs_pattern_no_session(instrumentation): """ Test the CrewAI < 0.105.0 pattern where end_session is called with only kwargs, but no session has been created. - + This should log a warning but not fail. """ import agentops - + # Initialize with test API key agentops.init(api_key="test-api-key") - + # We don't need to explicitly clear the session state # Just make sure we start with a clean state by calling init - + # Test the CrewAI < 0.105.0 pattern - calling end_session with only kwargs # when no session exists. This should not raise an error. - agentops.end_session( - end_state="Success", - end_state_reason="Finished Execution", - is_auto_end=True - ) + agentops.end_session(end_state="Success", end_state_reason="Finished Execution", is_auto_end=True) def test_crewai_kwargs_force_flush(): """ Test that when using the CrewAI < 0.105.0 pattern (end_session with kwargs), the spans are properly exported to the backend with force_flush. - + This is a more comprehensive test that ensures spans are actually sent to the backend when using the CrewAI integration pattern. """ import agentops from agentops.sdk.core import TracingCore import time - + # Initialize AgentOps with API key agentops.init(api_key="test-api-key") - + # Create a session session = agentops.start_session(tags=["test", "crewai-integration"]) - + # Simulate some work time.sleep(0.1) - + # End session with kwargs (CrewAI < 0.105.0 pattern) - agentops.end_session( - end_state="Success", - end_state_reason="Test Finished", - is_auto_end=True - ) - + agentops.end_session(end_state="Success", end_state_reason="Test Finished", is_auto_end=True) + # Explicitly ensure the core isn't already shut down for the test - assert TracingCore.get_instance()._initialized, "TracingCore should still be initialized" \ No newline at end of file + assert TracingCore.get_instance()._initialized, "TracingCore should still be initialized"