Merge pull request #42 from basedosdados/feat/refresh-tokens

feat: refresh tokens

Author: vrtornisiello

frontend/api/api_client.py

@@ -1,30 +1,101 @@
+from datetime import datetime, timedelta
 from typing import Iterator
 from uuid import UUID, uuid4
 
 import httpx
+import jwt
+import streamlit as st
 from loguru import logger
 
 from frontend.datatypes import (EventData, MessagePair, StreamEvent, Thread,
                                 UserMessage)
+from frontend.exceptions import SessionExpiredException
 
 
 class APIClient:
     def __init__(self, base_url: str):
         self.base_url = base_url
         self.logger = logger.bind(classname=self.__class__.__name__)
 
-    def authenticate(self, email: str, password: str) -> tuple[str|None, str]:
+    def _is_token_expired(self, token: str) -> bool:
+        """Check if a JWT token is expired or about to expire (within 1 minute).
+
+        Args:
+            token (str): The token.
+
+        Returns:
+            bool: Whether the token is expired or not.
+        """
+        if not token:
+            return True
+
+        try:
+            payload: dict = jwt.decode(token, options={"verify_signature": False})
+            exp = payload.get("exp")
+            if not exp:
+                return True
+            expiration = datetime.fromtimestamp(exp)
+            return datetime.now() >= expiration - timedelta(seconds=60)
+        except Exception:
+            return True
+
+    def _refresh_access_token(self, refresh_token: str) -> str:
+        """Refresh the access token using a refresh token.
+
+        Args:
+            refresh_token (str): The refresh token.
+
+        Returns:
+            str: A refreshed access token.
+        """
+        response = httpx.post(
+            f"{self.base_url}/chatbot/token/refresh/",
+            json={"refresh": refresh_token}
+        )
+        response.raise_for_status()
+        access_token = response.json()["access"]
+        return access_token
+
+    def _get_headers(self, access_token: str, refresh_token: str) -> dict[str, str]:
+        """Get authorization headers, refreshing access token as needed.
+
+        Args:
+            access_token (str): The access token.
+            refresh_token (str): The refresh token.
+
+        Raises:
+            SessionExpiredException: If refresh token is expired (401).
+
+        Returns:
+            dict[str, str]: The authorization headers,
+        """
+        if self._is_token_expired(access_token):
+            self.logger.info("[AUTH] Access token expired, refreshing...")
+            try:
+                access_token = self._refresh_access_token(refresh_token)
+                st.session_state["access_token"] = access_token
+                self.logger.success("[AUTH] Access token refreshed successfully")
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == httpx.codes.UNAUTHORIZED:
+                    self.logger.info("[AUTH] Refresh token expired")
+                    raise SessionExpiredException from e
+                raise  # Re-raise other HTTP errors
+
+        return {"Authorization": f"Bearer {access_token}"}
+
+    def authenticate(self, email: str, password: str) -> tuple[str|None, str|None, str]:
         """Send a post request to the authentication endpoint.
 
         Args:
             email (str): The email.
             password (str): The password.
 
         Returns:
-            tuple[str|None, str]: A tuple containing the access token and a status message.
+            tuple[str|None, str|None, str]:
+                A tuple containing the access token, the refresh token and a status message.
         """
         access_token = None
-
+        refresh_token = None
         message = "Ops! Ocorreu um erro durante o login. Por favor, tente novamente."
 
         try:
@@ -38,28 +109,30 @@ def authenticate(self, email: str, password: str) -> tuple[str|None, str]:
             response.raise_for_status()
 
             access_token = response.json().get("access")
+            refresh_token = response.json().get("refresh")
 
-            if access_token:
-                self.logger.success(f"[LOGIN] Successfully logged in")
+            if access_token and refresh_token:
+                self.logger.success(f"[AUTH] Successfully logged in")
                 message = "Conectado com sucesso!"
             else:
-                self.logger.error(f"[LOGIN] No access token returned")
+                self.logger.error(f"[AUTH] No access and refresh tokens returned")
         except httpx.HTTPStatusError:
             if response.status_code == httpx.codes.UNAUTHORIZED:
-                self.logger.warning(f"[LOGIN] Invalid credentials")
+                self.logger.warning(f"[AUTH] Invalid credentials")
                 message = "Usuário ou senha incorretos."
             else:
-                self.logger.exception(f"[LOGIN] HTTP error:")
+                self.logger.exception(f"[AUTH] HTTP error:")
         except Exception:
-            self.logger.exception(f"[LOGIN] Login error:")
+            self.logger.exception(f"[AUTH] Login error:")
 
-        return access_token, message
+        return access_token, refresh_token, message
 
-    def create_thread(self, access_token: str, title: str) -> Thread|None:
+    def create_thread(self, access_token: str, refresh_token: str, title: str) -> Thread|None:
         """Create a thread.
 
         Args:
             access_token (str): User access token.
+            refresh_token (str): User refresh token.
             title (str): The thread title.
 
         Returns:
@@ -71,21 +144,24 @@ def create_thread(self, access_token: str, title: str) -> Thread|None:
             response = httpx.post(
                 url=f"{self.base_url}/chatbot/threads/",
                 json={"title": title},
-                headers={"Authorization": f"Bearer {access_token}"},
+                headers=self._get_headers(access_token, refresh_token),
             )
             response.raise_for_status()
             thread = Thread(**response.json())
             self.logger.success(f"[THREAD] Thread created successfully for user {thread.account}")
             return thread
+        except SessionExpiredException:
+            raise
         except Exception:
             self.logger.exception(f"[THREAD] Error on thread creation:")
             return None
 
-    def get_threads(self, access_token: str) -> list[Thread]|None:
+    def get_threads(self, access_token: str, refresh_token: str) -> list[Thread]|None:
         """Get all threads from a user.
 
         Args:
             access_token (str): User access token.
+            refresh_token (str): User refresh token.
 
         Returns:
             list[Thread]|None: A list of Thread objects if any thread was found. None otherwise.
@@ -95,39 +171,54 @@ def get_threads(self, access_token: str) -> list[Thread]|None:
             response = httpx.get(
                 url=f"{self.base_url}/chatbot/threads/",
                 params={"order_by": "created_at"},
-                headers={"Authorization": f"Bearer {access_token}"}
+                headers=self._get_headers(access_token, refresh_token)
             )
             response.raise_for_status()
             threads = [Thread(**thread) for thread in response.json()]
             self.logger.success(f"[THREAD] Threads retrieved successfully")
             return threads
+        except SessionExpiredException:
+            raise
         except Exception:
             self.logger.exception(f"[THREAD] Error on threads retrieval:")
             return None
 
-    def get_message_pairs(self, access_token: str, thread_id: UUID) -> list[MessagePair]|None:
+    def get_message_pairs(self, access_token: str, refresh_token: str, thread_id: UUID) -> list[MessagePair]|None:
+        """Get all messages from a thread.
+
+        Args:
+            access_token (str): User access token.
+            refresh_token (str): User refresh token.
+            thread_id (UUID): Thread unique identifier.
+
+        Returns:
+            list[MessagePair]|None: A list of MessagePair objects if any message was found. None otherwise.
+        """
         self.logger.info(f"[MESSAGE] Retrieving message pairs for thread {thread_id}")
         try:
             response = httpx.get(
                 url=f"{self.base_url}/chatbot/threads/{thread_id}/messages/",
                 params={"order_by": "created_at"},
-                headers={"Authorization": f"Bearer {access_token}"}
+                headers=self._get_headers(access_token, refresh_token)
             )
             response.raise_for_status()
             message_pairs = [MessagePair(**pair) for pair in response.json()]
             self.logger.success(f"[MESSAGE] Message pairs retrieved successfully for thread {thread_id}")
             return message_pairs
+        except SessionExpiredException:
+            raise
         except Exception:
             self.logger.exception(f"[MESSAGE] Error on message pairs retrieval for thread {thread_id}:")
             return None
 
-    def send_message(self, access_token: str, message: str, thread_id: UUID) -> Iterator[StreamEvent]:
+    def send_message(self, access_token: str, refresh_token: str, message: str, thread_id: UUID) -> Iterator[StreamEvent]:
         """Send a user message and stream the assistant's response.
 
         Args:
-            access_token (str): The user's access token.
+            access_token (str): User access token.
+            refresh_token (str): User refresh token.
             message (str): The message sent by the user.
-            thread_id (UUID): The unique identifier of the thread.
+            thread_id (UUID):Thread unique identifier.
 
         Yields:
             Iterator[StreamEvent]: Iterator of `StreamEvent` objects.
@@ -143,7 +234,7 @@ def send_message(self, access_token: str, message: str, thread_id: UUID) -> Iter
             with httpx.stream(
                 method="POST",
                 url=f"{self.base_url}/chatbot/threads/{thread_id}/messages/",
-                headers={"Authorization": f"Bearer {access_token}"},
+                headers=self._get_headers(access_token, refresh_token),
                 json=user_message.model_dump(mode="json"),
                 timeout=httpx.Timeout(5.0, read=300.0),
             ) as response:
@@ -161,6 +252,8 @@ def send_message(self, access_token: str, message: str, thread_id: UUID) -> Iter
                         stream_completed = True
 
                     yield event
+        except SessionExpiredException:
+            raise
         except httpx.ReadTimeout:
             self.logger.exception(f"[MESSAGE] Timeout error on sending user message:")
             error_message=(
@@ -194,11 +287,12 @@ def send_message(self, access_token: str, message: str, thread_id: UUID) -> Iter
                 data=EventData(run_id=uuid4())
             )
 
-    def send_feedback(self, access_token: str, message_pair_id: UUID, rating: int, comments: str) -> bool:
+    def send_feedback(self, access_token: str, refresh_token: str, message_pair_id: UUID, rating: int, comments: str) -> bool:
         """Send a feedback.
 
         Args:
             access_token (str): User access token.
+            refresh_token (str): User refresh token.
             message_pair_id (UUID): The message pair unique identifier.
             rating (int): The rating (0 or 1).
             comments (str): The comments.
@@ -212,20 +306,23 @@ def send_feedback(self, access_token: str, message_pair_id: UUID, rating: int, c
             response = httpx.put(
                 url=f"{self.base_url}/chatbot/message-pairs/{message_pair_id}/feedbacks/",
                 json={"rating": rating, "comment": comments},
-                headers={"Authorization": f"Bearer {access_token}"}
+                headers=self._get_headers(access_token, refresh_token)
             )
             response.raise_for_status()
             self.logger.success(f"[FEEDBACK] Feedback sent successfully")
             return True
+        except SessionExpiredException:
+            raise
         except Exception:
             self.logger.exception(f"[FEEDBACK] Error on sending feedback:")
             return False
 
-    def delete_thread(self, access_token: str, thread_id: UUID) -> bool:
+    def delete_thread(self, access_token: str, refresh_token: str, thread_id: UUID) -> bool:
         """Soft delete a thread and hard delete all its checkpoints.
 
         Args:
             access_token (str): User access token.
+            refresh_token (str): User refresh token.
             thread_id (UUID): Thread unique identifier.
 
         Returns:
@@ -236,12 +333,14 @@ def delete_thread(self, access_token: str, thread_id: UUID) -> bool:
         try:
             response = httpx.delete(
                 url=f"{self.base_url}/chatbot/threads/{thread_id}/",
-                headers={"Authorization": f"Bearer {access_token}"},
+                headers=self._get_headers(access_token, refresh_token),
                 timeout=httpx.Timeout(5.0, read=60.0)
             )
             response.raise_for_status()
             self.logger.success(f"[CLEAR] Assistant memory cleared successfully")
             return True
+        except SessionExpiredException:
+            raise
         except Exception:
             self.logger.exception("[CLEAR] Error on clearing assistant memory:")
             return False