Add support for Google's GenAI models (e.g., gemini-pro)

Author: gustavocidornelas

openlayer/model_runners/ll_model_runners.py

@@ -15,6 +15,7 @@
 import pandas as pd
 import pybars
 import requests
+from google import generativeai
 from tqdm import tqdm
 
 from .. import constants
@@ -502,6 +503,84 @@ def _get_cost_estimate(self, response: Dict[str, Any]) -> None:
             )
 
 
+class GoogleGenAIModelRunner(LLModelRunner):
+    """Wraps Google's Gen AI models."""
+
+    def __init__(
+        self,
+        logger: Optional[logging.Logger] = None,
+        **kwargs,
+    ):
+        super().__init__(logger, **kwargs)
+        if kwargs.get("google_api_key") is None:
+            raise openlayer_exceptions.OpenlayerMissingLlmApiKey(
+                "Please pass your Google API key generated with "
+                "https://makersuite.google.com/  as the keyword argument"
+                " 'google_api_key'"
+            )
+        self.google_api_key = kwargs["google_api_key"]
+
+        self._initialize_llm()
+
+        self.cost: List[float] = []
+
+    def _initialize_llm(self):
+        """Initializes the OpenAI chat completion model."""
+        if self.model_config.get("model") is None:
+            warnings.warn("No model specified. Defaulting to model 'gemini-pro'.")
+        if self.model_config.get("model_parameters") is None:
+            warnings.warn("No model parameters specified. Using default parameters.")
+        # Check if API key is valid
+        try:
+            generativeai.configure(api_key=self.google_api_key)
+            self.model = generativeai.GenerativeModel(
+                self.model_config.get("model", "gemini-pro")
+            )
+        except Exception as e:
+            raise openlayer_exceptions.OpenlayerInvalidLlmApiKey(
+                "Please pass your Google API key generated with "
+                "https://makersuite.google.com/  as the keyword argument"
+                f" 'google_api_key' \n Error message: {e}"
+            ) from e
+
+    def _get_llm_input(
+        self, injected_prompt: List[Dict[str, str]]
+    ) -> List[Dict[str, str]]:
+        """Prepares the input for Google's model."""
+        llm_input = ""
+        for message in injected_prompt:
+            if message["role"] == "system":
+                llm_input += f"S: {message['content']} \n"
+            elif message["role"] == "assistant":
+                llm_input += f"A: {message['content']} \n"
+            elif message["role"] == "user":
+                llm_input += f"U: {message['content']} \n"
+            else:
+                raise ValueError(
+                    "Message role must be either 'system', 'assistant' or 'user'. "
+                    f"Got: {message['role']}"
+                )
+        llm_input += "A:"
+        return llm_input
+
+    def _make_request(self, llm_input: List[Dict[str, str]]) -> Dict[str, Any]:
+        """Make the request to Google's model
+        for a given input."""
+        response = self.model.generate_content(
+            contents=llm_input,
+            **self.model_config.get("model_parameters", {}),
+        )
+        return response
+
+    def _get_output(self, response: Dict[str, Any]) -> str:
+        """Gets the output from the response."""
+        return response.text
+
+    def _get_cost_estimate(self, response: Dict[str, Any]) -> None:
+        """Estimates the cost from the response."""
+        return 0
+
+
 class SelfHostedLLModelRunner(LLModelRunner):
     """Wraps a self-hosted LLM."""
 

openlayer/models.py

@@ -111,6 +111,7 @@ class ModelRunnerFactory:
         "OpenAI": ll_model_runners.OpenAIChatCompletionRunner,
         "SelfHosted": ll_model_runners.SelfHostedLLModelRunner,
         "HuggingFace": ll_model_runners.HuggingFaceModelRunner,
+        "Google": ll_model_runners.GoogleGenAIModelRunner,
     }
     _MODEL_RUNNERS = {
         tasks.TaskType.TabularClassification.value: traditional_ml_model_runners.ClassificationModelRunner,
@@ -176,8 +177,8 @@ def _create_ll_model_runner(
             raise exceptions.OpenlayerUnsupportedLlmProvider(
                 provider=model_provider,
                 message="\nCurrently, the supported providers are: 'OpenAI', 'Cohere',"
-                " 'Anthropic', 'SelfHosted', 'HuggingFace'. Reach out if you'd like us"
-                " to support your use case.",
+                " 'Anthropic', 'SelfHosted', 'HuggingFace', and 'Google'."
+                " Reach out if you'd like us to support your use case.",
             )
 
         model_runner_class = ModelRunnerFactory._LL_MODEL_RUNNERS[task_type.value][

setup.cfg

@@ -41,6 +41,7 @@ packages =
 install_requires =
     anthropic
     cohere
+    google-generativeai
     marshmallow
     marshmallow_oneofschema
     openai>=1.0.0