Traces/logs in Jupyter Notebooks

agentops/__init__.py

@@ -15,7 +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.sdk.decorators import trace, session, agent, task, workflow, operation, in_guardrail, out_guardrail
 
 from agentops.logging.config import logger
 
@@ -106,6 +106,13 @@ def init(
     elif default_tags:
         merged_tags = default_tags
 
+    # Check if in a Jupyter Notebook (manual start/end_trace())
+    try:
+        __IPYTHON__ # type: ignore
+        auto_start_session = False
+    except NameError:
+        auto_start_session = True
+
     return _client.init(
         api_key=api_key,
         endpoint=endpoint,
@@ -247,4 +254,6 @@ def end_trace(trace_context: Optional[TraceContext] = None, end_state: str = "Su
     "task",
     "workflow",
     "operation",
+    "in_guardrail",
+    "out_guardrail",
 ]

agentops/sdk/decorators/__init__.py

@@ -19,6 +19,8 @@
 session = create_entity_decorator(SpanKind.SESSION)
 tool = create_entity_decorator(SpanKind.TOOL)
 operation = task
+in_guardrail = create_entity_decorator(SpanKind.INPUT_GUARDRAIL)
+out_guardrail = create_entity_decorator(SpanKind.OUTPUT_GUARDRAIL)
 
 # For backward compatibility: @session decorator calls @trace decorator
 @functools.wraps(trace)
@@ -37,4 +39,4 @@ def session(*args, **kwargs):
 # For now, keeping the alias as it was, assuming it was intentional for `operation` to be `task`.
 operation = task
 
-__all__ = ["agent", "task", "workflow", "trace", "session", "operation", "tool"]
+__all__ = ["agent", "task", "workflow", "trace", "session", "operation", "tool", "in_guardrail", "out_guardrail"]

agentops/sdk/decorators/factory.py

@@ -48,7 +48,7 @@ def __init__(self, *args: Any, **kwargs: Any):
 
                     self._agentops_active_span = self._agentops_span_context_manager.__enter__()
                     try:
-                        _record_entity_input(self._agentops_active_span, args, kwargs)
+                        _record_entity_input(self._agentops_active_span, args, kwargs, entity_kind=entity_kind)
                     except Exception as e:
                         logger.warning(f"Failed to record entity input for class {op_name}: {e}")
                     super().__init__(*args, **kwargs)
@@ -64,7 +64,7 @@ async def __aenter__(self) -> "WrappedClass":
                 async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
                     if hasattr(self, "_agentops_active_span") and hasattr(self, "_agentops_span_context_manager"):
                         try:
-                            _record_entity_output(self._agentops_active_span, self)
+                            _record_entity_output(self._agentops_active_span, self, entity_kind=entity_kind)
                         except Exception as e:
                             logger.warning(f"Failed to record entity output for class instance: {e}")
                         self._agentops_span_context_manager.__exit__(exc_type, exc_val, exc_tb)
@@ -107,12 +107,12 @@ async def _wrapped_session_async() -> Any:
                                 )
                                 return await wrapped_func(*args, **kwargs)
                             try:
-                                _record_entity_input(trace_context.span, args, kwargs)
+                                _record_entity_input(trace_context.span, args, kwargs, entity_kind=entity_kind)
                             except Exception as e:
                                 logger.warning(f"Input recording failed for @trace '{operation_name}': {e}")
                             result = await wrapped_func(*args, **kwargs)
                             try:
-                                _record_entity_output(trace_context.span, result)
+                                _record_entity_output(trace_context.span, result, entity_kind=entity_kind)
                             except Exception as e:
                                 logger.warning(f"Output recording failed for @trace '{operation_name}': {e}")
                             TracingCore.get_instance().end_trace(trace_context, "Success")
@@ -139,12 +139,12 @@ async def _wrapped_session_async() -> Any:
                             )
                             return wrapped_func(*args, **kwargs)
                         try:
-                            _record_entity_input(trace_context.span, args, kwargs)
+                            _record_entity_input(trace_context.span, args, kwargs, entity_kind=entity_kind)
                         except Exception as e:
                             logger.warning(f"Input recording failed for @trace '{operation_name}': {e}")
                         result = wrapped_func(*args, **kwargs)
                         try:
-                            _record_entity_output(trace_context.span, result)
+                            _record_entity_output(trace_context.span, result, entity_kind=entity_kind)
                         except Exception as e:
                             logger.warning(f"Output recording failed for @trace '{operation_name}': {e}")
                         TracingCore.get_instance().end_trace(trace_context, "Success")
@@ -203,7 +203,7 @@ async def _wrapped_async() -> Any:
                         attributes={CoreAttributes.TAGS: tags} if tags else None,
                     ) as span:
                         try:
-                            _record_entity_input(span, args, kwargs)
+                            _record_entity_input(span, args, kwargs, entity_kind=entity_kind)
                             # Set cost attribute if tool
                             if entity_kind == "tool" and cost is not None:
                                 span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost)
@@ -212,7 +212,7 @@ async def _wrapped_async() -> Any:
                         try:
                             result = await wrapped_func(*args, **kwargs)
                             try:
-                                _record_entity_output(span, result)
+                                _record_entity_output(span, result, entity_kind=entity_kind)
                             except Exception as e:
                                 logger.warning(f"Output recording failed for '{operation_name}': {e}")
                             return result
@@ -230,7 +230,7 @@ async def _wrapped_async() -> Any:
                     attributes={CoreAttributes.TAGS: tags} if tags else None,
                 ) as span:
                     try:
-                        _record_entity_input(span, args, kwargs)
+                        _record_entity_input(span, args, kwargs, entity_kind=entity_kind)
                         # Set cost attribute if tool
                         if entity_kind == "tool" and cost is not None:
                             span.set_attribute(SpanAttributes.LLM_USAGE_TOOL_COST, cost)
@@ -239,7 +239,7 @@ async def _wrapped_async() -> Any:
                     try:
                         result = wrapped_func(*args, **kwargs)
                         try:
-                            _record_entity_output(span, result)
+                            _record_entity_output(span, result, entity_kind=entity_kind)
                         except Exception as e:
                             logger.warning(f"Output recording failed for '{operation_name}': {e}")
                         return result

agentops/sdk/decorators/utility.py

@@ -192,27 +192,27 @@ def _make_span(
     return span, ctx, token
 
 
-def _record_entity_input(span: trace.Span, args: tuple, kwargs: Dict[str, Any]) -> None:
+def _record_entity_input(span: trace.Span, args: tuple, kwargs: Dict[str, Any], entity_kind: str = "entity") -> None:
     """Record operation input parameters to span if content tracing is enabled"""
     try:
         input_data = {"args": args, "kwargs": kwargs}
         json_data = safe_serialize(input_data)
 
         if _check_content_size(json_data):
-            span.set_attribute(SpanAttributes.AGENTOPS_ENTITY_INPUT, json_data)
+            span.set_attribute(SpanAttributes.AGENTOPS_DECORATOR_INPUT.format(entity_kind=entity_kind), json_data)
         else:
             logger.debug("Operation input exceeds size limit, not recording")
     except Exception as err:
         logger.warning(f"Failed to serialize operation input: {err}")
 
 
-def _record_entity_output(span: trace.Span, result: Any) -> None:
+def _record_entity_output(span: trace.Span, result: Any, entity_kind: str = "entity") -> None:
     """Record operation output value to span if content tracing is enabled"""
     try:
         json_data = safe_serialize(result)
 
         if _check_content_size(json_data):
-            span.set_attribute(SpanAttributes.AGENTOPS_ENTITY_OUTPUT, json_data)
+            span.set_attribute(SpanAttributes.AGENTOPS_DECORATOR_OUTPUT.format(entity_kind=entity_kind), json_data)
         else:
             logger.debug("Operation output exceeds size limit, not recording")
     except Exception as err:

agentops/semconv/span_attributes.py

@@ -88,7 +88,9 @@ class SpanAttributes:
     AGENTOPS_ENTITY_INPUT = "agentops.entity.input"
     AGENTOPS_SPAN_KIND = "agentops.span.kind"
     AGENTOPS_ENTITY_NAME = "agentops.entity.name"
-
+    AGENTOPS_DECORATOR_INPUT = "agentops.{entity_kind}.input"
+    AGENTOPS_DECORATOR_OUTPUT = "agentops.{entity_kind}.output"
+
     # Operation attributes
     OPERATION_NAME = "operation.name"
     OPERATION_VERSION = "operation.version"

agentops/semconv/span_kinds.py

@@ -27,6 +27,8 @@ class SpanKind:
     UNKNOWN = "unknown"
     CHAIN = "chain"
     TEXT = "text"
+    INPUT_GUARDRAIL = "guardrail_input"
+    OUTPUT_GUARDRAIL = "guardrail_output"
 
 
 class AgentOpsSpanKindValues(Enum):

docs/v2/usage/sdk-reference.mdx

@@ -25,7 +25,7 @@ Initializes the AgentOps SDK and automatically starts tracking your application.
 - `default_tags` (List[str], optional): Default tags for the sessions that can be used for grouping or sorting later (e.g. ["GPT-4"]).
 - `tags` (List[str], optional): [Deprecated] Use `default_tags` instead.
 - `instrument_llm_calls` (bool, optional): Whether to instrument LLM calls automatically. Defaults to True.
-- `auto_start_session` (bool, optional): Whether to start a session automatically when the client is created. Defaults to True.
+- `auto_start_session` (bool, optional): Whether to start a session automatically when the client is created. Set to False if running in a Jupyter Notebook. Defaults to True.
 - `auto_init` (bool, optional): Whether to automatically initialize the client on import. Defaults to True.
 - `skip_auto_end_session` (bool, optional): Don't automatically end session based on your framework's decision-making. Defaults to False.
 - `env_data_opt_out` (bool, optional): Whether to opt out of collecting environment data. Defaults to False.
@@ -109,6 +109,8 @@ These functions help you manage the lifecycle of tracking traces.
 
 Starts a new AgentOps trace manually. This is useful when you've disabled automatic session creation or need multiple separate traces.
 
+Manually managing traces is required when running in a Jupyter Notebook as there is no end state.
+
 **Parameters**:
 
 - `trace_name` (str, optional): Name for the trace. If not provided, a default name will be used.
@@ -134,6 +136,8 @@ trace = agentops.start_trace("customer-service-workflow", tags=["customer-query"
 
 Ends a specific trace or all active traces.
 
+Manually managing traces is required when running in a Jupyter Notebook as there is no end state.
+
 **Parameters**:
 
 - `trace` (TraceContext, optional): The specific trace to end. If not provided, all active traces will be ended.