From 303df621da8aef3801b6724ed131e96c8780ba72 Mon Sep 17 00:00:00 2001
From: UZQueen <157540577+HanaokaYuzu@users.noreply.github.com>
Date: Sat, 7 Sep 2024 15:39:25 -0500
Subject: [PATCH] fix: image generation response structure

ref #34
---
 src/gemini_webapi/client.py | 31 ++++++++++++++++++-------------
 tests/test_save_image.py    |  2 +-
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/gemini_webapi/client.py b/src/gemini_webapi/client.py
index 1c29329..35ec81c 100644
--- a/src/gemini_webapi/client.py
+++ b/src/gemini_webapi/client.py
@@ -345,12 +345,14 @@ async def generate_content(
             )
         else:
             try:
+                response_json = json.loads(response.text.split("\n")[2])
+
                 # Plain request
-                body = json.loads(json.loads(response.text.split("\n")[2])[0][2])
+                body = json.loads(response_json[0][2])
 
                 if not body[4]:
                     # Request with Gemini extensions enabled
-                    body = json.loads(json.loads(response.text.split("\n")[2])[4][2])
+                    body = json.loads(response_json[4][2])
 
                 if not body[4]:
                     raise Exception
@@ -363,9 +365,11 @@ async def generate_content(
 
             try:
                 candidates = []
-                for candidate in body[4]:
+                for i, candidate in enumerate(body[4]):
                     text = candidate[1][0]
-                    if re.match(r"^http://googleusercontent.com/card_content/\d+$", text):
+                    if re.match(
+                        r"^http://googleusercontent.com/card_content/\d+$", text
+                    ):
                         text = candidate[22] and candidate[22][0] or text
 
                     web_images = (
@@ -383,11 +387,11 @@ async def generate_content(
                         or []
                     )
 
-                    generated_images = (
-                        candidate[12]
-                        and candidate[12][7]
-                        and candidate[12][7][0]
-                        and [
+                    generated_images = []
+                    if candidate[12] and candidate[12][7] and candidate[12][7][0]:
+                        image_generation_body = json.loads(response_json[1][2])
+                        image_generation_candidate = image_generation_body[4][i]
+                        generated_images = [
                             GeneratedImage(
                                 url=image[0][3][3],
                                 title=f"[Generated Image {image[3][6]}]",
@@ -397,10 +401,10 @@ async def generate_content(
                                 proxies=self.proxies,
                                 cookies=self.cookies,
                             )
-                            for i, image in enumerate(candidate[12][7][0])
-                        ]
-                        or []
-                    )
+                            for i, image in enumerate(
+                                image_generation_candidate[12][7][0]
+                            )
+                        ] or []
 
                     candidates.append(
                         Candidate(
@@ -417,6 +421,7 @@ async def generate_content(
 
                 output = ModelOutput(metadata=body[1], candidates=candidates)
             except (TypeError, IndexError):
+                logger.debug(f"Invalid response: {response.text}")
                 raise APIError(
                     "Failed to parse response body. Data structure is invalid."
                 )
diff --git a/tests/test_save_image.py b/tests/test_save_image.py
index d8f40e5..211c706 100644
--- a/tests/test_save_image.py
+++ b/tests/test_save_image.py
@@ -24,7 +24,7 @@ async def asyncSetUp(self):
 
     async def test_save_web_image(self):
         response = await self.geminiclient.generate_content(
-            "Send me 10 pictures of random subjects"
+            "Show me some pictures of random subjects"
         )
         self.assertTrue(response.images)
         for i, image in enumerate(response.images):