Better root span management

agentops/__init__.py

@@ -12,9 +12,12 @@
     LLMEvent,
 )  # type: ignore
 
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict, Any
 from agentops.client import Client
+from agentops.sdk.core import TracingCore, TraceContext
+from agentops.sdk.decorators import trace, session, agent, task, workflow, operation
 
+from agentops.logging.config import logger
 
 # Client global instance; one per process runtime
 _client = Client()
@@ -53,6 +56,7 @@
     max_queue_size: Optional[int] = None,
     tags: Optional[List[str]] = None,
     default_tags: Optional[List[str]] = None,
+    trace_name: Optional[str] = None,
     instrument_llm_calls: Optional[bool] = None,
     auto_start_session: Optional[bool] = None,
     auto_init: Optional[bool] = None,
@@ -78,6 +82,7 @@
         max_queue_size (int, optional): The maximum size of the event queue. Defaults to 512.
         tags (List[str], optional): [Deprecated] Use `default_tags` instead.
         default_tags (List[str], optional): Default tags for the sessions that can be used for grouping or sorting later (e.g. ["GPT-4"]).
+        trace_name (str, optional): Name for the default trace/session. If none is provided, defaults to "default".
         instrument_llm_calls (bool): Whether to instrument LLM calls and emit LLMEvents.
         auto_start_session (bool): Whether to start a session automatically when the client is created.
         auto_init (bool): Whether to automatically initialize the client on import. Defaults to True.
@@ -108,6 +113,7 @@
         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,
@@ -165,26 +171,80 @@
     # Check for invalid parameters
     invalid_params = set(kwargs.keys()) - valid_params
     if invalid_params:
-        from .logging.config import logger
-
         logger.warning(f"Invalid configuration parameters: {invalid_params}")
 
     _client.configure(**kwargs)
 
 
+def start_trace(
+    trace_name: str = "session", tags: Optional[Union[Dict[str, Any], List[str]]] = None
+) -> Optional[TraceContext]:
+    """
+    Starts a new trace (root span) and returns its context.
+    This allows for multiple concurrent, user-managed traces.
+
+    Args:
+        trace_name: Name for the trace (e.g., "session", "my_custom_task").
+        tags: Optional tags to attach to the trace span (list of strings or dict).
+
+    Returns:
+        A TraceContext object containing the span and context token, or None if SDK not initialized.
+    """
+    tracing_core = TracingCore.get_instance()
+    if not tracing_core.initialized:
+        # Optionally, attempt to initialize the client if not already, or log a more severe warning.
+        # For now, align with legacy start_session that would try to init.
+        # However, explicit init is preferred before starting traces.
+        logger.warning("AgentOps SDK not initialized. Attempting to initialize with defaults before starting trace.")
+        try:
+            init()  # Attempt to initialize with environment variables / defaults
+            if not tracing_core.initialized:
+                logger.error("SDK initialization failed. Cannot start trace.")
+                return None
+        except Exception as e:
+            logger.error(f"SDK auto-initialization failed during start_trace: {e}. Cannot start trace.")
+            return None
+
+    return tracing_core.start_trace(trace_name=trace_name, tags=tags)
+
+
+def end_trace(trace_context: Optional[TraceContext] = None, end_state: str = "Success") -> None:
+    """
+    Ends a trace (its root span) and finalizes it.
+    If no trace_context is provided, ends all active session spans.
+
+    Args:
+        trace_context: The TraceContext object returned by start_trace. If None, ends all active traces.
+        end_state: The final state of the trace (e.g., "Success", "Failure", "Error").
+    """
+    tracing_core = TracingCore.get_instance()
+    if not tracing_core.initialized:
+        logger.warning("AgentOps SDK not initialized. Cannot end trace.")
+        return
+    tracing_core.end_trace(trace_context=trace_context, end_state=end_state)
+
+
 __all__ = [
     "init",
     "configure",
     "get_client",
     "record",
+    "start_trace",
+    "end_trace",
     "start_session",
     "end_session",
     "track_agent",
     "track_tool",
     "end_all_sessions",
-    "Session",
     "ToolEvent",
     "ErrorEvent",
     "ActionEvent",
     "LLMEvent",
+    "Session",
+    "trace",
+    "session",
+    "agent",
+    "task",
+    "workflow",
+    "operation",
 ]

agentops/client/api/versions/v3.py

@@ -30,28 +30,26 @@
 
         r = self.post(path, data, headers)
 
+        if r.status_code != 200:
+            error_msg = f"Authentication failed: {r.status_code}"
+            try:
+                error_data = r.json()
+                if "error" in error_data:
+                    error_msg = f"{error_data['error']}"
+            except Exception:
+                pass
+            logger.error(f"{error_msg} - Perhaps an invalid API key?")
+            raise ApiServerException(error_msg)
+
         try:
-            if r.status_code != 200:
-                error_msg = f"Authentication failed: {r.status_code}"
-                try:
-                    error_data = r.json()
-                    if "error" in error_data:
-                        error_msg = f"{error_data['error']}"
-                except Exception:
-                    pass
-                raise ApiServerException(error_msg)
+            jr = r.json()
+            token = jr.get("token")
+            if not token:
+                raise ApiServerException("No token in authentication response")
 
-            try:
-                jr = r.json()
-                token = jr.get("token")
-                if not token:
-                    raise ApiServerException("No token in authentication response")
-
-                return jr
-            except Exception as e:
-                raise ApiServerException(f"Failed to process authentication response: {str(e)}")
+            return jr
         except Exception as e:
-            logger.error(f"{str(e)} - Perhaps an invalid API key?")
-            return None
+            logger.error(f"Failed to process authentication response: {str(e)}")
+            raise ApiServerException(f"Failed to process authentication response: {str(e)}")
 
     # Add V3-specific API methods here