diff --git a/agentops/client.py b/agentops/client.py index fb3e1793..a4814fb6 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -15,7 +15,7 @@ import traceback from decimal import Decimal from functools import cached_property -from typing import List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Set, Tuple, Union from uuid import UUID, uuid4 from termcolor import colored @@ -247,7 +247,7 @@ def end_session( end_state_reason: Optional[str] = None, video: Optional[str] = None, is_auto_end: Optional[bool] = None, - ) -> Optional[Decimal]: + ) -> Optional[Dict[str, Any]]: """ End the current session with the AgentOps service. @@ -258,17 +258,15 @@ def end_session( is_auto_end (bool, optional): is this an automatic use of end_session and should be skipped with skip_auto_end_session Returns: - Decimal: The token cost of the session. Returns 0 if the cost is unknown. + Dict[str, Any]: Session statistics including duration, cost, and event counts. """ session = self._safe_get_session() if session is None: - return + return None if is_auto_end and self._config.skip_auto_end_session: - return - - token_cost = session.end_session(end_state=end_state, end_state_reason=end_state_reason, video=video) + return None - return token_cost + return session.end_session(end_state=end_state, end_state_reason=end_state_reason, video=video) def create_agent( self, diff --git a/agentops/session.py b/agentops/session.py index 535bc997..69c55ce8 100644 --- a/agentops/session.py +++ b/agentops/session.py @@ -284,7 +284,7 @@ def end_session( end_state: str = "Indeterminate", end_state_reason: Optional[str] = None, video: Optional[str] = None, - ) -> Union[Decimal, None]: + ) -> Union[Dict[str, Any], None]: with self._end_session_lock: if not self.is_running: return None @@ -324,8 +324,9 @@ def end_session( finally: del self._span_processor - # 5. Final session update - if not (analytics_stats := self.get_analytics()): + # 5. Final session update and get analytics + analytics_stats = self.get_analytics() + if not analytics_stats: return None analytics = ( @@ -341,6 +342,7 @@ def end_session( except Exception as e: logger.exception(f"Error during session end: {e}") + return None finally: active_sessions.remove(self) # First thing, get rid of the session @@ -350,7 +352,7 @@ def end_session( "blue", ) ) - return self.token_cost + return analytics_stats def add_tags(self, tags: List[str]) -> None: """