use system instruction for gemini

djcopley · Feb 9, 2025 · c54d6ee · c54d6ee
1 parent 462def9
commit c54d6ee
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/src/shelloracle/providers/google.py b/src/shelloracle/providers/google.py
@@ -9,23 +9,19 @@ class Google(Provider):
     name = "Google"
 
     api_key = Setting(default="")
-    model = Setting(default="gemini-pro")  # Assuming a default model name
+    model = Setting(default="gemini-2.0-flash")  # Assuming a default model name
 
     def __init__(self):
         if not self.api_key:
             msg = "No API key provided"
             raise ProviderError(msg)
         genai.configure(api_key=self.api_key)
-        self.model_instance = genai.GenerativeModel(self.model)
+        self.model_instance = genai.GenerativeModel(self.model, system_instruction=system_prompt)
 
     async def generate(self, prompt: str) -> AsyncIterator[str]:
         try:
             response = await self.model_instance.generate_content_async(
-                [
-                    {"role": "user", "parts": [system_prompt]},
-                    {"role": "model", "parts": ["Okay."]},  # Gemini requires a model response before user input
-                    {"role": "user", "parts": [prompt]},
-                ],
+                [prompt],
                 stream=True,
             )