agentbasis 0.1.0__py3-none-any.whl

agentbasis/__init__.py

@@ -0,0 +1,87 @@
+from typing import Optional
+from .client import AgentBasis
+from .decorators import trace
+from .context import (
+    context,
+    set_user,
+    set_session,
+    set_conversation,
+    set_metadata,
+    with_context,
+    AgentBasisContext,
+)
+
+
+def init(api_key: Optional[str] = None, agent_id: Optional[str] = None) -> AgentBasis:
+    """
+    Initialize the AgentBasis SDK.
+    
+    Args:
+        api_key: Your AgentBasis API Key. If not provided, reads from AGENTBASIS_API_KEY env var.
+        agent_id: The ID of the agent to track. If not provided, reads from AGENTBASIS_AGENT_ID env var.
+    Returns:
+        The initialized AgentBasis client instance.
+    """
+    return AgentBasis.initialize(api_key=api_key, agent_id=agent_id)
+
+
+def flush(timeout_millis: int = 30000) -> bool:
+    """
+    Force flush all pending telemetry data.
+    
+    This is useful when you want to ensure all traces are sent before
+    a critical operation, at specific checkpoints, or before exiting.
+    
+    Note: The SDK automatically flushes on normal Python exit via atexit.
+    
+    Args:
+        timeout_millis: Maximum time to wait for flush (default 30 seconds).
+        
+    Returns:
+        True if flush completed successfully, False if timed out or not initialized.
+        
+    Example:
+        >>> agentbasis.init(api_key="...", agent_id="...")
+        >>> # ... your agent code ...
+        >>> agentbasis.flush()  # Ensure all data is sent
+    """
+    try:
+        client = AgentBasis.get_instance()
+        return client.flush(timeout_millis)
+    except RuntimeError:
+        # SDK not initialized
+        return False
+
+
+def shutdown():
+    """
+    Manually shut down the SDK and flush all pending data.
+    
+    This is automatically called on Python exit, but can be called
+    manually if you need to shut down the SDK before the process ends.
+    
+    This method is idempotent - calling it multiple times is safe.
+    """
+    try:
+        client = AgentBasis.get_instance()
+        client.shutdown()
+    except RuntimeError:
+        # SDK not initialized
+        pass
+
+
+__all__ = [
+    "init",
+    "AgentBasis",
+    "trace",
+    "flush",
+    "shutdown",
+    # Context management
+    "context",
+    "set_user",
+    "set_session",
+    "set_conversation",
+    "set_metadata",
+    "with_context",
+    "AgentBasisContext",
+]

agentbasis/client.py

@@ -0,0 +1,134 @@
+from typing import Optional
+import atexit
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+
+from .config import Config
+
+
+class AgentBasis:
+    """
+    The main AgentBasis client.
+    Manages OpenTelemetry configuration and data transmission.
+    """
+    _instance: Optional['AgentBasis'] = None
+    _shutdown_registered: bool = False
+
+    def __init__(self, config: Config):
+        self.config = config
+        self._is_shutdown = False
+        
+        # 1. Create Resource (Metadata about who is sending data)
+        attributes = {
+            "service.name": "agentbasis-python-sdk",
+            # We can add more metadata here like environment
+        }
+        
+        # Add agent_id if present (it should be, as Config validates it)
+        if config.agent_id:
+            attributes["service.instance.id"] = config.agent_id
+            # Also adding a custom attribute just in case we want to query by it explicitly later
+            attributes["agentbasis.agent.id"] = config.agent_id
+
+        resource = Resource.create(attributes=attributes)
+
+        # 2. Initialize Tracer Provider
+        self.tracer_provider = TracerProvider(resource=resource)
+
+        # 3. Configure Exporter
+        endpoint = f"{config.api_url}/api/v1/traces" 
+        exporter = OTLPSpanExporter(
+            endpoint=endpoint,
+            headers={"x-api-key": config.api_key}
+        )
+
+        # 4. Add Batch Processor (Background thread for sending)
+        self._processor = BatchSpanProcessor(exporter)
+        self.tracer_provider.add_span_processor(self._processor)
+
+        # 5. Register as Global Tracer
+        # This allows trace.get_tracer(__name__) to work anywhere in the user's code
+        trace.set_tracer_provider(self.tracer_provider)
+        
+        # 6. Register atexit handler for graceful shutdown
+        self._register_atexit()
+
+    def _register_atexit(self):
+        """
+        Register the shutdown handler to run when Python exits.
+        Only registers once to avoid duplicate handlers.
+        """
+        if not AgentBasis._shutdown_registered:
+            atexit.register(self._atexit_handler)
+            AgentBasis._shutdown_registered = True
+
+    def _atexit_handler(self):
+        """
+        Handler called when Python exits. Flushes and shuts down gracefully.
+        """
+        self.shutdown()
+
+    @classmethod
+    def initialize(cls, api_key: Optional[str] = None, agent_id: Optional[str] = None) -> 'AgentBasis':
+        """
+        Initializes the global AgentBasis client.
+        """
+        config = Config(api_key=api_key, agent_id=agent_id)
+        config.validate()
+        
+        cls._instance = cls(config)
+        return cls._instance
+
+    @classmethod
+    def get_instance(cls) -> 'AgentBasis':
+        """
+        Returns the global AgentBasis client instance.
+        """
+        if cls._instance is None:
+            raise RuntimeError(
+                "AgentBasis is not initialized. "
+                "Please call `agentbasis.init(api_key='...', agent_id='...')` first."
+            )
+        return cls._instance
+
+    def flush(self, timeout_millis: int = 30000) -> bool:
+        """
+        Forces a flush of all pending spans.
+        
+        This is useful when you want to ensure all telemetry is sent before
+        a critical operation, or at specific checkpoints in your application.
+        
+        Args:
+            timeout_millis: Maximum time to wait for flush to complete (default 30 seconds).
+            
+        Returns:
+            True if flush completed successfully, False if timed out.
+        """
+        if self._is_shutdown:
+            return False
+            
+        if self.tracer_provider and hasattr(self.tracer_provider, 'force_flush'):
+            return self.tracer_provider.force_flush(timeout_millis)
+        return True
+
+    def shutdown(self):
+        """
+        Flushes remaining spans and shuts down the provider.
+        
+        This method is idempotent - calling it multiple times is safe.
+        It's automatically called when Python exits via atexit.
+        """
+        if self._is_shutdown:
+            return
+            
+        self._is_shutdown = True
+        
+        if self.tracer_provider:
+            try:
+                self.tracer_provider.shutdown()
+            except Exception:
+                # Silently ignore shutdown errors to avoid noise during exit
+                pass

agentbasis/config.py

@@ -0,0 +1,33 @@
+import os
+from typing import Optional
+
+class Config:
+    """
+    Configuration settings for the AgentBasis SDK.
+    Handles API keys and Agent IDs.
+    """
+    def __init__(self, api_key: Optional[str] = None, agent_id: Optional[str] = None):
+        self.api_key = api_key or os.environ.get("AGENTBASIS_API_KEY")
+        self.agent_id = agent_id or os.environ.get("AGENTBASIS_AGENT_ID")
+        
+        # Backend endpoint - not user-configurable, but can be overridden via env var for testing
+        self.api_url = os.environ.get("AGENTBASIS_API_URL", "https://api.agentbasis.co")
+
+    def validate(self):
+        """
+        Checks if the configuration is valid (i.e., has an API key and Agent ID).
+        Raises a ValueError if the key is missing.
+        """
+        if not self.api_key:
+            raise ValueError(
+                "AgentBasis API Key is missing. "
+                "Please provide it via `agentbasis.init(api_key='...')` "
+                "or set the `AGENTBASIS_API_KEY` environment variable."
+            )
+            
+        if not self.agent_id:
+            raise ValueError(
+                "AgentBasis Agent ID is missing. "
+                "Please provide it via `agentbasis.init(agent_id='...')` "
+                "or set the `AGENTBASIS_AGENT_ID` environment variable."
+            )

agentbasis/context.py

@@ -0,0 +1,259 @@
+"""
+Context management for AgentBasis.
+
+This module provides context tracking for user sessions, allowing developers
+to associate traces with specific users, sessions, and conversations.
+"""
+
+from contextvars import ContextVar
+from contextlib import contextmanager
+from typing import Optional, Dict, Any
+import functools
+import json
+
+
+# Context variables for storing user/session info
+# These are thread-safe and async-safe
+_user_id: ContextVar[Optional[str]] = ContextVar('user_id', default=None)
+_session_id: ContextVar[Optional[str]] = ContextVar('session_id', default=None)
+_conversation_id: ContextVar[Optional[str]] = ContextVar('conversation_id', default=None)
+_metadata: ContextVar[Optional[Dict[str, Any]]] = ContextVar('metadata', default=None)
+
+
+class AgentBasisContext:
+    """
+    Context manager for setting user/session context.
+    
+    All spans created within this context will automatically include
+    the user_id, session_id, and other context attributes.
+    
+    Example:
+        with AgentBasisContext(user_id="user-123", session_id="sess-456"):
+            response = client.chat.completions.create(...)
+            # This span will have user_id and session_id attributes
+    """ 
+    
+    def __init__(
+        self,
+        user_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        self.user_id = user_id
+        self.session_id = session_id
+        self.conversation_id = conversation_id
+        self.metadata = metadata
+        
+        # Store tokens to reset on exit
+        self._tokens = []
+    
+    def __enter__(self):
+        # Set context variables and store tokens for cleanup
+        if self.user_id is not None:
+            self._tokens.append(('user_id', _user_id.set(self.user_id)))
+        if self.session_id is not None:
+            self._tokens.append(('session_id', _session_id.set(self.session_id)))
+        if self.conversation_id is not None:
+            self._tokens.append(('conversation_id', _conversation_id.set(self.conversation_id)))
+        if self.metadata is not None:
+            self._tokens.append(('metadata', _metadata.set(self.metadata)))
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Reset context variables to their previous values
+        for name, token in self._tokens:
+            if name == 'user_id':
+                _user_id.reset(token)
+            elif name == 'session_id':
+                _session_id.reset(token)
+            elif name == 'conversation_id':
+                _conversation_id.reset(token)
+            elif name == 'metadata':
+                _metadata.reset(token)
+        return False  # Don't suppress exceptions
+
+
+# Convenience function (alias for AgentBasisContext)
+@contextmanager
+def context(
+    user_id: Optional[str] = None,
+    session_id: Optional[str] = None,
+    conversation_id: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None
+):
+    """
+    Context manager for setting user/session context.
+    
+    Example:
+        with agentbasis.context(user_id="user-123"):
+            response = client.chat.completions.create(...)
+    """
+    with AgentBasisContext(
+        user_id=user_id,
+        session_id=session_id,
+        conversation_id=conversation_id,
+        metadata=metadata
+    ):
+        yield
+
+
+# Global setters for simpler usage
+def set_user(user_id: Optional[str]):
+    """
+    Set the current user ID globally.
+    
+    This affects all subsequent spans until changed or cleared.
+    For scoped context, use the context() context manager instead.
+    
+    Example:
+        agentbasis.set_user("user-123")
+        response = client.chat.completions.create(...)  # Has user_id
+        agentbasis.set_user(None)  # Clear
+    """
+    _user_id.set(user_id)
+
+
+def set_session(session_id: Optional[str]):
+    """
+    Set the current session ID globally.
+    
+    Example:
+        agentbasis.set_session("sess-456")
+    """
+    _session_id.set(session_id)
+
+
+def set_conversation(conversation_id: Optional[str]):
+    """
+    Set the current conversation ID globally.
+    
+    Example:
+        agentbasis.set_conversation("conv-789")
+    """
+    _conversation_id.set(conversation_id)
+
+
+def set_metadata(metadata: Optional[Dict[str, Any]]):
+    """
+    Set custom metadata globally.
+    
+    Example:
+        agentbasis.set_metadata({"plan": "pro", "feature_flag": "new_ui"})
+    """
+    _metadata.set(metadata)
+
+
+# Getters for internal use by instrumentations
+def get_user() -> Optional[str]:
+    """Get the current user ID from context."""
+    return _user_id.get()
+
+
+def get_session() -> Optional[str]:
+    """Get the current session ID from context."""
+    return _session_id.get()
+
+
+def get_conversation() -> Optional[str]:
+    """Get the current conversation ID from context."""
+    return _conversation_id.get()
+
+
+def get_metadata() -> Optional[Dict[str, Any]]:
+    """Get the current metadata from context."""
+    return _metadata.get()
+
+
+def get_context_attributes() -> Dict[str, Any]:
+    """
+    Get all context attributes as a dictionary.
+    
+    This is used internally by instrumentations to inject context into spans.
+    Only returns attributes that are set (not None).
+    """
+    attributes = {}
+    
+    user_id = get_user()
+    if user_id:
+        attributes["agentbasis.user.id"] = user_id
+    
+    session_id = get_session()
+    if session_id:
+        attributes["agentbasis.session.id"] = session_id
+    
+    conversation_id = get_conversation()
+    if conversation_id:
+        attributes["agentbasis.conversation.id"] = conversation_id
+    
+    metadata = get_metadata()
+    if metadata:
+        attributes["agentbasis.metadata"] = json.dumps(metadata)
+    
+    return attributes
+
+
+# Decorator for function-level context
+def with_context(
+    user_id: Optional[str] = None,
+    session_id: Optional[str] = None,
+    conversation_id: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None
+):
+    """
+    Decorator to set context for a function.
+    
+    Example:
+        @agentbasis.with_context(user_id="user-123")
+        def my_agent_function():
+            response = client.chat.completions.create(...)
+            return response
+    """
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            with AgentBasisContext(
+                user_id=user_id,
+                session_id=session_id,
+                conversation_id=conversation_id,
+                metadata=metadata
+            ):
+                return func(*args, **kwargs)
+        
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            with AgentBasisContext(
+                user_id=user_id,
+                session_id=session_id,
+                conversation_id=conversation_id,
+                metadata=metadata
+            ):
+                return await func(*args, **kwargs)
+        
+        # Return appropriate wrapper based on function type
+        if _is_async_function(func):
+            return async_wrapper
+        return wrapper
+    
+    return decorator
+
+
+def _is_async_function(func) -> bool:
+    """Check if a function is async."""
+    import asyncio
+    return asyncio.iscoroutinefunction(func)
+
+
+def inject_context_to_span(span) -> None:
+    """
+    Inject current context attributes into a span.
+    
+    This is called by LLM instrumentations to automatically add
+    user/session context to every span.
+    
+    Args:
+        span: An OpenTelemetry Span object
+    """
+    attributes = get_context_attributes()
+    for key, value in attributes.items():
+        span.set_attribute(key, value)

agentbasis/decorators.py

@@ -0,0 +1,80 @@
+import functools
+import asyncio
+from typing import Callable
+from opentelemetry import trace as otel_trace
+from opentelemetry.trace import Status, StatusCode
+
+from agentbasis.context import inject_context_to_span
+
+
+def trace(func: Callable) -> Callable:
+    """
+    Decorator to track the execution of a function as an OTel Span.
+    
+    Works with both sync and async functions. Automatically injects
+    user/session context from agentbasis.context.
+    
+    Example:
+        @agentbasis.trace
+        def my_function():
+            ...
+        
+        @agentbasis.trace
+        async def my_async_function():
+            ...
+    """
+    
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        # Get the tracer at runtime, so it uses the configured provider
+        tracer = otel_trace.get_tracer("agentbasis")
+        
+        with tracer.start_as_current_span(func.__name__) as span:
+            # Inject user/session context
+            inject_context_to_span(span)
+            
+            # Record Inputs
+            span.set_attribute("code.function", func.__name__)
+            span.set_attribute("input.args", str(args))
+            span.set_attribute("input.kwargs", str(kwargs))
+
+            try:
+                result = func(*args, **kwargs)
+                span.set_attribute("output", str(result))
+                span.set_status(Status(StatusCode.OK))
+                return result
+
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    @functools.wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        # Get the tracer at runtime
+        tracer = otel_trace.get_tracer("agentbasis")
+        
+        with tracer.start_as_current_span(func.__name__) as span:
+            # Inject user/session context
+            inject_context_to_span(span)
+            
+            # Record Inputs
+            span.set_attribute("code.function", func.__name__)
+            span.set_attribute("input.args", str(args))
+            span.set_attribute("input.kwargs", str(kwargs))
+
+            try:
+                result = await func(*args, **kwargs)
+                span.set_attribute("output", str(result))
+                span.set_status(Status(StatusCode.OK))
+                return result
+
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    # Return appropriate wrapper based on function type
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    return sync_wrapper

agentbasis/frameworks/langchain/__init__.py

@@ -0,0 +1,109 @@
+from .callback import AgentBasisCallbackHandler
+
+# Module-level handler instance for convenience
+_handler_instance = None
+
+
+def get_callback_handler() -> AgentBasisCallbackHandler:
+    """
+    Get a new AgentBasis callback handler for LangChain.
+    
+    Use this to get a handler instance that you can pass to your
+    LangChain chains, agents, or LLMs.
+    
+    Returns:
+        A new AgentBasisCallbackHandler instance
+        
+    Example:
+        from agentbasis.frameworks.langchain import get_callback_handler
+        
+        handler = get_callback_handler()
+        
+        # Pass to chain
+        chain.invoke({"query": "..."}, config={"callbacks": [handler]})
+        
+        # Or pass to LLM
+        llm = ChatOpenAI(callbacks=[handler])
+        
+        # Or pass to agent
+        agent_executor = AgentExecutor(agent=agent, tools=tools, callbacks=[handler])
+    """
+    return AgentBasisCallbackHandler()
+
+
+def instrument() -> AgentBasisCallbackHandler:
+    """
+    Get the global AgentBasis callback handler for LangChain.
+    
+    This returns a singleton handler instance that can be reused across
+    your application. For most use cases, you should pass this handler
+    explicitly to your LangChain components.
+    
+    Returns:
+        The global AgentBasisCallbackHandler instance
+        
+    Example - Explicit callbacks (recommended):
+        from agentbasis.frameworks.langchain import instrument
+        
+        handler = instrument()
+        
+        # Option 1: Pass to invoke/run
+        chain.invoke({"query": "..."}, config={"callbacks": [handler]})
+        
+        # Option 2: Pass to constructor
+        llm = ChatOpenAI(model="gpt-4", callbacks=[handler])
+        
+        # Option 3: Pass to agent executor
+        agent = AgentExecutor(
+            agent=agent,
+            tools=tools,
+            callbacks=[handler]
+        )
+        
+    Example - With RunnableConfig:
+        from langchain_core.runnables import RunnableConfig
+        
+        handler = instrument()
+        config = RunnableConfig(callbacks=[handler])
+        
+        result = chain.invoke({"query": "..."}, config=config)
+    
+    Note:
+        Unlike OpenAI/Anthropic instrumentation which patches the SDK globally,
+        LangChain requires explicit callback passing. This is because LangChain
+        has its own callback system that doesn't support global monkey-patching.
+    """
+    global _handler_instance
+    
+    if _handler_instance is None:
+        _handler_instance = AgentBasisCallbackHandler()
+    
+    return _handler_instance
+
+
+def get_callback_config():
+    """
+    Get a LangChain RunnableConfig with AgentBasis callbacks pre-configured.
+    
+    This is a convenience function for users who prefer working with
+    RunnableConfig objects.
+    
+    Returns:
+        A dict that can be passed as the `config` parameter to invoke/batch/stream
+        
+    Example:
+        from agentbasis.frameworks.langchain import get_callback_config
+        
+        config = get_callback_config()
+        result = chain.invoke({"query": "..."}, config=config)
+    """
+    handler = instrument()
+    return {"callbacks": [handler]}
+
+
+__all__ = [
+    "AgentBasisCallbackHandler",
+    "instrument",
+    "get_callback_handler",
+    "get_callback_config",
+]