Google PaLM support

prr/service_registry.py

@@ -1,4 +1,6 @@
 from prr.services.providers.anthropic.complete import ServiceAnthropicComplete
+from prr.services.providers.google.chat import ServiceGoogleChat
+from prr.services.providers.google.complete import ServiceGoogleComplete
 from prr.services.providers.openai.chat import ServiceOpenAIChat
 
 
@@ -13,6 +15,8 @@ def register(self, service_class):
     def register_all_services(self):
         self.register(ServiceOpenAIChat)
         self.register(ServiceAnthropicComplete)
+        self.register(ServiceGoogleComplete)
+        self.register(ServiceGoogleChat)
 
     def service_for_service_config(self, service_config):
         key = service_config.service_key()

prr/services/providers/anthropic/complete.py

@@ -26,7 +26,7 @@ def run(self, prompt, service_config):
 
         prompt_text = self.prompt_text()
 
-        client = anthropic.Client(config["ANTHROPIC_API_KEY"])
+        client = anthropic.Client(config.get("ANTHROPIC_API_KEY", None))
 
         service_request = ServiceRequest(self.service_config, prompt_text)
 

prr/services/providers/google/chat.py

@@ -0,0 +1,104 @@
+from google.cloud import aiplatform
+from vertexai.preview.language_models import ChatModel, InputOutputTextPair
+
+from prr.config import load_config
+from prr.request import ServiceRequest
+from prr.response import ServiceResponse
+
+config = load_config()
+
+aiplatform.init(
+    # your Google Cloud Project ID or number
+    # environment default used is not set
+    project=config.get("GOOGLE_PROJECT"),
+    # the Vertex AI region you will use
+    # defaults to us-central1
+    location=config.get("GOOGLE_LOCATION"),
+    # custom google.auth.credentials.Credentials
+    # environment default creds used if not set
+    # credentials=config[my_credentials],
+)
+
+
+class ServiceGoogleChat:
+    provider = "google"
+    service = "chat"
+
+    def run(self, prompt, service_config):
+        self.service_config = service_config
+        options = self.service_config.options
+        self.prompt = prompt
+
+        messages = self.messages_from_prompt()
+
+        model = ChatModel.from_pretrained(self.service_config.model_name())
+
+        service_request = ServiceRequest(self.service_config, {"messages": messages})
+
+        parameters = {
+            "temperature": options.temperature,
+            "max_output_tokens": options.max_tokens,
+            "top_p": options.top_p,
+            "top_k": options.top_k,
+        }
+
+        chat = model.start_chat(
+            context=f"""{self.context_from_messages()}""",
+            examples=[],
+        )
+
+        response = chat.send_message(
+            f"""{self.message_from_messages()}""", **parameters
+        )
+
+        service_response = ServiceResponse(response.text, {})
+
+        return service_request, service_response
+
+    def messages_from_prompt(self):
+        messages = self.prompt.messages
+
+        # prefer messages in prompt if they exist
+        if messages:
+            return messages
+
+        return [{"role": "user", "content": self.prompt.text()}]
+
+    def context_from_messages(self):
+        messages = self.prompt.messages
+
+        if messages:
+            for message in messages:
+                if message["role"] == "system":
+                    context = message["content"]
+        else:
+            context = None
+
+        return context
+
+    def examples_from_messages(self):
+        messages = self.prompt.messages
+
+        if messages:
+            examples = []
+            for message in messages:
+                if message["role"] == "examples":
+                    examples.append(
+                        InputOutputTextPair(message["input"], message["output"])
+                    )
+        else:
+            examples = []
+
+        return examples
+
+    def message_from_messages(self):
+        messages = self.prompt.messages
+
+        if messages:
+            for message in messages:
+                if message["role"] == "user":
+                    message = message["content"]
+        else:
+            message = None
+
+        return message

prr/services/providers/google/complete.py

@@ -0,0 +1,70 @@
+from google.cloud import aiplatform
+from vertexai.preview.language_models import TextGenerationModel
+
+from prr.config import load_config
+from prr.request import ServiceRequest
+from prr.response import ServiceResponse
+
+config = load_config()
+
+aiplatform.init(
+    # your Google Cloud Project ID or number
+    # environment default used is not set
+    project=config.get("GOOGLE_PROJECT", None),
+    # the Vertex AI region you will use
+    # defaults to us-central1
+    location=config.get("GOOGLE_LOCATION", None),
+    # custom google.auth.credentials.Credentials
+    # environment default creds used if not set
+    # credentials=config[my_credentials],
+)
+
+
+class ServiceGoogleComplete:
+    provider = "google"
+    service = "complete"
+
+    def run(self, prompt, service_config):
+        self.service_config = service_config
+        options = self.service_config.options
+        self.prompt = prompt
+
+        prompt_text = self.prompt_text()
+
+        client = TextGenerationModel.from_pretrained(self.service_config.model_name())
+
+        service_request = ServiceRequest(self.service_config, prompt_text)
+
+        options = self.service_config.options
+
+        response = client.predict(
+            prompt_text,
+            max_output_tokens=options.max_tokens,
+            temperature=options.temperature,
+            top_k=options.top_k,
+            top_p=options.top_p,
+        )
+
+        service_response = ServiceResponse(
+            str(response),
+            {
+                "details": response,
+            },
+        )
+
+        return service_request, service_response
+
+    def prompt_text(self):
+        messages = self.prompt.messages
+
+        prompt_text = ""
+
+        if messages:
+            for message in messages:
+                if message["role"] != "assistant":
+                    prompt_text += " " + message["content"]
+
+        else:
+            prompt_text = self.prompt.text()
+
+        return messages

prr/services/providers/openai/chat.py

@@ -5,7 +5,7 @@
 from prr.response import ServiceResponse
 
 config = load_config()
-openai.api_key = config["OPENAI_API_KEY"]
+openai.api_key = config.get("OPENAI_API_KEY", None)
 
 
 # OpenAI model provider class

pyproject.toml

@@ -14,6 +14,7 @@ anthropic = "^0.2.7"
 rich = "^13.3.5"
 Jinja2 = "^3.1.2"
 pyyaml = "^6.0"
+google-cloud-aiplatform = "^1.25.0"
 
 [tool.poetry.dev-dependencies]
 pre-commit = "^2.12.0"