From b0c968bfe5ec0a16e3754cc5da67a3f792f0960d Mon Sep 17 00:00:00 2001
From: Ivan Leo <ivanleomk@gmail.com>
Date: Wed, 25 Dec 2024 10:09:10 +0800
Subject: [PATCH] Adding support for Gemini-8b and 2.0 (#1274)

---
 .github/workflows/test.yml    | 5 +++++
 instructor/reask.py           | 9 +++++++++
 tests/llm/test_gemini/util.py | 2 +-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8dcdab16f..f325a570f 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -43,6 +43,11 @@ jobs:
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
           COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
 
+      - name: Run Gemini Tests
+        run: poetry run pytest tests/llm/test_gemini
+        env:
+          GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+
       - name: Generate coverage report
         if: matrix.python-version == '3.11'
         run: |
diff --git a/instructor/reask.py b/instructor/reask.py
index 1684d48b4..42da71a6d 100644
--- a/instructor/reask.py
+++ b/instructor/reask.py
@@ -107,6 +107,15 @@ def reask_gemini_tools(
     from google.ai import generativelanguage as glm  # type: ignore
 
     reask_msgs = [
+        {
+            "role": "model",
+            "parts": [
+                glm.FunctionCall(
+                    name=response.parts[0].function_call.name,
+                    args=response.parts[0].function_call.args,
+                )
+            ],
+        },
         {
             "role": "function",
             "parts": [
diff --git a/tests/llm/test_gemini/util.py b/tests/llm/test_gemini/util.py
index c19ea9f7f..0c4757151 100644
--- a/tests/llm/test_gemini/util.py
+++ b/tests/llm/test_gemini/util.py
@@ -1,4 +1,4 @@
 import instructor
 
-models: list[str] = ["models/gemini-1.5-flash-latest"]
+models: list[str] = ["models/gemini-1.5-flash-8b"]
 modes = [instructor.Mode.GEMINI_TOOLS, instructor.Mode.GEMINI_JSON]