[DO NOT MERGE] Refactoring core session management

agentops/client.py

@@ -36,7 +36,7 @@ def __init__(self):
         self._pre_init_messages: List[str] = []
         self._initialized: bool = False
         self._llm_tracker: Optional[LlmTracker] = None
-        self._sessions: List[Session] = active_sessions
+        self._sessions: List[Session] = list(active_sessions)
         self._config = Configuration()
         self._pre_init_queue = {"agents": []}
 

agentops/config.py

@@ -1,74 +1,53 @@
-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 
 
+@dataclass
 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
+    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
+    graceful_shutdown_wait_time: int = 2000
+    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,
-        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,
+        **kwargs
     ):
-        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
+        # 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)