From e7c5f9f4511e8decb47a556b25bd645991985e83 Mon Sep 17 00:00:00 2001
From: KCaverly <kylebcaverly@gmail.com>
Date: Tue, 20 Feb 2024 11:35:04 -0500
Subject: [PATCH 1/3] fix: update imports to allow global dspy.Google

---
 dsp/modules/__init__.py | 1 +
 dspy/__init__.py        | 1 +
 2 files changed, 2 insertions(+)

diff --git a/dsp/modules/__init__.py b/dsp/modules/__init__.py
index 1fb45bdac..469f6fee4 100644
--- a/dsp/modules/__init__.py
+++ b/dsp/modules/__init__.py
@@ -9,6 +9,7 @@
 from .ollama import *
 from .clarifai import *
 from .bedrock import *
+from .google import *
 
 
 from .hf_client import HFClientTGI
diff --git a/dspy/__init__.py b/dspy/__init__.py
index c30295518..b5fadb4e0 100644
--- a/dspy/__init__.py
+++ b/dspy/__init__.py
@@ -21,6 +21,7 @@
 ColBERTv2 = dsp.ColBERTv2
 Pyserini = dsp.PyseriniRetriever
 Clarifai = dsp.ClarifaiLLM
+Google = dsp.Google
 
 HFClientTGI = dsp.HFClientTGI
 HFClientVLLM = HFClientVLLM

From 22ea0da7eb1fc730318a7e618cb52b67858e3f8f Mon Sep 17 00:00:00 2001
From: kcaverly <kylebcaverly@gmail.com>
Date: Tue, 20 Feb 2024 11:30:31 -0500
Subject: [PATCH 2/3] Removed print google install error

---
 dsp/modules/google.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dsp/modules/google.py b/dsp/modules/google.py
index 581c98414..4cf15033b 100644
--- a/dsp/modules/google.py
+++ b/dsp/modules/google.py
@@ -8,7 +8,7 @@
     import google.generativeai as genai
 except ImportError:
     google_api_error = Exception
-    print("Not loading Google because it is not installed.")
+    # print("Not loading Google because it is not installed.")
 
 def backoff_hdlr(details):
     """Handler from https://pypi.org/project/backoff/"""

From 96949a76db635501a3ea45a9db94096160e90fb5 Mon Sep 17 00:00:00 2001
From: KCaverly <kylebcaverly@gmail.com>
Date: Wed, 21 Feb 2024 15:25:19 -0500
Subject: [PATCH 3/3] fix api_key error

---
 dsp/modules/google.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/dsp/modules/google.py b/dsp/modules/google.py
index 4cf15033b..9b4516e9c 100644
--- a/dsp/modules/google.py
+++ b/dsp/modules/google.py
@@ -10,6 +10,7 @@
     google_api_error = Exception
     # print("Not loading Google because it is not installed.")
 
+
 def backoff_hdlr(details):
     """Handler from https://pypi.org/project/backoff/"""
     print(
@@ -33,10 +34,7 @@ class Google(LM):
     """
 
     def __init__(
-        self,
-        model: str = "gemini-pro-1.0",
-        api_key: Optional[str] = None,
-        **kwargs
+        self, model: str = "gemini-pro-1.0", api_key: Optional[str] = None, **kwargs
     ):
         """
         Parameters
@@ -51,15 +49,17 @@ def __init__(
             Additional arguments to pass to the API provider.
         """
         super().__init__(model)
-        self.google = genai.configure(api_key=self.api_key)
+        self.google = genai.configure(api_key=api_key)
         self.provider = "google"
         self.kwargs = {
             "model_name": model,
-            "temperature": 0.0 if "temperature" not in kwargs else kwargs["temperature"],
+            "temperature": 0.0
+            if "temperature" not in kwargs
+            else kwargs["temperature"],
             "max_output_tokens": 2048,
             "top_p": 1,
             "top_k": 1,
-            **kwargs
+            **kwargs,
         }
 
         self.history: list[dict[str, Any]] = []
@@ -85,7 +85,7 @@ def basic_request(self, prompt: str, **kwargs):
 
     @backoff.on_exception(
         backoff.expo,
-        (google_api_error),
+        (Exception),
         max_time=1000,
         on_backoff=backoff_hdlr,
         giveup=giveup_hdlr,
@@ -99,6 +99,6 @@ def __call__(
         prompt: str,
         only_completed: bool = True,
         return_sorted: bool = False,
-        **kwargs
+        **kwargs,
     ):
         return self.request(prompt, **kwargs)