agentreplay 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

agentreplay/__init__.py

@@ -12,7 +12,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Agentreplay Python SDK - Agent Trace Engine for LLM Agents."""
+"""
+Agentreplay Python SDK - Agent Trace Engine for LLM Agents.
+
+Modern, ergonomic SDK for observability in LLM applications.
+
+Quick Start:
+    >>> import agentreplay
+    >>> 
+    >>> # Initialize with env vars (recommended)
+    >>> agentreplay.init()
+    >>> 
+    >>> # Or with explicit config
+    >>> agentreplay.init(
+    ...     api_key="your-key",
+    ...     project_id="my-project",
+    ... )
+    >>> 
+    >>> # Trace functions with decorator
+    >>> @agentreplay.traceable
+    >>> def my_llm_function(query: str) -> str:
+    ...     return call_llm(query)
+    >>> 
+    >>> # Wrap OpenAI for automatic tracing
+    >>> from openai import OpenAI
+    >>> client = agentreplay.wrap_openai(OpenAI())
+    >>> 
+    >>> # Ensure traces are sent before exit
+    >>> agentreplay.flush()
+"""
 
 from agentreplay.client import AgentreplayClient
 from agentreplay.models import SpanType, AgentFlowEdge
@@ -32,44 +60,158 @@ from agentreplay.exceptions import (
 )
 
 # Agent Context Tracking
-from agentreplay.context import AgentContext
+from agentreplay.context import (
+    AgentContext,
+    set_context,
+    get_global_context,
+    clear_context,
+    with_context,
+)
+
+# Auto-instrumentation (Pure OpenTelemetry) - Optional
+try:
+    from agentreplay.auto_instrument import auto_instrument, setup_instrumentation
+except ImportError:
+    auto_instrument = None  # type: ignore
+    setup_instrumentation = None  # type: ignore
+
+# OTEL Bridge & Bootstrap - Optional (requires opentelemetry-sdk)
+try:
+    from agentreplay.bootstrap import init_otel_instrumentation, is_initialized
+    from agentreplay.otel_bridge import get_tracer
+except ImportError:
+    init_otel_instrumentation = None  # type: ignore
+    is_initialized = lambda: False  # type: ignore
+    get_tracer = None  # type: ignore
+
+# =============================================================================
+# Ergonomic Top-Level API (v0.4+)
+# =============================================================================
 
-# Auto-instrumentation (Pure OpenTelemetry)
-from agentreplay.auto_instrument import auto_instrument, setup_instrumentation
+# SDK Lifecycle
+from agentreplay.sdk import (
+    init,
+    get_client,
+    get_batching_client,
+    flush,
+    shutdown,
+    reset,
+    get_stats,
+    ping,
+)
+
+# Decorators & Tracing
+from agentreplay.decorators import (
+    traceable,
+    observe,  # Langfuse-style alias
+    trace,
+    start_span,
+    get_current_span,
+    SpanKind,
+    ActiveSpan,
+)
 
-# OTEL Bridge & Bootstrap
-from agentreplay.bootstrap import init_otel_instrumentation, is_initialized
-from agentreplay.otel_bridge import get_tracer
+# Client Wrappers
+from agentreplay.wrappers import (
+    wrap_openai,
+    wrap_anthropic,
+    wrap_method,
+)
+
+# Privacy
+from agentreplay.privacy import (
+    configure_privacy,
+    redact_payload,
+    redact_string,
+    hash_pii,
+    add_pattern,
+    add_scrub_path,
+    privacy_context,
+)
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 __all__ = [
+    # ==========================================================================
+    # MODERN ERGONOMIC API (Recommended)
+    # ==========================================================================
+    
+    # Initialization (one-liner setup)
+    "init",
+    "flush",
+    "shutdown",
+    "reset",
+    "get_stats",
+    "ping",
+    
+    # Tracing decorators
+    "traceable",
+    "observe",  # Langfuse-style alias
+    "trace",
+    "start_span",
+    "get_current_span",
+    "SpanKind",
+    "ActiveSpan",
+    
+    # Client wrappers (one-liner instrumentation)
+    "wrap_openai",
+    "wrap_anthropic",
+    "wrap_method",
+    
+    # Context management
+    "set_context",
+    "get_global_context",
+    "clear_context",
+    "with_context",
+    "AgentContext",  # Class-based context
+    
+    # Privacy
+    "configure_privacy",
+    "redact_payload",
+    "redact_string",
+    "hash_pii",
+    "add_pattern",
+    "add_scrub_path",
+    "privacy_context",
+    
+    # ==========================================================================
+    # CORE API (Advanced Usage)
+    # ==========================================================================
+    
     # Core client
+    "get_client",
+    "get_batching_client",
     "AgentreplayClient",
     "BatchingAgentreplayClient",
+    
     # Models
     "SpanType",
     "AgentFlowEdge",
     "Span",
+    
     # Configuration
     "AgentreplayConfig",
     "get_config",
     "set_config",
     "reset_config",
+    
     # Session management
     "Session",
+    
     # Retry utilities
     "retry_with_backoff",
-    # Agent Context
-    "AgentContext",
+    
     # Auto-instrumentation (Pure OpenTelemetry)
     "auto_instrument",
     "setup_instrumentation",
+    
     # OTEL Initialization
     "init_otel_instrumentation",
     "is_initialized",
+    
     # OTEL Bridge
     "get_tracer",
+    
     # Exceptions
     "AgentreplayError",
     "AuthenticationError",

agentreplay/context.py

@@ -166,3 +166,136 @@ def set_user_id(user_id: str):
         Token that can be used to reset the context
     """
     return _user_id.set(user_id)
+
+
+# =============================================================================
+# Ergonomic set_context Function
+# =============================================================================
+
+# Global context for all spans (merged with span-level context)
+_global_context: dict = {}
+
+
+def set_context(
+    *,
+    user_id: Optional[str] = None,
+    session_id: Optional[str] = None,
+    agent_id: Optional[str] = None,
+    workflow_id: Optional[str] = None,
+    **extra: str,
+) -> None:
+    """Set global context that applies to all subsequent spans.
+    
+    This is a convenience function for setting context that should
+    apply to all traces/spans. The context persists until reset.
+    
+    For request-scoped context, use the `AgentContext` context manager
+    or pass context directly to `@traceable` or `trace()`.
+    
+    Args:
+        user_id: User identifier
+        session_id: Session identifier
+        agent_id: Agent identifier
+        workflow_id: Workflow identifier
+        **extra: Additional key-value pairs
+        
+    Example:
+        >>> from agentreplay import init, set_context
+        >>> 
+        >>> init()
+        >>> set_context(user_id="user-123", session_id="session-456")
+        >>> 
+        >>> # All subsequent traces include this context
+        >>> @traceable
+        >>> def my_function():
+        ...     return "hello"
+    """
+    global _global_context
+    
+    if user_id is not None:
+        _global_context["user_id"] = user_id
+        set_user_id(user_id)
+    
+    if session_id is not None:
+        _global_context["session_id"] = session_id
+        set_session_id(session_id)
+    
+    if agent_id is not None:
+        _global_context["agent_id"] = agent_id
+        set_agent_id(agent_id)
+    
+    if workflow_id is not None:
+        _global_context["workflow_id"] = workflow_id
+        set_workflow_id(workflow_id)
+    
+    # Store extra context
+    _global_context.update(extra)
+
+
+def get_global_context() -> dict:
+    """Get the current global context.
+    
+    Returns:
+        Dictionary of global context key-value pairs
+    """
+    return _global_context.copy()
+
+
+def clear_context() -> None:
+    """Clear all global context."""
+    global _global_context
+    _global_context = {}
+
+
+def with_context(**context) -> "ContextScope":
+    """Create a context scope for the current async context.
+    
+    This is like `set_context` but only applies within the current
+    async task or thread, and automatically cleans up when the scope exits.
+    
+    Args:
+        **context: Key-value pairs to set in context
+        
+    Returns:
+        ContextScope context manager
+        
+    Example:
+        >>> async def handle_request(user_id: str):
+        ...     with with_context(user_id=user_id, request_id="req-123"):
+        ...         # Context only applies within this block
+        ...         result = await my_traced_function()
+        ...     return result
+    """
+    return ContextScope(**context)
+
+
+class ContextScope:
+    """Context manager for request-scoped context.
+    
+    Automatically sets and resets context variables.
+    """
+    
+    def __init__(self, **context):
+        self.context = context
+        self.tokens = []
+    
+    def __enter__(self):
+        if "user_id" in self.context:
+            self.tokens.append(_user_id.set(self.context["user_id"]))
+        if "session_id" in self.context:
+            self.tokens.append(_session_id.set(self.context["session_id"]))
+        if "agent_id" in self.context:
+            self.tokens.append(_agent_id.set(self.context["agent_id"]))
+        if "workflow_id" in self.context:
+            self.tokens.append(_workflow_id.set(self.context["workflow_id"]))
+        return self
+    
+    def __exit__(self, *args):
+        for token in reversed(self.tokens):
+            token.var.reset(token)
+    
+    async def __aenter__(self):
+        return self.__enter__()
+    
+    async def __aexit__(self, *args):
+        return self.__exit__(*args)