Refactoring core session management

agentops/__init__.py

@@ -1,22 +1,21 @@
 # agentops/__init__.py
+import json
 import sys
-from typing import Optional, List, Union
+import threading
+from importlib.metadata import version as get_version
+from typing import List, Optional, Union
+
+from packaging import version
 
 from .client import Client
-from .event import Event, ActionEvent, LLMEvent, ToolEvent, ErrorEvent
-from .decorators import record_action, track_agent, record_tool, record_function
+from .decorators import record_action, record_function, record_tool, track_agent
+from .event import ActionEvent, ErrorEvent, Event, LLMEvent, ToolEvent
 from .helpers import check_agentops_update
 from .log_config import logger
-from .session import Session
-import threading
-from importlib.metadata import version as get_version
-from packaging import version
+from .session import Session, active_sessions
 
 try:
-    from .partners.langchain_callback_handler import (
-        LangchainCallbackHandler,
-        AsyncLangchainCallbackHandler,
-    )
+    from .partners.langchain_callback_handler import AsyncLangchainCallbackHandler, LangchainCallbackHandler
 except ModuleNotFoundError:
     pass
 
@@ -150,7 +149,7 @@ def configure(
 def start_session(
     tags: Optional[List[str]] = None,
     inherited_session_id: Optional[str] = None,
-) -> Union[Session, None]:
+) -> Session:
     """
     Start a new session for recording events.
 
@@ -321,3 +320,24 @@ def get_session(session_id: str):
 # prevents unexpected sessions on new tests
 def end_all_sessions() -> None:
     return Client().end_all_sessions()
+
+
+def flush():
+    """
+    Publish all pending events to API and return the flushed events
+
+    Returns:
+        bool: Whether the flush was successful
+    """
+    # from agentops.session.exporter import get_tracer_provider
+    from opentelemetry import trace
+
+    from agentops.client import Client
+
+    # assert Client()
+    # Get the provider and force flush
+    provider = trace.get_tracer_provider()
+
+    print("@@@PROVIDER:" + str(id(provider)))
+    breakpoint()
+    return provider.force_flush()  # !!

agentops/client.py

@@ -5,29 +5,36 @@
     Client: Provides methods to interact with the AgentOps service.
 """
 
-import inspect
 import atexit
+import inspect
 import logging
 import os
 import signal
 import sys
 import threading
 import traceback
 from decimal import Decimal
+from functools import cached_property
+from typing import List, Optional, Tuple, Union
 from uuid import UUID, uuid4
-from typing import Optional, List, Union, Tuple
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
 from termcolor import colored
 
-from .event import Event, ErrorEvent
-from .singleton import (
-    conditional_singleton,
-)
-from .session import Session, active_sessions
+from .config import Configuration
+from .event import ErrorEvent, Event
 from .host_env import get_host_env
+from .llms import LlmTracker
 from .log_config import logger
 from .meta_client import MetaClient
-from .config import Configuration
-from .llms import LlmTracker
+from .session import Session, active_sessions
+from .singleton import conditional_singleton
+
+# This can only be done once, a warning will be logged if any further attempt is made.
+_provider = TracerProvider()
+print("@@@PROVIDER:" + str(id(_provider)))
+trace.set_tracer_provider(_provider)
 
 
 @conditional_singleton
@@ -39,6 +46,7 @@ def __init__(self):
         self._sessions: List[Session] = active_sessions
         self._config = Configuration()
         self._pre_init_queue = {"agents": []}
+        self._host_env = None  # Cache host env data
 
         self.configure(
             api_key=os.environ.get("AGENTOPS_API_KEY"),
@@ -111,6 +119,7 @@ def initialize(self) -> Union[Session, None]:
     def _initialize_autogen_logger(self) -> None:
         try:
             import autogen
+
             from .partners.autogen_logger import AutogenLogger
 
             autogen.runtime_logging.start(logger=AutogenLogger())
@@ -196,7 +205,7 @@ def start_session(
         self,
         tags: Optional[List[str]] = None,
         inherited_session_id: Optional[str] = None,
-    ) -> Union[Session, None]:
+    ) -> Session:
         """
         Start a new session for recording events.
 
@@ -224,7 +233,7 @@ def start_session(
         session = Session(
             session_id=session_id,
             tags=list(session_tags),
-            host_env=get_host_env(self._config.env_data_opt_out),
+            host_env=self.host_env,
             config=self._config,
         )
 
@@ -430,3 +439,8 @@ def api_key(self):
     @property
     def parent_key(self):
         return self._config.parent_key
+
+    @cached_property
+    def host_env(self):
+        """Cache and reuse host environment data"""
+        return get_host_env(self._config.env_data_opt_out)

agentops/config.py

@@ -1,74 +1,49 @@
-from typing import List, Optional
+from dataclasses import dataclass, field
+from typing import Optional, Set
 from uuid import UUID
 
 from .log_config import logger
 
+# TODO: Use annotations to clarify the purpose of each attribute.
+# Details are defined in a docstrings found in __init__.py, but
+# it's good to have those right on the fields at class definition
 
-class Configuration:
-    def __init__(self):
-        self.api_key: Optional[str] = None
-        self.parent_key: Optional[str] = None
-        self.endpoint: str = "https://api.agentops.ai"
-        self.max_wait_time: int = 5000
-        self.max_queue_size: int = 512
-        self.default_tags: set[str] = set()
-        self.instrument_llm_calls: bool = True
-        self.auto_start_session: bool = True
-        self.skip_auto_end_session: bool = False
-        self.env_data_opt_out: bool = False
-
-    def configure(
-        self,
-        client,
-        api_key: Optional[str] = None,
-        parent_key: Optional[str] = None,
-        endpoint: Optional[str] = None,
-        max_wait_time: Optional[int] = None,
-        max_queue_size: Optional[int] = None,
-        default_tags: Optional[List[str]] = None,
-        instrument_llm_calls: Optional[bool] = None,
-        auto_start_session: Optional[bool] = None,
-        skip_auto_end_session: Optional[bool] = None,
-        env_data_opt_out: Optional[bool] = None,
-    ):
-        if api_key is not None:
-            try:
-                UUID(api_key)
-                self.api_key = api_key
-            except ValueError:
-                message = f"API Key is invalid: {{{api_key}}}.\n\t    Find your API key at https://app.agentops.ai/settings/projects"
-                client.add_pre_init_warning(message)
-                logger.error(message)
-
-        if parent_key is not None:
-            try:
-                UUID(parent_key)
-                self.parent_key = parent_key
-            except ValueError:
-                message = f"Parent Key is invalid: {parent_key}"
-                client.add_pre_init_warning(message)
-                logger.warning(message)
-
-        if endpoint is not None:
-            self.endpoint = endpoint
-
-        if max_wait_time is not None:
-            self.max_wait_time = max_wait_time
-
-        if max_queue_size is not None:
-            self.max_queue_size = max_queue_size
 
-        if default_tags is not None:
-            self.default_tags.update(default_tags)
-
-        if instrument_llm_calls is not None:
-            self.instrument_llm_calls = instrument_llm_calls
-
-        if auto_start_session is not None:
-            self.auto_start_session = auto_start_session
-
-        if skip_auto_end_session is not None:
-            self.skip_auto_end_session = skip_auto_end_session
-
-        if env_data_opt_out is not None:
-            self.env_data_opt_out = env_data_opt_out
+@dataclass
+class Configuration:
+    api_key: Optional[str] = None
+    parent_key: Optional[str] = None
+    endpoint: str = "https://api.agentops.ai"
+    max_wait_time: int = 5000
+    max_queue_size: int = 512
+    default_tags: Set[str] = field(default_factory=set)
+    instrument_llm_calls: bool = True
+    auto_start_session: bool = True
+    skip_auto_end_session: bool = False
+    env_data_opt_out: bool = False
+
+    def configure(self, client, **kwargs):
+        # Special handling for keys that need UUID validation
+        for key_name in ["api_key", "parent_key"]:
+            if key_name in kwargs and kwargs[key_name] is not None:
+                try:
+                    UUID(kwargs[key_name])
+                    setattr(self, key_name, kwargs[key_name])
+                except ValueError:
+                    message = (
+                        f"API Key is invalid: {{{kwargs[key_name]}}}.\n\t    Find your API key at https://app.agentops.ai/settings/projects"
+                        if key_name == "api_key"
+                        else f"Parent Key is invalid: {kwargs[key_name]}"
+                    )
+                    client.add_pre_init_warning(message)
+                    logger.error(message) if key_name == "api_key" else logger.warning(message)
+                kwargs.pop(key_name)
+
+        # Special handling for default_tags which needs update() instead of assignment
+        if "default_tags" in kwargs and kwargs["default_tags"] is not None:
+            self.default_tags.update(kwargs.pop("default_tags"))
+
+        # Handle all other attributes
+        for key, value in kwargs.items():
+            if value is not None and hasattr(self, key):
+                setattr(self, key, value)

agentops/decorators.py

@@ -1,5 +1,6 @@
 import functools
 import inspect
+from types import MethodType
 from typing import Optional, Union
 from uuid import uuid4
 
@@ -361,3 +362,99 @@ def new_init(self, *args, **kwargs):
         return obj
 
     return decorator
+
+
+# >>> multi-purpose atomic-ops decorator
+import threading
+import weakref
+from functools import wraps
+from typing import Dict, Union
+
+
+class atomic:
+    """A decorator class that provides thread-safe atomic operations.
+
+    This decorator ensures that decorated functions/methods execute atomically by using locks.
+    It supports both function and method decoration, with flexible locking strategies.
+
+    Key Features:
+    - Thread-safe execution of decorated functions/methods
+    - Supports both @atomic and @atomic() syntax
+    - Configurable locking scope via custom keys
+    - Optional reentrant locking
+    - Instance-specific locking for methods
+
+    Usage:
+        @atomic  # Uses function name as lock key
+        def my_func():
+            pass
+
+        @atomic(key="custom_key")  # Uses custom key for lock
+        def another_func():
+            pass
+
+        @atomic(key=lambda self: f"instance_{id(self)}")  # Instance-specific locks
+        def method(self):
+            pass
+
+    Args:
+        key: Optional[Union[str, Callable]]. Lock key or key generator function.
+             If None, uses function's qualified name.
+        reentrant: bool. If True, allows reentrant locking. Default False.
+    """
+
+    _locks = {}  # Class-level registry of all locks
+    _registry_lock = threading.Lock()  # Lock for thread-safe registry access
+
+    def __init__(self, func=None, *, key=None, reentrant=False):
+        # Support both @atomic and @atomic() syntax
+        self.func = func
+        self.key = key
+        self.reentrant = reentrant
+
+        if func is not None:  # @atomic case
+            self._wrapped = self._wrap(func)
+
+    def __call__(self, *args, **kwargs):
+        if self.func is None:  # @atomic() case
+            self.func = args[0]
+            self._wrapped = self._wrap(self.func)
+            return self._wrapped
+        return self._wrapped(*args, **kwargs)
+
+    def _wrap(self, func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Get instance for instance methods
+            instance = args[0] if args and not isinstance(args[0], type) else None
+
+            # Generate lock key
+            if callable(self.key) and instance:
+                # For instance-specific locks, use the key generator
+                lock_key = self.key(instance)
+            else:
+                # For shared locks or function locks, use the provided key or function qualname
+                lock_key = self.key or f"{func.__module__}.{func.__qualname__}"
+
+            # Get or create lock
+            with self._registry_lock:
+                if lock_key not in self._locks:
+                    self._locks[lock_key] = threading.RLock()  # THIS HAS TO BE A RLOCK NO MATTER WHAT
+                lock = self._locks[lock_key]
+
+            if not self.reentrant and lock._is_owned():
+                raise RuntimeError("Lock is not reentrant")
+
+            with lock:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    def __get__(self, obj, objtype=None):
+        """Support descriptor protocol for instance methods"""
+        if obj is None:
+            return self
+        return MethodType(self._wrapped, obj)
+
+
+# >>> @atomic