Introduced start_trace and end_trace functions for user-managed tracing, allowing concurrent traces.

Author: Dwij1704

agentops/__init__.py

@@ -12,9 +12,11 @@
     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.logging.config import logger
 
 # Client global instance; one per process runtime
 _client = Client()
@@ -51,18 +53,14 @@ def init(
     app_url: Optional[str] = None,
     max_wait_time: Optional[int] = None,
     max_queue_size: Optional[int] = None,
-    tags: Optional[List[str]] = None,
     default_tags: Optional[List[str]] = None,
     instrument_llm_calls: Optional[bool] = None,
     auto_start_session: Optional[bool] = None,
-    auto_init: Optional[bool] = None,
-    skip_auto_end_session: Optional[bool] = None,
     env_data_opt_out: Optional[bool] = None,
     log_level: Optional[Union[str, int]] = None,
-    fail_safe: Optional[bool] = None,
     exporter_endpoint: Optional[str] = None,
     **kwargs,
-):
+) -> None:
     """
     Initializes the AgentOps SDK.
 
@@ -76,45 +74,30 @@ def init(
         max_wait_time (int, optional): The maximum time to wait in milliseconds before flushing the queue.
             Defaults to 5,000 (5 seconds)
         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"]).
+        default_tags (List[str], optional): Default tags for the sessions that can be used for grouping or sorting later (e.g. [\"GPT-4\"]).
         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.
-        skip_auto_end_session (optional, bool): Don't automatically end session based on your framework's decision-making
-            (i.e. Crew determining when tasks are complete and ending the session)
+        auto_start_session (bool): Whether to start an initial trace automatically when the client is initialized. Defaults to True via Config.
         env_data_opt_out (bool): Whether to opt out of collecting environment data.
         log_level (str, int): The log level to use for the client. Defaults to 'CRITICAL'.
-        fail_safe (bool): Whether to suppress errors and continue execution when possible.
-        exporter_endpoint (str, optional): Endpoint for the exporter. If none is provided, key will
-            be read from the AGENTOPS_EXPORTER_ENDPOINT environment variable.
-        **kwargs: Additional configuration parameters to be passed to the client.
+        exporter_endpoint (str, optional): Endpoint for the OTLP exporter. Defaults to AgentOps collector.
+        **kwargs: Additional configuration parameters to be passed to the client's Config object.
     """
     global _client
 
-    # Merge tags and default_tags if both are provided
-    merged_tags = None
-    if tags and default_tags:
-        merged_tags = list(set(tags + default_tags))
-    elif tags:
-        merged_tags = tags
-    elif default_tags:
-        merged_tags = default_tags
-
-    return _client.init(
+    # The client.init() method now handles its own configuration loading and defaults.
+    # We pass through the relevant parameters.
+    # client.init() returns None.
+    _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,
+        default_tags=default_tags,
         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,
     )
@@ -165,18 +148,65 @@ def configure(**kwargs):
     # 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: TraceContext, end_state: str = "Success") -> None:
+    """
+    Ends a trace (its root span) and finalizes it.
+
+    Args:
+        trace_context: The TraceContext object returned by start_trace.
+        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",