Fix retrieving parts from response

ag2ai · Jan 29, 2025 · 1d558cd · 1d558cd
1 parent 166fa10
commit 1d558cd
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/autogen/oai/gemini.py b/autogen/oai/gemini.py
@@ -224,9 +224,9 @@ def create(self, params: dict) -> ChatCompletion:
             if autogen_term in params
         }
         if self.use_vertexai:
-            safety_settings = GeminiClient._to_vertexai_safety_settings(params.get("safety_settings", {}))
+            safety_settings = GeminiClient._to_vertexai_safety_settings(params.get("safety_settings", []))
         else:
-            safety_settings = params.get("safety_settings", {})
+            safety_settings = params.get("safety_settings", [])
 
         if stream:
             warnings.warn(
@@ -281,7 +281,11 @@ def create(self, params: dict) -> ChatCompletion:
         prev_function_calls = []
 
         if isinstance(response, GenerateContentResponse):
-            parts = response.parts
+            if len(response.candidates) != 1:
+                raise ValueError(
+                    f"Unexpected number of candidates in the response. Expected 1, got {len(response.candidates)}"
+                )
+            parts = response.candidates[0].content.parts
         elif isinstance(response, VertexAIGenerationResponse):  # or hasattr(response, "candidates"):
             # google.generativeai also raises an error len(candidates) != 1:
             if len(response.candidates) != 1: