Copy of PR #385 - Dead letter queue

.github/workflows/python-testing.yml

@@ -23,6 +23,9 @@ jobs:
       matrix:
         python-version: [3.7,3.8,3.9,3.10,3.11,3.12]
 
+    env:
+      ENVIRONMENT: test
+
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python

.gitignore

@@ -164,4 +164,5 @@ cython_debug/
 .DS_Store
 
 agentops_time_travel.json
-.agentops_time_travel.yaml
+.agentops_time_travel.yaml
+.agentops/

agentops/config.py

@@ -2,8 +2,10 @@
 from uuid import UUID
 
 from .log_config import logger
+from .singleton import singleton
 
 
+@singleton
 class Configuration:
     def __init__(self):
         self.api_key: Optional[str] = None

agentops/helpers.py

@@ -6,10 +6,10 @@
 import http.client
 import json
 from importlib.metadata import version, PackageNotFoundError
-
 from .log_config import logger
 from uuid import UUID
 from importlib.metadata import version
+import os
 
 
 def get_ISO_time():
@@ -179,3 +179,19 @@ def wrapper(self, *args, **kwargs):
         return func(self, *args, **kwargs)
 
     return wrapper
+
+
+def ensure_dead_letter_queue():
+    # Define file path
+    file_path = os.path.join(".agentops", "dead_letter_queue.json")
+
+    # Check if directory exists
+    if not os.path.exists(".agentops"):
+        os.makedirs(".agentops")
+
+    # Check if file exists
+    if not os.path.isfile(file_path):
+        with open(file_path, "w") as f:
+            json.dump({"messages": []}, f)
+
+    return file_path

agentops/http_client.py

@@ -1,9 +1,21 @@
+from datetime import datetime
 from enum import Enum
-from typing import Optional
+from typing import Optional, List, Union
+
+import jwt
 from requests.adapters import Retry, HTTPAdapter
 import requests
+from agentops.log_config import logger
+from .config import Configuration
 
 from .exceptions import ApiServerException
+from dotenv import load_dotenv
+import os
+
+from .helpers import ensure_dead_letter_queue, filter_unjsonable, safe_serialize
+import json
+
+load_dotenv()
 
 JSON_HEADER = {"Content-Type": "application/json; charset=UTF-8", "Accept": "*/*"}
 
@@ -21,8 +33,49 @@ class HttpStatus(Enum):
     UNKNOWN = -1
 
 
-class Response:
+class DeadLetterQueue:
+    def __init__(self):
+        self.queue: List[dict] = []
+        self.is_testing = os.environ.get("ENVIRONMENT") == "test"
+
+        # if not self.is_testing:
+        self.file_path = ensure_dead_letter_queue()
+
+    def read_queue(self):
+        if not self.is_testing:
+            with open(self.file_path, "r") as f:
+                return json.load(f)["messages"]
+        else:
+            return []
+
+    def write_queue(self):
+        if not self.is_testing:
+            with open(self.file_path, "w") as f:
+                json.dump({"messages": safe_serialize(self.queue)}, f)
+
+    def add(self, request_data: dict):
+        if not self.is_testing:
+            self.queue.append(request_data)
+            self.write_queue()
+
+    def get_all(self) -> List[dict]:
+        return self.queue
+
+    def remove(self, request_data: dict):
+        if not self.is_testing:
+            if request_data in self.queue:
+                self.queue.remove(request_data)
+                self.write_queue()
+
+    def clear(self):
+        self.queue.clear()
+        self.write_queue()
 
+
+dead_letter_queue = DeadLetterQueue()
+
+
+class Response:
     def __init__(
         self, status: HttpStatus = HttpStatus.UNKNOWN, body: Optional[dict] = None
     ):
@@ -57,15 +110,13 @@ def get_status(code: int) -> HttpStatus:
 
 
 class HttpClient:
-
     @staticmethod
     def post(
         url: str,
         payload: bytes,
         api_key: Optional[str] = None,
         parent_key: Optional[str] = None,
-        jwt: Optional[str] = None,
-        header=None,
+        token: Optional[str] = None,
     ) -> Response:
         result = Response()
         try:
@@ -79,34 +130,47 @@ def post(
             if parent_key is not None:
                 JSON_HEADER["X-Agentops-Parent-Key"] = parent_key
 
-            if jwt is not None:
-                JSON_HEADER["Authorization"] = f"Bearer {jwt}"
+            if token is not None:
+                decoded_jwt = jwt.decode(
+                    token,
+                    algorithms=["HS256"],
+                    options={"verify_signature": False},
+                )
+
+                # if token is expired, reauth
+                if datetime.fromtimestamp(decoded_jwt["exp"]) < datetime.now():
+                    new_jwt = reauthorize_jwt(
+                        token,
+                        api_key,
+                        decoded_jwt["session_id"],
+                    )
+                    token = new_jwt
+
+                JSON_HEADER["Authorization"] = f"Bearer {token}"
 
             res = request_session.post(
                 url, data=payload, headers=JSON_HEADER, timeout=20
             )
 
             result.parse(res)
-        except requests.exceptions.Timeout:
-            result.code = 408
-            result.status = HttpStatus.TIMEOUT
-            raise ApiServerException(
-                "Could not reach API server - connection timed out"
+
+            if result.code == 200:
+                HttpClient._retry_dlq_requests()
+
+        except (requests.exceptions.Timeout, requests.exceptions.HTTPError) as e:
+            HttpClient._handle_failed_request(
+                url, payload, api_key, parent_key, token, type(e).__name__
             )
-        except requests.exceptions.HTTPError as e:
-            try:
-                result.parse(e.response)
-            except Exception:
-                result = Response()
-                result.code = e.response.status_code
-                result.status = Response.get_status(e.response.status_code)
-                result.body = {"error": str(e)}
-                raise ApiServerException(f"HTTPError: {e}")
+            raise ApiServerException(f"{type(e).__name__}: {e}")
         except requests.exceptions.RequestException as e:
-            result.body = {"error": str(e)}
+            HttpClient._handle_failed_request(
+                url, payload, api_key, parent_key, token, "RequestException"
+            )
             raise ApiServerException(f"RequestException: {e}")
 
         if result.code == 401:
+            if result.body.get("message") == "Expired Token":
+                raise ApiServerException(f"API server: jwt token expired.")
             raise ApiServerException(
                 f"API server: invalid API key: {api_key}. Find your API key at https://app.agentops.ai/settings/projects"
             )
@@ -116,6 +180,9 @@ def post(
             else:
                 raise ApiServerException(f"API server: {result.body}")
         if result.code == 500:
+            HttpClient._handle_failed_request(
+                url, payload, api_key, parent_key, token, "ServerError"
+            )
             raise ApiServerException("API server: - internal server error")
 
         return result
@@ -124,8 +191,7 @@ def post(
     def get(
         url: str,
         api_key: Optional[str] = None,
-        jwt: Optional[str] = None,
-        header=None,
+        token: Optional[str] = None,
     ) -> Response:
         result = Response()
         try:
@@ -136,29 +202,25 @@ def get(
             if api_key is not None:
                 JSON_HEADER["X-Agentops-Api-Key"] = api_key
 
-            if jwt is not None:
-                JSON_HEADER["Authorization"] = f"Bearer {jwt}"
+            if token is not None:
+                JSON_HEADER["Authorization"] = f"Bearer {token}"
 
             res = request_session.get(url, headers=JSON_HEADER, timeout=20)
 
             result.parse(res)
-        except requests.exceptions.Timeout:
-            result.code = 408
-            result.status = HttpStatus.TIMEOUT
-            raise ApiServerException(
-                "Could not reach API server - connection timed out"
+
+            if result.code == 200:
+                HttpClient._retry_dlq_requests()
+
+        except (requests.exceptions.Timeout, requests.exceptions.HTTPError) as e:
+            HttpClient._handle_failed_request(
+                url, None, api_key, None, token, type(e).__name__
             )
-        except requests.exceptions.HTTPError as e:
-            try:
-                result.parse(e.response)
-            except Exception:
-                result = Response()
-                result.code = e.response.status_code
-                result.status = Response.get_status(e.response.status_code)
-                result.body = {"error": str(e)}
-                raise ApiServerException(f"HTTPError: {e}")
+            raise ApiServerException(f"{type(e).__name__}: {e}")
         except requests.exceptions.RequestException as e:
-            result.body = {"error": str(e)}
+            HttpClient._handle_failed_request(
+                url, None, api_key, None, token, "RequestException"
+            )
             raise ApiServerException(f"RequestException: {e}")
 
         if result.code == 401:
@@ -171,6 +233,79 @@ def get(
             else:
                 raise ApiServerException(f"API server: {result.body}")
         if result.code == 500:
+            HttpClient._handle_failed_request(
+                url, None, api_key, None, jwt, "ServerError"
+            )
             raise ApiServerException("API server: - internal server error")
 
         return result
+
+    @staticmethod
+    def _retry_dlq_requests():
+        """Retry requests in the DLQ"""
+        for failed_request in dead_letter_queue.get_all():
+            dead_letter_queue.clear()
+            try:
+                if "payload" in failed_request:
+                    # Retry POST request from DLQ
+                    HttpClient.post(
+                        failed_request["url"],
+                        json.dumps(filter_unjsonable(failed_request["payload"])).encode(
+                            "utf-8"
+                        ),
+                        failed_request["api_key"],
+                        failed_request["parent_key"],
+                        failed_request["jwt"],
+                    )
+                else:
+                    # Retry GET request from DLQ
+                    HttpClient.get(
+                        failed_request["url"],
+                        failed_request["api_key"],
+                        failed_request["jwt"],
+                    )
+            except ApiServerException as e:
+                dead_letter_queue.add(failed_request)
+                # If it still fails, keep it in the DLQ
+            except Exception as e:
+                dead_letter_queue.add(failed_request)
+
+    @staticmethod
+    def _handle_failed_request(
+        url: str,
+        payload: Optional[bytes],
+        api_key: Optional[str],
+        parent_key: Optional[str],
+        jwt: Optional[str],
+        error_type: str,
+    ):
+        dead_letter_queue.add(
+            {
+                "url": url,
+                "payload": payload,
+                "api_key": api_key,
+                "parent_key": parent_key,
+                "jwt": jwt,
+                "error_type": error_type,
+            }
+        )
+
+        logger.warning(
+            f"An error occurred while communicating with the server: {error_type}"
+        )
+
+
+def reauthorize_jwt(old_jwt: str, api_key: str, session_id: str) -> Union[str, None]:
+    payload = {"jwt": old_jwt, "session_id": session_id}
+    serialized_payload = json.dumps(filter_unjsonable(payload)).encode("utf-8")
+    config = Configuration()
+    res = HttpClient.post(
+        f"{config.endpoint}/v2/reauthorize_jwt",
+        serialized_payload,
+        api_key,
+    )
+
+    if res.code != 200:
+        return None
+
+    return res.body.get("jwt", None)

agentops/session.py

@@ -125,7 +125,7 @@ def format_duration(start_time, end_time):
                 res = HttpClient.post(
                     f"{self.config.endpoint}/v2/update_session",
                     json.dumps(filter_unjsonable(payload)).encode("utf-8"),
-                    jwt=self.jwt,
+                    token=self.jwt,
                 )
             except ApiServerException as e:
                 return logger.error(f"Could not end session - {e}")
@@ -289,7 +289,7 @@ def _update_session(self) -> None:
                 res = HttpClient.post(
                     f"{self.config.endpoint}/v2/update_session",
                     json.dumps(filter_unjsonable(payload)).encode("utf-8"),
-                    jwt=self.jwt,
+                    token=self.jwt,
                 )
             except ApiServerException as e:
                 return logger.error(f"Could not update session - {e}")
@@ -311,7 +311,7 @@ def _flush_queue(self) -> None:
                     HttpClient.post(
                         f"{self.config.endpoint}/v2/create_events",
                         serialized_payload,
-                        jwt=self.jwt,
+                        token=self.jwt,
                     )
                 except ApiServerException as e:
                     return logger.error(f"Could not post events - {e}")
@@ -360,7 +360,7 @@ def create_agent(self, name, agent_id):
             HttpClient.post(
                 f"{self.config.endpoint}/v2/create_agent",
                 serialized_payload,
-                jwt=self.jwt,
+                token=self.jwt,
             )
         except ApiServerException as e:
             return logger.error(f"Could not create agent - {e}")