fix hooks

src/novu_py/_hooks/novuhook.py

@@ -1,60 +1,65 @@
-import requests
-import time
-import random
 import json
-import io
-from typing import Optional, Tuple, Union
+import random
+import time
+from typing import Union
+
+import httpx
 
 from .types import (
-    BeforeRequestContext,
     BeforeRequestHook,
-    AfterSuccessContext,
     AfterSuccessHook,
-    AfterErrorContext,
-    AfterErrorHook,
-    SDKInitHook
+    BeforeRequestContext,
+    AfterSuccessContext
 )
 
-class NovuHooks(SDKInitHook, BeforeRequestHook, AfterSuccessHook, AfterErrorHook):
-    def sdk_init(self, base_url: str, client: requests.Session) -> Tuple[str, requests.Session]:
-        """
-        Modify the base_url or wrap the client used by the SDK here and return the updated values.
 
-        :param base_url: The original base URL
-        :param client: The requests Session object
-        :return: Tuple of potentially modified base_url and client
-        """
-        return base_url, client
+def generate_idempotency_key():
+    """
+    Generate a unique idempotency key.
+
+    :return: A unique idempotency key as a string
+    """
+    timestamp = int(time.time() * 1000)  # Current time in milliseconds
+    random_string = ''.join(random.sample('abcdefghijklmnopqrstuvwxyz0123456789', 9))  # Unique alphanumeric string
+    return f"{timestamp}{random_string}"
 
-    def before_request(self, hook_ctx: BeforeRequestContext, request: requests.PreparedRequest) -> Union[requests.PreparedRequest, Exception]:
+
+class NovuHooks(BeforeRequestHook, AfterSuccessHook):
+    def before_request(self, hook_ctx: BeforeRequestContext, request: httpx.Request) -> Union[httpx.Request, Exception]:
         """
-        Modify the request headers before sending it.
+        Modify the request before sending.
 
         :param hook_ctx: Context for the before request hook
-        :param request: The prepared request to be modified
-        :return: Modified request or an exception
+        :param request: The request to be modified
+        :return: Modified request
         """
         auth_key = 'Authorization'
         idempotency_key = 'idempotency-key'
         api_key_prefix = 'ApiKey'
 
-        # Ensure headers exist and are a dictionary-like object
-        if not hasattr(request, 'headers') or request.headers is None:
-            request.headers = {}
+        # Create a copy of headers
+        headers = dict(request.headers or {})
 
         # Check and modify authorization header
-        if auth_key in request.headers:
-            key = request.headers[auth_key]
+        if auth_key in headers:
+            key = headers[auth_key]
             if key and not key.startswith(api_key_prefix):
-                request.headers[auth_key] = f"{api_key_prefix} {key}"
+                headers[auth_key] = f"{api_key_prefix} {key}"
 
         # Add idempotency key if not present
-        if idempotency_key not in request.headers or not request.headers[idempotency_key]:
-            request.headers[idempotency_key] = self.generate_idempotency_key()
-
-        return request
-
-    def after_success(self, hook_ctx: AfterSuccessContext, response: requests.Response) -> Union[requests.Response, Exception]:
+        if idempotency_key not in headers or not headers[idempotency_key]:
+            headers[idempotency_key] = generate_idempotency_key()
+
+        # Recreate the request with modified headers
+        return httpx.Request(
+            method=request.method,
+            url=request.url,
+            headers=headers,
+            content=request.content,
+            extensions=request.extensions
+        )
+
+    def after_success(self, hook_ctx: AfterSuccessContext, response: httpx.Response) -> Union[httpx.Response, Exception]:
         """
         Modify the response after a successful request.
 
@@ -76,44 +81,13 @@ def after_success(self, hook_ctx: AfterSuccessContext, response: requests.Respon
 
         # Check if the response contains a single 'data' key
         if isinstance(json_response, dict) and len(json_response) == 1 and 'data' in json_response:
-            # Create a new response manually
-            new_response = requests.Response()
-            new_response.status_code = response.status_code
-            new_response.headers = response.headers.copy()
-            new_response.reason = response.reason
-            new_response.url = response.url
-
-            # Serialize the 'data' content
-            data_content = json.dumps(json_response['data']).encode('utf-8')
-
-            # Use io.BytesIO to set the raw content
-            new_response.raw = io.BytesIO(data_content)
-
+            # Create a new HTTPX response
+            new_response = httpx.Response(
+                status_code=response.status_code,
+                headers=response.headers,
+                content=json.dumps(json_response['data']).encode('utf-8'),
+                request=response.request
+            )
             return new_response
 
         return response
-
-    def after_error(self, hook_ctx: AfterErrorContext, response: Optional[requests.Response], error: Optional[Exception]) -> Union[Tuple[Optional[requests.Response], Optional[Exception]], Exception]:
-        """
-        Modify the response or error after a failed request.
-
-        :param hook_ctx: Context for the after error hook
-        :param response: The response object (if any)
-        :param error: The exception that occurred
-        :return: Tuple of potentially modified response and error
-        """
-        # Validate input
-        if response is None and error is None:
-            return ValueError("Both response and error cannot be None")
-
-        return (response, error)
-
-    def generate_idempotency_key(self) -> str:
-        """
-        Generate a unique idempotency key using a timestamp and a random string.
-
-        :return: A unique idempotency key
-        """
-        timestamp = int(time.time() * 1000)  # Current time in milliseconds
-        random_string = ''.join(random.sample('abcdefghijklmnopqrstuvwxyz0123456789', 9))  # Unique alphanumeric string
-        return f"{timestamp}{random_string}"

src/novu_py/_hooks/registration.py

@@ -12,5 +12,3 @@ def init_hooks(hooks: Hooks):
     my_hook = NovuHooks()
     hooks.register_before_request_hook(my_hook)
     hooks.register_after_success_hook(my_hook)
-    hooks.register_after_error_hook(my_hook)
-    hooks.register_sdk_init_hook(my_hook)