feat: set up exception classes for better error handling

src/gemini/client.py

@@ -4,11 +4,19 @@
 from asyncio import Task
 from typing import Any, Optional
 
-from httpx import AsyncClient
+from httpx import AsyncClient, ReadTimeout
 from loguru import logger
 
 from .consts import HEADERS
-from .types import Image, Candidate, ModelOutput
+from .types import (
+    Image,
+    Candidate,
+    ModelOutput,
+    AuthError,
+    APIError,
+    GeminiError,
+    TimeoutError,
+)
 
 
 def running(func) -> callable:
@@ -102,7 +110,7 @@ async def init(
             response = await self.client.get("https://gemini.google.com/app")
 
             if response.status_code != 200:
-                raise Exception(
+                raise APIError(
                     f"Failed to initiate client. Request failed with status code {response.status_code}"
                 )
             else:
@@ -112,7 +120,7 @@ async def init(
                     self.running = True
                     logger.success("Gemini client initiated successfully.")
                 else:
-                    raise Exception(
+                    raise AuthError(
                         "Failed to initiate client. SNlM0e not found in response, make sure cookie values are valid."
                     )
 
@@ -171,18 +179,23 @@ async def generate_content(
         if self.auto_close:
             await self.reset_close_task()
 
-        response = await self.client.post(
-            "https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate",
-            data={
-                "at": self.access_token,
-                "f.req": json.dumps(
-                    [None, json.dumps([[prompt], None, chat and chat.metadata])]
-                ),
-            },
-        )
+        try:
+            response = await self.client.post(
+                "https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate",
+                data={
+                    "at": self.access_token,
+                    "f.req": json.dumps(
+                        [None, json.dumps([[prompt], None, chat and chat.metadata])]
+                    ),
+                },
+            )
+        except ReadTimeout:
+            raise TimeoutError(
+                "Request timed out, please try again. If the problem persists, consider setting a higher `timeout` value when initiating GeminiClient."
+            )
 
         if response.status_code != 200:
-            raise Exception(
+            raise APIError(
                 f"Failed to generate contents. Request failed with status code {response.status_code}"
             )
         else:
@@ -202,7 +215,7 @@ async def generate_content(
                     Candidate(rcid=candidate[0], text=candidate[1][0], images=images)
                 )
             if not candidates:
-                raise Exception(
+                raise GeminiError(
                     "Failed to generate contents. No output data found in response."
                 )
 

src/gemini/types.py

@@ -1,18 +1,29 @@
-from typing import Optional
-
 from pydantic import BaseModel
 
 
 class Image(BaseModel):
+    """
+    A single image object returned from Gemini.
+
+    Parameters
+    ----------
+    url: `str`
+        URL of the image
+    title: `str`, optional
+        Title of the image, by default is "[Image]"
+    alt: `str`, optional
+        Optional description of the image
+    """
+
     url: str
-    title: Optional[str] = "[Image]"
-    alt: Optional[str] = ""
+    title: str = "[Image]"
+    alt: str = ""
 
     def __str__(self):
         return f"{self.title}({self.url}) - {self.alt}"
 
     def __repr__(self):
-        return f"Image(title='{self.title}', url='{len(self.url)<=20 and self.url or self.url[:8] + '...' + self.url[-12:]}', alt='{self.alt}')"
+        return f"""Image(title='{self.title}', url='{len(self.url) <= 20 and self.url or self.url[:8] + '...' + self.url[-12:]}', alt='{self.alt}')"""
 
 
 class Candidate(BaseModel):
@@ -24,7 +35,7 @@ def __str__(self):
         return self.text
 
     def __repr__(self):
-        return f"Candidate(rcid='{self.rcid}', text='{len(self.text)<=20 and self.text or self.text[:20] + '...'}', images={self.images})"
+        return f"Candidate(rcid='{self.rcid}', text='{len(self.text) <= 20 and self.text or self.text[:20] + '...'}', images={self.images})"
 
 
 class ModelOutput(BaseModel):
@@ -62,3 +73,35 @@ def images(self):
     @property
     def rcid(self):
         return self.candidates[self.chosen].rcid
+
+
+class AuthError(Exception):
+    """
+    Exception for authentication errors caused by invalid credentials/cookies.
+    """
+
+    pass
+
+
+class APIError(Exception):
+    """
+    Exception for package-level errors which need to be fixed in the future development (e.g. validation errors).
+    """
+
+    pass
+
+
+class GeminiError(Exception):
+    """
+    Exception for errors returned from Gemini server which are not handled by the package.
+    """
+
+    pass
+
+
+class TimeoutError(GeminiError):
+    """
+    Exception for request timeouts.
+    """
+
+    pass