Merge pull request #419 from KCaverly/google_import_fix

fix: update imports to allow global dspy.Google

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

dsp/modules/google.py

@@ -8,7 +8,8 @@
     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/"""
@@ -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)

dspy/__init__.py

@@ -21,6 +21,7 @@
 ColBERTv2 = dsp.ColBERTv2
 Pyserini = dsp.PyseriniRetriever
 Clarifai = dsp.ClarifaiLLM
+Google = dsp.Google
 
 HFClientTGI = dsp.HFClientTGI
 HFClientVLLM = HFClientVLLM