cleanup

Author: Dwij1704

agentops/__init__.py

@@ -5,7 +5,7 @@
     track_agent,
     track_tool,
     end_all_sessions,
-    Session,
+    Session as LegacySession,
     ToolEvent,
     ErrorEvent,
     ActionEvent,
@@ -104,7 +104,7 @@ def init(
     )
 
 
-def configure(**kwargs):
+def configure(**kwargs: Any) -> None:
     """Update client configuration
 
     Args:
@@ -219,7 +219,7 @@ def end_trace(trace_context: TraceContext, end_state: str = "Success") -> None:
     "track_agent",
     "track_tool",
     "end_all_sessions",
-    "Session",
+    "LegacySession",
     "ToolEvent",
     "ErrorEvent",
     "ActionEvent",

agentops/client/client.py

@@ -1,5 +1,5 @@
 import atexit
-from typing import Optional
+from typing import Optional, Any
 
 from agentops.client.api import ApiClient
 from agentops.config import Config
@@ -49,7 +49,7 @@ class Client:
 
     api: ApiClient
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: Any, **kwargs: Any) -> "Client":
         if cls.__instance is None:
             cls.__instance = super(Client, cls).__new__(cls)
             # Initialize instance variables that should only be set once per instance
@@ -69,7 +69,7 @@ def __init__(self):
             # self._init_trace_context = None # Already done in __new__
             # self._legacy_session_for_init_trace = None # Already done in __new__
 
-    def init(self, **kwargs) -> None:  # Return type updated to None
+    def init(self, **kwargs: Any) -> None:  # Return type updated to None
         # Recreate the Config object to parse environment variables at the time of initialization
         # This allows re-init with new env vars if needed, though true singletons usually init once.
         self.config = Config()
@@ -168,7 +168,7 @@ def init(self, **kwargs) -> None:  # Return type updated to None
             self._initialized = True  # Successfully initialized, just no auto-trace
             return None  # No auto-session, so return None
 
-    def configure(self, **kwargs):
+    def configure(self, **kwargs: Any) -> None:
         """Update client configuration"""
         self.config.configure(**kwargs)
 
@@ -177,7 +177,7 @@ def initialized(self) -> bool:
         return self._initialized
 
     @initialized.setter
-    def initialized(self, value: bool):
+    def initialized(self, value: bool) -> None:
         if self._initialized and self._initialized != value:
             # Allow re-setting to False if we are intentionally re-initializing
             # This logic is now partly in init() to handle re-init cases