feat: v0.4.0 release

- fix: update index of GeneratedImage.alt
- fix: now secure_1psidts will only be passed to cookies if it is not empty
- feat: add an optional parameter to Image.save() to skip images with invalid file names
- feat: Image.save() now returns saved file path on success
- refactor: code reorganization and cleanup
- test: updated unit tests

README.md

@@ -117,13 +117,13 @@ Note: by default, when asked to send images (like the previous example), Gemini
 
 ### Save images to local files
 
-You can save images returned from Gemini to local files under `/temp` by calling `Image.save()`. Optionally, you can specify the file path and file name by passing `path` and `filename` arguments to the function. Works for both `WebImage` and `GeneratedImage`.
+You can save images returned from Gemini to local files under `/temp` by calling `Image.save()`. Optionally, you can specify the file path and file name by passing `path` and `filename` arguments to the function and skip images with invalid file names by passing `skip_invalid_filename=True`. Works for both `WebImage` and `GeneratedImage`.
 
 ```python
 async def main():
  response = await client.generate_content("Generate some pictures of cats")
  for i, image in enumerate(response.images):
- await image.save(path="temp/", filename=f"cat_{i}.png")
+ await image.save(path="temp/", filename=f"cat_{i}.png", verbose=True)
 
 asyncio.run(main())
 ```

src/gemini/__init__.py

@@ -1,2 +1,3 @@
 from .client import GeminiClient, ChatSession # noqa: F401
+from .exceptions import * # noqa: F401, F403
 from .types import * # noqa: F401, F403

src/gemini/client.py

@@ -7,17 +7,9 @@
 from httpx import AsyncClient, ReadTimeout
 from loguru import logger
 
-from .consts import HEADERS
-from .types import (
- WebImage,
- GeneratedImage,
- Candidate,
- ModelOutput,
- AuthError,
- APIError,
- GeminiError,
- TimeoutError,
-)
+from .types import WebImage, GeneratedImage, Candidate, ModelOutput
+from .exceptions import APIError, AuthError, TimeoutError, GeminiError
+from .constant import HEADERS
 
 
 def running(func) -> callable:
@@ -71,10 +63,7 @@ def __init__(
  secure_1psidts: Optional[str] = None,
  proxy: Optional[dict] = None,
  ):
- self.cookies = {
- "__Secure-1PSID": secure_1psid,
- "__Secure-1PSIDTS": secure_1psidts,
- }
+ self.cookies = {"__Secure-1PSID": secure_1psid}
  self.proxy = proxy
  self.client: AsyncClient | None = None
  self.access_token: Optional[str] = None
@@ -83,6 +72,9 @@ def __init__(
  self.close_delay: float = 300
  self.close_task: Task | None = None
 
+ if secure_1psidts:
+ self.cookies["__Secure-1PSIDTS"] = secure_1psidts
+
  async def init(
  self, timeout: float = 30, auto_close: bool = False, close_delay: float = 300
  ) -> None:
@@ -248,7 +240,9 @@ async def generate_content(
  GeneratedImage(
  url=image[0][3][3],
  title=f"[Generated Image {image[3][6]}]",
- alt=image[3][5][i],
+ alt=len(image[3][5]) > i
+ and image[3][5][i]
+ or image[3][5][0],
  cookies=self.cookies,
  )
  for i, image in enumerate(candidate[12][7][0])

src/gemini/exceptions.py

@@ -0,0 +1,30 @@
+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

src/gemini/types/__init__.py

@@ -0,0 +1,3 @@
+from .image import Image, WebImage, GeneratedImage # noqa: F401
+from .candidate import Candidate # noqa: F401
+from .modeloutput import ModelOutput # noqa: F401

src/gemini/types/candidate.py

@@ -0,0 +1,35 @@
+from pydantic import BaseModel
+
+from .image import Image, WebImage, GeneratedImage
+
+
+class Candidate(BaseModel):
+ """
+ A single reply candidate object in the model output. A full response from Gemini usually contains multiple reply candidates.
+
+ Parameters
+ ----------
+ rcid: `str`
+ Reply candidate ID to build the metadata
+ text: `str`
+ Text output
+ web_images: `list[WebImage]`, optional
+ List of web images in reply, can be empty.
+ generated_images: `list[GeneratedImage]`, optional
+ List of generated images in reply, can be empty
+ """
+
+ rcid: str
+ text: str
+ web_images: list[WebImage] = []
+ generated_images: list[GeneratedImage] = []
+
+ 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})"
+
+ @property
+ def images(self) -> list[Image]:
+ return self.web_images + self.generated_images

src/gemini/types.py → src/gemini/types/image.py

@@ -36,8 +36,9 @@ async def save(
  path: str = "temp",
  filename: str | None = None,
  cookies: dict | None = None,
- verbose: bool = False
- ) -> None:
+ verbose: bool = False,
+ skip_invalid_filename: bool = False,
+ ) -> str | None:
  """
  Save the image to disk.
 
@@ -46,22 +47,27 @@ async def save(
  path: `str`, optional
  Path to save the image, by default will save to ./temp
  filename: `str`, optional
- Filename to save the image, by default will use the original filename from the URL
+ File name to save the image, by default will use the original file name from the URL
  cookies: `dict`, optional
  Cookies used for requesting the content of the image
  verbose : `bool`, optional
- If True, print the path of the saved file, by default False
+ If True, print the path of the saved file or warning for invalid file name, by default False
+ skip_invalid_filename: `bool`, optional
+ If True, will only save the image if the file name and extension are valid, by default False
+
+ Returns
+ -------
+ `str | None`
+ Absolute path of the saved image if successful, None if filename is invalid and `skip_invalid_filename` is True
  """
+ filename = filename or self.url.split("/")[-1].split("?")[0]
  try:
- filename = (
- filename
- or (
- re.search(r"^(.*\.\w+)", self.url.split("/")[-1])
- or re.search(r"^(.*)\?", self.url.split("/")[-1])
- ).group()
- )
+ filename = re.search(r"^(.*\.\w+)", filename).group()
  except AttributeError:
- filename = self.url.split("/")[-1]
+ if verbose:
+ logger.warning(f"Invalid filename: {filename}")
+ if skip_invalid_filename:
+ return None
 
  async with AsyncClient(follow_redirects=True, cookies=cookies) as client:
  response = await client.get(self.url)
@@ -80,6 +86,8 @@ async def save(
 
  if verbose:
  logger.info(f"Image saved as {dest.resolve()}")
+
+ return dest.resolve()
  else:
  raise HTTPError(
  f"Error downloading image: {response.status_code} {response.reason_phrase}"
@@ -110,129 +118,28 @@ class GeneratedImage(Image):
  @field_validator("cookies")
  @classmethod
  def validate_cookies(cls, v: dict) -> dict:
- if "__Secure-1PSID" not in v or "__Secure-1PSIDTS" not in v:
+ if len(v) == 0:
  raise ValueError(
- "Cookies must contain '__Secure-1PSID' and '__Secure-1PSIDTS'"
+ "GeneratedImage is designed to be initiated with same cookies as GeminiClient."
  )
  return v
 
  # @override
- async def save(self, path: str = "temp/", filename: str = None) -> None:
+ async def save(self, **kwargs) -> None:
  """
  Save the image to disk.
 
  Parameters
  ----------
- path: `str`
- Path to save the image
  filename: `str`, optional
  Filename to save the image, generated images are always in .png format, but file extension will not be included in the URL.
  And since the URL ends with a long hash, by default will use timestamp + end of the hash as the filename
+ **kwargs: `dict`, optional
+ Other arguments to pass to `Image.save`
  """
  await super().save(
- path,
- filename
+ filename=kwargs.pop("filename", None)
  or f"{datetime.now().strftime('%Y%m%d%H%M%S')}_{self.url[-10:]}.png",
- self.cookies,
+ cookies=self.cookies,
+ **kwargs,
  )
-
-
-class Candidate(BaseModel):
- """
- A single reply candidate object in the model output. A full response from Gemini usually contains multiple reply candidates.
-
- Parameters
- ----------
- rcid: `str`
- Reply candidate ID to build the metadata
- text: `str`
- Text output
- web_images: `list[WebImage]`, optional
- List of web images in reply, can be empty.
- generated_images: `list[GeneratedImage]`, optional
- List of generated images in reply, can be empty
- """
-
- rcid: str
- text: str
- web_images: list[WebImage] = []
- generated_images: list[GeneratedImage] = []
-
- 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})"
-
- @property
- def images(self) -> list[Image]:
- return self.web_images + self.generated_images
-
-
-class ModelOutput(BaseModel):
- """
- Classified output from gemini.google.com
-
- Parameters
- ----------
- metadata: `list[str]`
- List of chat metadata `[cid, rid, rcid]`, can be shorter than 3 elements, like `[cid, rid]` or `[cid]` only
- candidates: `list[Candidate]`
- List of all candidates returned from gemini
- chosen: `int`, optional
- Index of the chosen candidate, by default will choose the first one
- """
-
- metadata: list[str]
- candidates: list[Candidate]
- chosen: int = 0
-
- def __str__(self):
- return self.text
-
- def __repr__(self):
- return f"ModelOutput(metadata={self.metadata}, chosen={self.chosen}, candidates={self.candidates})"
-
- @property
- def text(self) -> str:
- return self.candidates[self.chosen].text
-
- @property
- def images(self) -> list[Image]:
- return self.candidates[self.chosen].images
-
- @property
- def rcid(self) -> str:
- 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

src/gemini/types/modeloutput.py

@@ -0,0 +1,41 @@
+from pydantic import BaseModel
+
+from .image import Image
+from .candidate import Candidate
+
+
+class ModelOutput(BaseModel):
+ """
+ Classified output from gemini.google.com
+
+ Parameters
+ ----------
+ metadata: `list[str]`
+ List of chat metadata `[cid, rid, rcid]`, can be shorter than 3 elements, like `[cid, rid]` or `[cid]` only
+ candidates: `list[Candidate]`
+ List of all candidates returned from gemini
+ chosen: `int`, optional
+ Index of the chosen candidate, by default will choose the first one
+ """
+
+ metadata: list[str]
+ candidates: list[Candidate]
+ chosen: int = 0
+
+ def __str__(self):
+ return self.text
+
+ def __repr__(self):
+ return f"ModelOutput(metadata={self.metadata}, chosen={self.chosen}, candidates={self.candidates})"
+
+ @property
+ def text(self) -> str:
+ return self.candidates[self.chosen].text
+
+ @property
+ def images(self) -> list[Image]:
+ return self.candidates[self.chosen].images
+
+ @property
+ def rcid(self) -> str:
+ return self.candidates[self.chosen].rcid

tests/test_client_features.py

@@ -67,7 +67,11 @@ async def test_reply_candidates(self):
  response = await chat.send_message(
  "What's the best Japanese dish? Recommend one only."
  )
- self.assertTrue(len(response.candidates) > 1)
+
+ if len(response.candidates) == 1:
+ logger.debug(response.candidates[0])
+ self.skipTest("Only one candidate was returned. Test skipped")
+
  for candidate in response.candidates:
  logger.debug(candidate)