From 3432d6e9e224be2e756df3188f8214e95c1225e2 Mon Sep 17 00:00:00 2001
From: Howard Gil <howardbgil@gmail.com>
Date: Fri, 15 Dec 2023 13:26:39 -0600
Subject: [PATCH] Adding support for screenshots and video (#55)

* Adding support for multion screenshots and video

* Added support for screenshot images and screen recording videos

* Changed screenshot type from PIL Image to string (base64)

* Updated arg definition to reflect that screenshot can be base64 or url

* typo
---
 agentops/client.py |  9 ++++++---
 agentops/event.py  | 10 ++++++----
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/agentops/client.py b/agentops/client.py
index 4acb2055..cf206785 100644
--- a/agentops/client.py
+++ b/agentops/client.py
@@ -245,19 +245,22 @@ def start_session(self, tags: Optional[List[str]] = None):
         self.worker.start_session(self.session)
 
     def end_session(self, end_state: str = Field("Indeterminate",
-                                                 description="End state of the session",
-                                                 pattern="^(Success|Fail|Indeterminate)$"),
-                    rating: Optional[str] = None):
+                                                description="End state of the session",
+                                                pattern="^(Success|Fail|Indeterminate)$"),
+                                                rating: Optional[str] = None,
+                                                video: Optional[str] = None):
         """
         End the current session with the AgentOps service.
 
         Args:
             end_state (str, optional): The final state of the session.
             rating (str, optional): The rating for the session.
+            video (str, optional): The video screen recording of the session
         """
         if not self.session.has_ended:
             self.session.end_session(end_state, rating)
             self.worker.end_session(self.session)
+            self.session.video = video
         else:
             logging.info("Warning: The session has already been ended.")
 
diff --git a/agentops/event.py b/agentops/event.py
index 1bc0b2ab..e33b76d6 100644
--- a/agentops/event.py
+++ b/agentops/event.py
@@ -8,7 +8,6 @@
 from typing import Optional, List
 from pydantic import Field
 
-
 class Event:
     """
     Represents a discrete event to be recorded.
@@ -18,7 +17,7 @@ class Event:
         params (str, optional): The parameters passed to the operation.
         returns (str, optional): The output of the operation.
         result (str, optional): Result of the operation, e.g., "Success", "Fail", "Indeterminate". Defaults to "Indeterminate".
-        action_type (str, optional): Type of action of the event e.g. 'action', 'llm', 'api'. Defaults to 'action'.
+        action_type (str, optional): Type of action of the event e.g. 'action', 'llm', 'api', 'screenshot'. Defaults to 'action'.
         model (Models, optional): The model used during the event if an LLM is used (i.e. GPT-4).
                 For models, see the types available in the Models enum. 
                 If a model is set but an action_type is not, the action_type will be coerced to 'llm'. 
@@ -27,6 +26,7 @@ class Event:
         tags (List[str], optional): Tags that can be used for grouping or sorting later. e.g. ["my_tag"]. Defaults to None.
         init_timestamp (float, optional): The timestamp for when the event was initiated, represented as seconds since the epoch.
                 Defaults to the end timestamp.
+        screenshot (str, optional): A screenshot of the webpage at the time of the event. Base64 string or URL. Defaults to None.
 
     Attributes:
         event_type (str): Type of the event.
@@ -49,11 +49,12 @@ def __init__(self, event_type: str,
                                      pattern="^(Success|Fail|Indeterminate)$"),
                  action_type: Optional[str] = Field("action",
                                                     description="Type of action that the user is recording",
-                                                    pattern="^(action|api|llm)$"),
+                                                    pattern="^(action|api|llm|screenshot)$"),
                  model: Optional[Models] = None,
                  prompt: Optional[str] = None,
                  tags: Optional[List[str]] = None,
-                 init_timestamp: Optional[float] = None
+                 init_timestamp: Optional[float] = None,
+                 screenshot: Optional[str] = None
                  ):
         self.event_type = event_type
         self.params = params
@@ -65,3 +66,4 @@ def __init__(self, event_type: str,
         self.prompt = prompt
         self.end_timestamp = get_ISO_time()
         self.init_timestamp = init_timestamp if init_timestamp else self.end_timestamp
+        self.screenshot = screenshot
\ No newline at end of file