feat: Context management for using with that manages trace(s) automatically (#1019)

* add context management to use `with` with `start_trace`

* set to `Success` when no error encountered

* add examples

* add tests for context management

* add docs

* add `TraceState` for handling status codes and handle thread management for client

* modify to use updated status codes

* heckin' ruff

* gh copilot said do it so i did

* global tracer everywher

* should fix linting

* Convert end_state to string using helper function

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Revert "Convert end_state to string using helper function"

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: dot-agi

agentops/__init__.py

@@ -15,18 +15,28 @@
 from typing import List, Optional, Union, Dict, Any
 from agentops.client import Client
 from agentops.sdk.core import TraceContext, tracer
-from agentops.sdk.decorators import trace, session, agent, task, workflow, operation
+from agentops.sdk.decorators import trace, session, agent, task, workflow, operation, tool
+from agentops.enums import TraceState, SUCCESS, ERROR, UNSET
+from opentelemetry.trace.status import StatusCode
 
 from agentops.logging.config import logger
+import threading
 
-# Client global instance; one per process runtime
-_client = Client()
+# Thread-safe client management
+_client_lock = threading.Lock()
+_client = None
 
 
 def get_client() -> Client:
-    """Get the singleton client instance"""
+    """Get the singleton client instance in a thread-safe manner"""
     global _client
 
+    # Double-checked locking pattern for thread safety
+    if _client is None:
+        with _client_lock:
+            if _client is None:
+                _client = Client()
+
     return _client
 
 
@@ -106,24 +116,31 @@ def init(
     elif default_tags:
         merged_tags = default_tags
 
-    return _client.init(
-        api_key=api_key,
-        endpoint=endpoint,
-        app_url=app_url,
-        max_wait_time=max_wait_time,
-        max_queue_size=max_queue_size,
-        default_tags=merged_tags,
-        trace_name=trace_name,
-        instrument_llm_calls=instrument_llm_calls,
-        auto_start_session=auto_start_session,
-        auto_init=auto_init,
-        skip_auto_end_session=skip_auto_end_session,
-        env_data_opt_out=env_data_opt_out,
-        log_level=log_level,
-        fail_safe=fail_safe,
-        exporter_endpoint=exporter_endpoint,
+    # Prepare initialization arguments
+    init_kwargs = {
+        "api_key": api_key,
+        "endpoint": endpoint,
+        "app_url": app_url,
+        "max_wait_time": max_wait_time,
+        "max_queue_size": max_queue_size,
+        "default_tags": merged_tags,
+        "trace_name": trace_name,
+        "instrument_llm_calls": instrument_llm_calls,
+        "auto_start_session": auto_start_session,
+        "auto_init": auto_init,
+        "skip_auto_end_session": skip_auto_end_session,
+        "env_data_opt_out": env_data_opt_out,
+        "log_level": log_level,
+        "fail_safe": fail_safe,
+        "exporter_endpoint": exporter_endpoint,
         **kwargs,
-    )
+    }
+
+    # Get the current client instance (creates new one if needed)
+    client = get_client()
+
+    # Initialize the client directly
+    return client.init(**init_kwargs)
 
 
 def configure(**kwargs):
@@ -173,7 +190,8 @@ def configure(**kwargs):
     if invalid_params:
         logger.warning(f"Invalid configuration parameters: {invalid_params}")
 
-    _client.configure(**kwargs)
+    client = get_client()
+    client.configure(**kwargs)
 
 
 def start_trace(
@@ -207,7 +225,9 @@ def start_trace(
     return tracer.start_trace(trace_name=trace_name, tags=tags)
 
 
-def end_trace(trace_context: Optional[TraceContext] = None, end_state: str = "Success") -> None:
+def end_trace(
+    trace_context: Optional[TraceContext] = None, end_state: Union[TraceState, StatusCode, str] = TraceState.SUCCESS
+) -> None:
     """
     Ends a trace (its root span) and finalizes it.
     If no trace_context is provided, ends all active session spans.
@@ -246,4 +266,12 @@ def end_trace(trace_context: Optional[TraceContext] = None, end_state: str = "Su
     "workflow",
     "operation",
     "tracer",
+    "tool",
+    # Trace state enums
+    "TraceState",
+    "SUCCESS",
+    "ERROR",
+    "UNSET",
+    # OpenTelemetry status codes (for advanced users)
+    "StatusCode",
 ]

agentops/client/client.py

@@ -24,7 +24,7 @@ def _end_init_trace_atexit():
     if _client_init_trace_context is not None:
         logger.debug("Auto-ending client's init trace during shutdown.")
         try:
-            # Use TracingCore to end the trace directly
+            # Use global tracer to end the trace directly
             if tracer.initialized and _client_init_trace_context.span.is_recording():
                 tracer.end_trace(_client_init_trace_context, end_state="Shutdown")
         except Exception as e:
@@ -104,7 +104,7 @@ def init(self, **kwargs: Any) -> None:  # Return type updated to None
         try:
             response = self.api.v3.fetch_auth_token(self.config.api_key)
             if response is None:
-                # If auth fails, we cannot proceed with TracingCore initialization that depends on project_id
+                # If auth fails, we cannot proceed with tracer initialization that depends on project_id
                 logger.error("Failed to fetch auth token. AgentOps SDK will not be initialized.")
                 return None  # Explicitly return None if auth fails
         except Exception as e:
@@ -161,8 +161,8 @@ def init(self, **kwargs: Any) -> None:  # Return type updated to None
 
                 else:
                     logger.error("Failed to start the auto-init trace.")
-                    # Even if auto-start fails, core services up to TracingCore might be initialized.
-                    # Set self.initialized to True if TracingCore is up, but return None.
+                    # Even if auto-start fails, core services up to the tracer might be initialized.
+                    # Set self.initialized to True if tracer is up, but return None.
                     self._initialized = tracer.initialized
                     return None  # Failed to start trace
 

agentops/enums.py

@@ -0,0 +1,36 @@
+"""
+AgentOps enums for user-friendly API.
+
+This module provides simple enums that users can import from agentops
+without needing to know about OpenTelemetry internals.
+"""
+
+from enum import Enum
+from opentelemetry.trace.status import StatusCode
+
+
+class TraceState(Enum):
+    """
+    Enum for trace end states.
+
+    This provides a user-friendly interface that maps to OpenTelemetry StatusCode internally.
+    Users can simply use agentops.TraceState.SUCCESS instead of importing OpenTelemetry.
+    """
+
+    SUCCESS = StatusCode.OK
+    ERROR = StatusCode.ERROR
+    UNSET = StatusCode.UNSET
+
+    def __str__(self) -> str:
+        """Return the name for string representation."""
+        return self.name
+
+    def to_status_code(self) -> StatusCode:
+        """Convert to OpenTelemetry StatusCode."""
+        return self.value
+
+
+# For backward compatibility, also provide these as module-level constants
+SUCCESS = TraceState.SUCCESS
+ERROR = TraceState.ERROR
+UNSET = TraceState.UNSET

agentops/instrumentation/__init__.py

@@ -444,7 +444,7 @@ def instrument_one(loader: InstrumentorLoader) -> Optional[BaseInstrumentor]:
 
     instrumentor = loader.get_instance()
     try:
-        instrumentor.instrument(tracer_provider=TracingCore.get_instance()._provider)
+        instrumentor.instrument(tracer_provider=tracer._provider)
         logger.info(
             f"AgentOps: Successfully instrumented '{loader.class_name}' for package '{loader.package_name or loader.module_name}'."
         )

agentops/legacy/__init__.py

@@ -65,7 +65,7 @@ def start_session(
 ) -> Session:
     """
     @deprecated Use agentops.start_trace() instead.
-    Starts a legacy AgentOps session. Calls TracingCore.start_trace internally.
+    Starts a legacy AgentOps session. Calls tracer.start_trace internally.
     """
     global _current_session, _current_trace_context
 
@@ -89,7 +89,7 @@ def start_session(
 
     trace_context = tracer.start_trace(trace_name="session", tags=tags)
     if trace_context is None:
-        logger.error("Failed to start trace via TracingCore. Returning dummy session.")
+        logger.error("Failed to start trace via global tracer. Returning dummy session.")
         dummy_session = Session(None)
         _current_session = dummy_session
         _current_trace_context = None
@@ -124,13 +124,13 @@ def _set_span_attributes(span: Any, attributes: Dict[str, Any]) -> None:
 def end_session(session_or_status: Any = None, **kwargs: Any) -> None:
     """
     @deprecated Use agentops.end_trace() instead.
-    Ends a legacy AgentOps session. Calls TracingCore.end_trace internally.
+    Ends a legacy AgentOps session. Calls tracer.end_trace internally.
     Supports multiple calling patterns for backward compatibility.
     """
     global _current_session, _current_trace_context
 
     if not tracer.initialized:
-        logger.debug("Ignoring end_session: TracingCore not initialized.")
+        logger.debug("Ignoring end_session: global tracer not initialized.")
         return
 
     target_trace_context: Optional[TraceContext] = None
@@ -189,7 +189,7 @@ def end_session(session_or_status: Any = None, **kwargs: Any) -> None:
 def end_all_sessions() -> None:
     """@deprecated Ends all active sessions/traces."""
     if not tracer.initialized:
-        logger.debug("Ignoring end_all_sessions: TracingCore not initialized.")
+        logger.debug("Ignoring end_all_sessions: global tracer not initialized.")
         return
 
     # Use the new end_trace functionality to end all active traces

agentops/sdk/README.md

@@ -176,11 +176,11 @@ flowchart TD
 ## Example Usage
 
 ```python
-from agentops import Session, agent, tool
-from agentops.sdk import TracingCore, TracingConfig
+from agentops import Session, agent, tool, tracer
+from agentops.sdk import TracingConfig
 
-# Initialize the tracing core with a dedicated configuration
-TracingCore.get_instance().initialize(
+# Initialize the global tracer with a dedicated configuration
+tracer.initialize(
     service_name="my-service",
     exporter_endpoint="https://my-exporter-endpoint.com",
     max_queue_size=1000,

agentops/sdk/__init__.py

@@ -5,23 +5,23 @@
 for different types of operations in AI agent workflows.
 """
 
-# 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
 
+from opentelemetry.trace.status import StatusCode
+
 __all__ = [
     # Core components
-    "TracingCore",
     "TracingConfig",
     # Decorators
     "session",
     "operation",
     "agent",
     "task",
     "workflow",
+    # OpenTelemetry status codes
+    "StatusCode",
 ]