From ffa5a15e0c2ef1c3422052d94c3aec7be40ee6ee Mon Sep 17 00:00:00 2001
From: Heeth Jain <143497789+heethjain21@users.noreply.github.com>
Date: Wed, 18 Sep 2024 07:39:38 +0530
Subject: [PATCH] feat: session end analytics: tools, llms, duration (#392)

* feat: session replay additional info: tools, llms, duration

* Updated formatting

---------

Co-authored-by: Howard Gil <howardbgil@gmail.com>
---
 agentops/session.py | 65 ++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 55 insertions(+), 10 deletions(-)

diff --git a/agentops/session.py b/agentops/session.py
index cec18c3f..a82eeb35 100644
--- a/agentops/session.py
+++ b/agentops/session.py
@@ -7,6 +7,7 @@
 from termcolor import colored
 from typing import Optional, List, Union
 from uuid import UUID, uuid4
+from dateutil import parser
 
 from .exceptions import ApiServerException
 from .enums import EndState
@@ -52,6 +53,13 @@ def __init__(
         self.jwt = None
         self.lock = threading.Lock()
         self.queue = []
+        self.event_counts = {
+            "llms": 0,
+            "tools": 0,
+            "actions": 0,
+            "errors": 0,
+            "apis": 0,
+        }
 
         self.stop_flag = threading.Event()
         self.thread = threading.Thread(target=self._run)
@@ -97,6 +105,20 @@ def end_session(
         self.thread.join(timeout=1)
         self._flush_queue()
 
+        def format_duration(start_time, end_time):
+            duration = parser.parse(end_time) - parser.parse(start_time)
+            hours, remainder = divmod(duration.total_seconds(), 3600)
+            minutes, seconds = divmod(remainder, 60)
+
+            parts = []
+            if hours > 0:
+                parts.append(f"{int(hours)}h")
+            if minutes > 0:
+                parts.append(f"{int(minutes)}m")
+            parts.append(f"{seconds:.1f}s")
+
+            return " ".join(parts)
+
         with self.lock:
             payload = {"session": self.__dict__}
             try:
@@ -111,23 +133,31 @@ def end_session(
         logger.debug(res.body)
         token_cost = res.body.get("token_cost", "unknown")
 
+        formatted_duration = format_duration(self.init_timestamp, self.end_timestamp)
+
         if token_cost == "unknown" or token_cost is None:
-            logger.info("Could not determine cost of run.")
             token_cost_d = Decimal(0)
         else:
             token_cost_d = Decimal(token_cost)
-            logger.info(
-                "This run's cost ${}".format(
-                    "{:.2f}".format(token_cost_d)
-                    if token_cost_d == 0
-                    else "{:.6f}".format(
-                        token_cost_d.quantize(
-                            Decimal("0.000001"), rounding=ROUND_HALF_UP
-                        )
-                    )
+            formatted_cost = (
+                "{:.2f}".format(token_cost_d)
+                if token_cost_d == 0
+                else "{:.6f}".format(
+                    token_cost_d.quantize(Decimal("0.000001"), rounding=ROUND_HALF_UP)
                 )
             )
 
+        analytics = (
+            f"Analytics for this run - "
+            f"LLM calls: {self.event_counts['llms']} | "
+            f"Tool calls: {self.event_counts['tools']} | "
+            f"Actions: {self.event_counts['actions']} | "
+            f"Errors: {self.event_counts['errors']} | "
+            f"Duration: {formatted_duration} | "
+            f"Cost: ${formatted_cost}"
+        )
+        logger.info(analytics)
+
         logger.info(
             colored(
                 f"\x1b[34mSession Replay: https://app.agentops.ai/drilldown?session_id={self.session_id}\x1b[0m",
@@ -293,6 +323,21 @@ def _flush_queue(self) -> None:
                 logger.debug(serialized_payload)
                 logger.debug("</AGENTOPS_DEBUG_OUTPUT>\n")
 
+                # Count total events created based on type
+                events = payload["events"]
+                for event in events:
+                    event_type = event["event_type"]
+                    if event_type == "llms":
+                        self.event_counts["llms"] += 1
+                    elif event_type == "tools":
+                        self.event_counts["tools"] += 1
+                    elif event_type == "actions":
+                        self.event_counts["actions"] += 1
+                    elif event_type == "errors":
+                        self.event_counts["errors"] += 1
+                    elif event_type == "apis":
+                        self.event_counts["apis"] += 1
+
     def _run(self) -> None:
         while not self.stop_flag.is_set():
             time.sleep(self.config.max_wait_time / 1000)