Enhance tracing functionality by introducing new decorators and improving session handling. Added trace, session, agent, task, workflow, and operation decorators for better instrumentation. Updated log_trace_url to include titles for improved logging context. Refactored Client initialization trace name and adjusted end trace state handling. Improved error handling during trace logging in TracingCore and removed deprecated session decorator usage.

Author: Dwij1704

agentops/__init__.py

@@ -15,6 +15,7 @@
 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
 
@@ -207,6 +208,12 @@ def end_trace(trace_context: TraceContext, end_state: str = "Success") -> None:
     "record",
     "start_trace",
     "end_trace",
+    "trace",
+    "session",
+    "agent",
+    "task",
+    "workflow",
+    "operation",
     "start_session",
     "end_session",
     "track_agent",

agentops/client/client.py

@@ -27,7 +27,7 @@ def _end_init_trace_atexit():
             # Use TracingCore to end the trace directly
             tracing_core = TracingCore.get_instance()
             if tracing_core.initialized and _client_init_trace_context.span.is_recording():
-                tracing_core.end_trace(_client_init_trace_context, end_state="Aborted_AtExit")
+                tracing_core.end_trace(_client_init_trace_context, end_state="Shutdown")
         except Exception as e:
             logger.warning(f"Error ending client's init trace during shutdown: {e}")
         finally:
@@ -129,7 +129,7 @@ def init(self, **kwargs) -> None:  # Return type updated to None
             if self._init_trace_context is None or not self._init_trace_context.span.is_recording():
                 logger.debug("Auto-starting init trace.")
                 self._init_trace_context = tracing_core.start_trace(
-                    trace_name="agentops_init_trace",
+                    trace_name="default",
                     tags=list(self.config.default_tags) if self.config.default_tags else None,
                     is_init_trace=True,
                 )

agentops/helpers/dashboard.py

@@ -2,7 +2,7 @@
 Helpers for interacting with the AgentOps dashboard.
 """
 
-from typing import Union
+from typing import Union, Optional
 from termcolor import colored
 from opentelemetry.sdk.trace import Span, ReadableSpan
 from agentops.logging import logger
@@ -33,12 +33,12 @@ def get_trace_url(span: Union[Span, ReadableSpan]) -> str:
     return f"{app_url}/sessions?trace_id={trace_id}"
 
 
-def log_trace_url(span: Union[Span, ReadableSpan]) -> None:
+def log_trace_url(span: Union[Span, ReadableSpan], title: Optional[str] = None) -> None:
     """
     Log the trace URL for the AgentOps dashboard.
 
     Args:
         span: The span to log the URL for.
     """
     session_url = get_trace_url(span)
-    logger.info(colored(f"\x1b[34mSession Replay: {session_url}\x1b[0m", "blue"))
+    logger.info(colored(f"\x1b[34mSession Replay for {title} trace: {session_url}\x1b[0m", "blue"))

agentops/sdk/core.py

@@ -23,6 +23,7 @@
 from agentops.sdk.processors import InternalSpanProcessor
 from agentops.sdk.types import TracingConfig
 from agentops.semconv import ResourceAttributes, SpanKind, SpanAttributes
+from agentops.helpers.dashboard import log_trace_url
 
 # No need to create shortcuts since we're using our own ResourceAttributes class now
 
@@ -431,6 +432,13 @@ def start_trace(
         # It returns: span, context_object, context_token
         span, _, context_token = _make_span(trace_name, span_kind=SpanKind.SESSION, attributes=attributes)
         logger.debug(f"Trace '{trace_name}' started with span ID: {span.get_span_context().span_id}")
+
+        # Log the session replay URL for this new trace
+        try:
+            log_trace_url(span, title=trace_name)
+        except Exception as e:
+            logger.warning(f"Failed to log trace URL for '{trace_name}': {e}")
+
         return TraceContext(span, token=context_token, is_init_trace=is_init_trace)
 
     def end_trace(self, trace_context: TraceContext, end_state: str = "Success") -> None:
@@ -458,9 +466,17 @@ def end_trace(self, trace_context: TraceContext, end_state: str = "Success") ->
 
         try:
             span.set_attribute(SpanAttributes.AGENTOPS_SESSION_END_STATE, end_state)
-            # _finalize_span ends the span and detaches it from context if a token is provided
             _finalize_span(span, token=token)
             # For root spans (traces), we might want an immediate flush after they end.
             self._flush_span_processors()
+
+            # Log the session replay URL again after the trace has ended
+            # The span object should still contain the necessary context (trace_id)
+            try:
+                # Use span.name as the title, which should reflect the original trace_name
+                log_trace_url(span, title=span.name)
+            except Exception as e:
+                logger.warning(f"Failed to log trace URL after ending trace '{span.name}': {e}")
+
         except Exception as e:
             logger.error(f"Error ending trace: {e}", exc_info=True)

agentops/sdk/decorators/__init__.py

@@ -5,17 +5,38 @@
 and methods with appropriate span kinds. Decorators can be used with or without parentheses.
 """
 
+import functools
+from termcolor import colored
+
+from agentops.logging import logger
 from agentops.sdk.decorators.factory import create_entity_decorator
 from agentops.semconv.span_kinds import SpanKind
 
 # Create decorators for specific entity types using the factory
 agent = create_entity_decorator(SpanKind.AGENT)
 task = create_entity_decorator(SpanKind.TASK)
-operation = create_entity_decorator(SpanKind.OPERATION)
+operation_decorator = create_entity_decorator(SpanKind.OPERATION)
 workflow = create_entity_decorator(SpanKind.WORKFLOW)
-session = create_entity_decorator(SpanKind.SESSION)
+trace = create_entity_decorator(SpanKind.SESSION)
+
+
+# For backward compatibility: @session decorator calls @trace decorator
+@functools.wraps(trace)
+def session(*args, **kwargs):
+    # yellow color
+    logger.info(colored("@agentops.session decorator is deprecated. Please use @agentops.trace instead.", "yellow"))
+    # If called as @session or @session(...)
+    if not args or not callable(args[0]):  # called with kwargs like @session(name=...)
+        return trace(*args, **kwargs)
+    else:  # called as @session directly on a function
+        return trace(args[0], **kwargs)  # args[0] is the wrapped function
+
+
+# Note: The original `operation = task` was potentially problematic if `operation` was meant to be distinct.
+# Using operation_decorator for clarity if a distinct OPERATION kind decorator is needed.
+# For now, keeping the alias as it was, assuming it was intentional for `operation` to be `task`.
 operation = task
 
-__all__ = ["agent", "task", "workflow", "session", "operation"]
+__all__ = ["agent", "task", "workflow", "trace", "session", "operation"]
 
 # Create decorators task, workflow, session, agent

agentops/sdk/decorators/factory.py

@@ -6,6 +6,7 @@
 
 from agentops.logging import logger
 from agentops.sdk.core import TracingCore
+from agentops.semconv.span_kinds import SpanKind
 
 from .utility import (
     _create_as_current_span,
@@ -28,10 +29,10 @@ def create_entity_decorator(entity_kind: str):
         A decorator with optional arguments for name and version
     """
 
-    def decorator(wrapped=None, *, name=None, version=None):
+    def decorator(wrapped=None, *, name=None, version=None, tags=None):
         # Handle case where decorator is called with parameters
         if wrapped is None:
-            return functools.partial(decorator, name=name, version=version)
+            return functools.partial(decorator, name=name, version=version, tags=tags)
 
         # Handle class decoration
         if inspect.isclass(wrapped):
@@ -91,7 +92,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         @wrapt.decorator
         def wrapper(wrapped, instance, args, kwargs):
             # Skip instrumentation if tracer not initialized
-            if not TracingCore.get_instance()._initialized:
+            if not TracingCore.get_instance().initialized:
                 return wrapped(*args, **kwargs)
 
             # Use provided name or function name
@@ -102,8 +103,100 @@ def wrapper(wrapped, instance, args, kwargs):
             is_generator = inspect.isgeneratorfunction(wrapped)
             is_async_generator = inspect.isasyncgenfunction(wrapped)
 
+            # If it's a SESSION kind, we use start_trace/end_trace
+            if entity_kind == SpanKind.SESSION:
+                if is_generator or is_async_generator:
+                    # Using start_trace/end_trace for generators decorated with @session might be complex
+                    # due to the nature of yielding. For now, log a warning and fall back to existing generator handling OR disallow.
+                    # Let's keep existing generator handling for now, which creates a single span.
+                    # A true "session per generator invocation" would require more complex handling.
+                    logger.warning(
+                        f"@agentops.session decorator used on a generator function '{operation_name}'. \
+This will create a single span for the generator's instantiation, not a long-running trace for its entire execution."
+                    )
+                    # Fallthrough to existing generator logic below for non-session spans
+                    pass  # Explicitly fall through
+
+                elif is_async:
+
+                    async def _wrapped_session_async():
+                        trace_context = None
+                        try:
+                            trace_context = TracingCore.get_instance().start_trace(trace_name=operation_name, tags=tags)
+                            if not trace_context:
+                                logger.error(
+                                    f"Failed to start trace for @session '{operation_name}'. Executing function without AgentOps trace."
+                                )
+                                return await wrapped(*args, **kwargs)
+
+                            # Record input if possible (span is in trace_context.span)
+                            try:
+                                _record_entity_input(trace_context.span, args, kwargs)
+                            except Exception as e:
+                                logger.warning(f"Failed to record entity input for @session '{operation_name}': {e}")
+
+                            result = await wrapped(*args, **kwargs)
+
+                            try:
+                                _record_entity_output(trace_context.span, result)
+                            except Exception as e:
+                                logger.warning(f"Failed to record entity output for @session '{operation_name}': {e}")
+
+                            TracingCore.get_instance().end_trace(trace_context, "Success")
+                            return result
+                        except Exception:
+                            if trace_context:
+                                TracingCore.get_instance().end_trace(trace_context, "Failure")
+                            # record_exception on trace_context.span might be an option too
+                            # trace_context.span.record_exception(e) # If we want it on the span directly
+                            raise
+                        finally:
+                            # Ensure trace is ended if not already (e.g. early exit without exception but before success end_trace)
+                            if trace_context and trace_context.span.is_recording():
+                                logger.warning(
+                                    f"Trace for @session '{operation_name}' was not explicitly ended. Ending as 'Unknown'."
+                                )
+                                TracingCore.get_instance().end_trace(trace_context, "Unknown")
+
+                    return _wrapped_session_async()
+                else:  # Sync function for SpanKind.SESSION
+                    trace_context = None
+                    try:
+                        trace_context = TracingCore.get_instance().start_trace(trace_name=operation_name, tags=tags)
+                        if not trace_context:
+                            logger.error(
+                                f"Failed to start trace for @session '{operation_name}'. Executing function without AgentOps trace."
+                            )
+                            return wrapped(*args, **kwargs)
+
+                        try:
+                            _record_entity_input(trace_context.span, args, kwargs)
+                        except Exception as e:
+                            logger.warning(f"Failed to record entity input for @session '{operation_name}': {e}")
+
+                        result = wrapped(*args, **kwargs)
+
+                        try:
+                            _record_entity_output(trace_context.span, result)
+                        except Exception as e:
+                            logger.warning(f"Failed to record entity output for @session '{operation_name}': {e}")
+
+                        TracingCore.get_instance().end_trace(trace_context, "Success")
+                        return result
+                    except Exception:
+                        if trace_context:
+                            TracingCore.get_instance().end_trace(trace_context, "Failure")
+                        raise
+                    finally:
+                        if trace_context and trace_context.span.is_recording():
+                            logger.warning(
+                                f"Trace for @session '{operation_name}' was not explicitly ended. Ending as 'Unknown'."
+                            )
+                            TracingCore.get_instance().end_trace(trace_context, "Unknown")
+
+            # Existing logic for non-SESSION kinds or generators under @session (as per above warning)
             # Handle generator functions
-            if is_generator:
+            if is_generator:  # This 'if' will also catch generators decorated with @session due to fallthrough
                 # Use the old approach for generators
                 span, ctx, token = _make_span(operation_name, entity_kind, version)
                 try:
@@ -115,7 +208,7 @@ def wrapper(wrapped, instance, args, kwargs):
                 return _process_sync_generator(span, result)
 
             # Handle async generator functions
-            elif is_async_generator:
+            elif is_async_generator:  # This 'elif' will also catch async generators decorated with @session
                 # Use the old approach for async generators
                 span, ctx, token = _make_span(operation_name, entity_kind, version)
                 try:
@@ -126,8 +219,8 @@ def wrapper(wrapped, instance, args, kwargs):
                 result = wrapped(*args, **kwargs)
                 return _process_async_generator(span, token, result)
 
-            # Handle async functions
-            elif is_async:
+            # Handle async functions (non-SESSION)
+            elif is_async:  # This is for entity_kind != SpanKind.SESSION
 
                 async def _wrapped_async():
                     with _create_as_current_span(operation_name, entity_kind, version) as span:
@@ -149,8 +242,8 @@ async def _wrapped_async():
 
                 return _wrapped_async()
 
-            # Handle sync functions
-            else:
+            # Handle sync functions (non-SESSION)
+            else:  # This is for entity_kind != SpanKind.SESSION
                 with _create_as_current_span(operation_name, entity_kind, version) as span:
                     try:
                         _record_entity_input(span, args, kwargs)

agentops/sdk/processors.py

@@ -13,7 +13,6 @@
 from opentelemetry.sdk.trace.export import SpanExporter
 
 from agentops.logging import logger
-from agentops.helpers.dashboard import log_trace_url
 from agentops.semconv.core import CoreAttributes
 from agentops.logging import upload_logfile
 
@@ -108,7 +107,6 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
         if not self._root_span_id:
             self._root_span_id = span.context.span_id
             logger.debug(f"[agentops.InternalSpanProcessor] Found root span: {span.name}")
-            log_trace_url(span)
 
     def on_end(self, span: ReadableSpan) -> None:
         """
@@ -123,7 +121,6 @@ def on_end(self, span: ReadableSpan) -> None:
 
         if self._root_span_id and (span.context.span_id is self._root_span_id):
             logger.debug(f"[agentops.InternalSpanProcessor] Ending root span: {span.name}")
-            log_trace_url(span)
             try:
                 upload_logfile(span.context.trace_id)
             except Exception as e: