From 660be3b28f7eb51e5ba087345ca0408c4b59d631 Mon Sep 17 00:00:00 2001 From: michi-okahata Date: Tue, 3 Jun 2025 09:33:34 -0700 Subject: [PATCH 1/4] input/output guardrail decorator #1003 --- agentops/__init__.py | 4 +++- agentops/sdk/decorators/__init__.py | 4 +++- agentops/sdk/decorators/factory.py | 20 ++++++++++---------- agentops/sdk/decorators/utility.py | 8 ++++---- agentops/semconv/span_attributes.py | 4 ++++ agentops/semconv/span_kinds.py | 2 ++ 6 files changed, 26 insertions(+), 16 deletions(-) diff --git a/agentops/__init__.py b/agentops/__init__.py index 3b252759a..d769b1961 100755 --- a/agentops/__init__.py +++ b/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 @@ -247,4 +247,6 @@ def end_trace(trace_context: Optional[TraceContext] = None, end_state: str = "Su "task", "workflow", "operation", + "in_guardrail", + "out_guardrail", ] diff --git a/agentops/sdk/decorators/__init__.py b/agentops/sdk/decorators/__init__.py index f775b45d5..22667fa58 100644 --- a/agentops/sdk/decorators/__init__.py +++ b/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"] diff --git a/agentops/sdk/decorators/factory.py b/agentops/sdk/decorators/factory.py index cbf0e7026..f345c9ca7 100644 --- a/agentops/sdk/decorators/factory.py +++ b/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 diff --git a/agentops/sdk/decorators/utility.py b/agentops/sdk/decorators/utility.py index 5dee1d412..c0be4cf8b 100644 --- a/agentops/sdk/decorators/utility.py +++ b/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: diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index 67320bffd..a165df55a 100644 --- a/agentops/semconv/span_attributes.py +++ b/agentops/semconv/span_attributes.py @@ -89,6 +89,10 @@ class SpanAttributes: AGENTOPS_SPAN_KIND = "agentops.span.kind" AGENTOPS_ENTITY_NAME = "agentops.entity.name" + # Decorator + AGENTOPS_DECORATOR_INPUT = "agentops.{entity_kind}.input" + AGENTOPS_DECORATOR_OUTPUT = "agentops.{entity_kind}.output" + # Operation attributes OPERATION_NAME = "operation.name" OPERATION_VERSION = "operation.version" diff --git a/agentops/semconv/span_kinds.py b/agentops/semconv/span_kinds.py index 0d90a8cc9..bc8e806d2 100644 --- a/agentops/semconv/span_kinds.py +++ b/agentops/semconv/span_kinds.py @@ -27,6 +27,8 @@ class SpanKind: UNKNOWN = "unknown" CHAIN = "chain" TEXT = "text" + INPUT_GUARDRAIL = "input guardrail" + OUTPUT_GUARDRAIL = "output guardrail" class AgentOpsSpanKindValues(Enum): From 299db87946a7676fbdc6baf486b5ccf9e0f0370a Mon Sep 17 00:00:00 2001 From: michi-okahata Date: Thu, 5 Jun 2025 12:23:26 -0700 Subject: [PATCH 2/4] changed span_kinds name, confirmed no existing semconv for guardrails --- agentops/semconv/span_attributes.py | 2 -- agentops/semconv/span_kinds.py | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/agentops/semconv/span_attributes.py b/agentops/semconv/span_attributes.py index a165df55a..c2d1c64b8 100644 --- a/agentops/semconv/span_attributes.py +++ b/agentops/semconv/span_attributes.py @@ -88,8 +88,6 @@ class SpanAttributes: AGENTOPS_ENTITY_INPUT = "agentops.entity.input" AGENTOPS_SPAN_KIND = "agentops.span.kind" AGENTOPS_ENTITY_NAME = "agentops.entity.name" - - # Decorator AGENTOPS_DECORATOR_INPUT = "agentops.{entity_kind}.input" AGENTOPS_DECORATOR_OUTPUT = "agentops.{entity_kind}.output" diff --git a/agentops/semconv/span_kinds.py b/agentops/semconv/span_kinds.py index bc8e806d2..7e49d2b32 100644 --- a/agentops/semconv/span_kinds.py +++ b/agentops/semconv/span_kinds.py @@ -27,8 +27,8 @@ class SpanKind: UNKNOWN = "unknown" CHAIN = "chain" TEXT = "text" - INPUT_GUARDRAIL = "input guardrail" - OUTPUT_GUARDRAIL = "output guardrail" + INPUT_GUARDRAIL = "guardrail_input" + OUTPUT_GUARDRAIL = "guardrail_output" class AgentOpsSpanKindValues(Enum): From b50b7092e106a9d62a3501ad5816a59dc2382e49 Mon Sep 17 00:00:00 2001 From: michi-okahata Date: Thu, 5 Jun 2025 16:11:50 -0700 Subject: [PATCH 3/4] init() sets auto_start_session to false if in notebook --- agentops/__init__.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/agentops/__init__.py b/agentops/__init__.py index d769b1961..d4a0d0a75 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -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, From bfe1229ff3078c8f7b8988b8582ee5bf6f49210a Mon Sep 17 00:00:00 2001 From: michi-okahata Date: Thu, 5 Jun 2025 16:17:34 -0700 Subject: [PATCH 4/4] changed docs to reflect init() change and requirement for manual trace management --- docs/v2/usage/sdk-reference.mdx | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/v2/usage/sdk-reference.mdx b/docs/v2/usage/sdk-reference.mdx index de29cd30c..f11a6df3d 100644 --- a/docs/v2/usage/sdk-reference.mdx +++ b/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.