Skip to content

Commit

Permalink
refactor: Update end_session to return detailed statistics
Browse files Browse the repository at this point in the history
Co-Authored-By: Alex Reibman <[email protected]>
  • Loading branch information
devin-ai-integration[bot] and areibman committed Dec 18, 2024
1 parent 8c8fa36 commit bb3b056
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 12 deletions.
14 changes: 6 additions & 8 deletions agentops/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand Down
10 changes: 6 additions & 4 deletions agentops/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 = (
Expand All @@ -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

Expand All @@ -350,7 +352,7 @@ def end_session(
"blue",
)
)
return self.token_cost
return analytics_stats

def add_tags(self, tags: List[str]) -> None:
"""
Expand Down

0 comments on commit bb3b056

Please sign in to comment.