From 5f2802292acae04f3e0b14578e0d578dc5b77166 Mon Sep 17 00:00:00 2001
From: Yiheng Xu <github@ranpox.com>
Date: Fri, 22 Mar 2024 12:54:22 +0800
Subject: [PATCH] Update agent.py

---
 mm_agents/agent.py | 38 +++++++++++++++++++++-----------------
 1 file changed, 21 insertions(+), 17 deletions(-)

diff --git a/mm_agents/agent.py b/mm_agents/agent.py
index 4600628..4b27968 100644
--- a/mm_agents/agent.py
+++ b/mm_agents/agent.py
@@ -741,26 +741,30 @@ class PromptAgent:
             assert api_key is not None, "Please set the GENAI_API_KEY environment variable"
             genai.configure(api_key=api_key)
             logger.info("Generating content with Gemini model: %s", self.model)
-            response = genai.GenerativeModel(self.model).generate_content(
-                gemini_messages,
-                generation_config={
-                    "candidate_count": 1,
-                    "max_output_tokens": max_tokens,
-                    "top_p": top_p,
-                    "temperature": temperature
-                },
-                safety_settings={
-                    "harassment": "block_none",
-                    "hate": "block_none",
-                    "sex": "block_none",
-                    "danger": "block_none"
-                }
-            )
-
+            request_options = {"timeout": 120}
+            gemini_model = genai.GenerativeModel(self.model)
             try:
+                response = gemini_model.generate_content(
+                    gemini_messages,
+                    generation_config={
+                        "candidate_count": 1,
+                        "max_output_tokens": max_tokens,
+                        "top_p": top_p,
+                        "temperature": temperature
+                    },
+                    safety_settings={
+                        "harassment": "block_none",
+                        "hate": "block_none",
+                        "sex": "block_none",
+                        "danger": "block_none"
+                    },
+                    request_options=request_options
+                )
                 return response.text
             except Exception as e:
-                logger.error("Meet exception when calling Gemini API, " + str(e))
+                logger.error("Meet exception when calling Gemini API, " + str(e.__class__.__name__) + str(e))
+                logger.error(f"count_tokens: {gemini_model.count_tokens(gemini_messages)}")
+                logger.error(f"generation_config: {max_tokens}, {top_p}, {temperature}")
                 return ""
         elif self.model.startswith("qwen"):
             messages = payload["messages"]