specs

Author: andreibondarev

CHANGELOG.md

@@ -3,3 +3,4 @@
 ## [0.1.0] - 2023-05-23
 
 - Initial release
+  - Added `GooglePalmApi::Client` class that supports `embed()`, `generate_text()`, `generate_chat_message()`, `get_model()`, `list_models()` and `count_message_tokens()` methods

lib/google_palm_api/client.rb

@@ -8,7 +8,6 @@ class Client
     attr_reader :api_key, :connection
 
     ENDPOINT_URL = "https://generativelanguage.googleapis.com/"
-    AUTOPUSH_ENDPOINT_URL = "https://autopush-generativelanguage.sandbox.googleapis.com/"
 
     DEFAULTS = {
       temperature: 0.0,
@@ -51,11 +50,11 @@ def generate_text(
       stop_sequences: nil,
       client: nil
     )
-      response = connection(url: ENDPOINT_URL).post("/v1beta2/models/#{model || DEFAULTS[:completion_model_name]}:generateText") do |req|
+      response = connection.post("/v1beta2/models/#{model || DEFAULTS[:completion_model_name]}:generateText") do |req|
         req.params = {key: api_key}
 
         req.body = {prompt: { text: prompt }}
-        req.body[:temperature] = temperature if temperature
+        req.body[:temperature] = temperature || DEFAULTS[:temperature]
         req.body[:candidate_count] = candidate_count if candidate_count
         req.body[:max_output_tokens] = max_output_tokens if max_output_tokens
         req.body[:top_p] = top_p if top_p
@@ -99,15 +98,15 @@ def generate_chat_message(
       client: nil
     )
       # Overwrite the default ENDPOINT_URL for this method.
-      response = connection(url: AUTOPUSH_ENDPOINT_URL).post("/v1beta2/models/#{model || DEFAULTS[:chat_completion_model_name]}:generateMessage") do |req|
+      response = connection.post("/v1beta2/models/#{model || DEFAULTS[:chat_completion_model_name]}:generateMessage") do |req|
         req.params = {key: api_key}
 
         req.body = {prompt: { messages: [{content: prompt}] }}
         req.body[:model] = model if model
         req.body[:context] = context if context
         req.body[:examples] = examples if examples
         req.body[:messages] = messages if messages
-        req.body[:temperature] = temperature if temperature
+        req.body[:temperature] = temperature || DEFAULTS[:temperature]
         req.body[:candidate_count] = candidate_count if candidate_count
         req.body[:top_p] = top_p if top_p
         req.body[:top_k] = top_k if top_k
@@ -133,7 +132,7 @@ def embed(
       model: nil,
       client: nil
     )
-      response = connection(url: ENDPOINT_URL).post("/v1beta2/models/#{model || DEFAULTS[:embeddings_model_name]}:embedText") do |req|
+      response = connection.post("/v1beta2/models/#{model || DEFAULTS[:embeddings_model_name]}:embedText") do |req|
         req.params = {key: api_key}
 
         req.body = {text: text}
@@ -143,11 +142,57 @@ def embed(
       response.body
     end
 
+    #
+    # Lists models available through the API.
+    # 
+    # @param [Integer] page_size
+    # @param [String] page_token
+    # @return [Hash]
+    # 
+    def list_models(page_size: nil, page_token: nil)
+      response = connection.get("/v1beta2/models") do |req|
+        req.params = {key: api_key}
+
+        req.params[:pageSize] = page_size if page_size
+        req.params[:pageToken] = page_token if page_token
+      end
+      response.body      
+    end
+
+    # 
+    # Runs a model's tokenizer on a string and returns the token count.
+    # 
+    # @param [String] model
+    # @param [String] prompt
+    # @return [Hash]
+    # 
+    def count_message_tokens(model:, prompt:)
+      response = connection.post("/v1beta2/models/#{model}:countMessageTokens") do |req|
+        req.params = {key: api_key}
+
+        req.body = {prompt: { messages: [{content: prompt}] }}
+      end
+      response.body
+    end
+
+    #
+    # Gets information about a specific Model.
+    # 
+    # @param [String] name
+    # @return [Hash]
+    # 
+    def get_model(model:)
+      response = connection.get("/v1beta2/models/#{model}") do |req|
+        req.params = {key: api_key}
+      end
+      response.body
+    end 
+
     private
 
     # standard:disable Lint/DuplicateMethods
-    def connection(url:)
-      Faraday.new(url: url) do |faraday|
+    def connection
+      Faraday.new(url: ENDPOINT_URL) do |faraday|
         faraday.request :json
         faraday.response :json, content_type: /\bjson$/
         faraday.adapter Faraday.default_adapter

spec/fixtures/generate_chat_message.json

@@ -0,0 +1,14 @@
+{
+  "candidates": [
+    {
+      "author": "1",
+      "content": "Hello! How can I help you today?"
+    }
+  ],
+  "messages": [
+    {
+      "author": "0",
+      "content": "Hello!"
+    }
+  ]
+}

spec/fixtures/generate_text.json

@@ -0,0 +1,33 @@
+{
+  "candidates": [
+    {
+      "output": "A man walks into a library and asks for books about paranoia. The librarian whispers, \"They're right behind you!\"",
+      "safetyRatings": [
+        {
+          "category": "HARM_CATEGORY_DEROGATORY",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_TOXICITY",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_VIOLENCE",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_SEXUAL",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_MEDICAL",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_DANGEROUS",
+          "probability": "NEGLIGIBLE"
+        }
+      ]
+    }
+  ]
+}

spec/fixtures/list_models.json

@@ -0,0 +1,37 @@
+{
+  "models": [
+    {
+      "name": "models/chat-bison-001",
+      "version": "001",
+      "displayName": "Chat Bison",
+      "description": "Chat-optimized generative language model.",
+      "inputTokenLimit": 4096,
+      "outputTokenLimit": 1024,
+      "supportedGenerationMethods": ["generateMessage"],
+      "temperature": 0.25,
+      "topP": 0.95,
+      "topK": 40
+    },
+    {
+      "name": "models/text-bison-001",
+      "version": "001",
+      "displayName": "Text Bison",
+      "description": "Model targeted for text generation.",
+      "inputTokenLimit": 8196,
+      "outputTokenLimit": 1024,
+      "supportedGenerationMethods": ["generateText"],
+      "temperature": 0.7,
+      "topP": 0.95,
+      "topK": 40
+    },
+    {
+      "name": "models/embedding-gecko-001",
+      "version": "001",
+      "displayName": "Embedding Gecko",
+      "description": "Obtain a distributed representation of a text.",
+      "inputTokenLimit": 1024,
+      "outputTokenLimit": 1,
+      "supportedGenerationMethods": ["embedText"]
+    }
+  ]
+}

spec/google_palm_api/client_spec.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+require "cohere"
+
+RSpec.describe GooglePalmApi::Client do
+  subject { described_class.new(api_key: "123") }
+
+  describe "#list_models" do
+    let(:fixture) { JSON.parse(File.read("spec/fixtures/list_models.json")) }
+    let(:response) { OpenStruct.new(body: fixture) }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:get)
+        .with("/v1beta2/models")
+        .and_return(response)
+    end
+
+    it "returns a list of models" do
+      expect(subject.list_models.dig("models").count).to eq(3)
+    end
+  end
+
+  describe "#get_model" do
+    let(:fixture) { JSON.parse(File.read("spec/fixtures/list_models.json")).dig("models").first }
+    let(:response) { OpenStruct.new(body: fixture) }
+    let(:model) { "chat-bison-001" }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:get)
+        .with("/v1beta2/models/#{model}")
+        .and_return(response)
+    end
+
+    it "returns the model" do
+      expect(subject.get_model(model: model)).to eq(fixture)
+    end
+  end
+
+  describe "#count_message_tokens" do
+    let(:response) { OpenStruct.new(body: {"tokenCount" => 14}) }
+    let(:model) { "chat-bison-001" }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:post)
+        .with("/v1beta2/models/#{model}:countMessageTokens")
+        .and_return(response)
+    end
+
+    it "returns the token count" do
+      expect(subject.count_message_tokens(model: model, prompt: "Hello")).to eq(response.body)
+    end
+  end
+
+  describe "#embed" do
+    let(:response) { OpenStruct.new(body: {"embedding" => {"value" => [0.0071609155, 0.010057832, -0.016587045]}}) }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:post)
+        .with("/v1beta2/models/embedding-gecko-001:embedText")
+        .and_return(response)
+    end
+
+    it "returns the embedding" do
+      expect(subject.embed(text: "Hello world!")).to eq(response.body)
+    end
+  end
+
+  describe "#generate_text" do
+    let(:fixture) { JSON.parse(File.read("spec/fixtures/generate_text.json")) }
+    let(:response) { OpenStruct.new(body: fixture) }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:post)
+      .with("/v1beta2/models/text-bison-001:generateText")
+      .and_return(response)
+    end
+
+    it "returns the generated text" do
+      expect(subject.generate_text(prompt: "Hello")).to eq(fixture)
+    end
+  end
+
+  describe "#generate_chat_message" do
+    let(:fixture) { JSON.parse(File.read("spec/fixtures/generate_chat_message.json")) }
+    let(:response) { OpenStruct.new(body: fixture) }
+
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:post)
+      .with("/v1beta2/models/chat-bison-001:generateMessage")
+      .and_return(response)
+    end
+
+    it "returns the generated text" do
+      expect(subject.generate_chat_message(prompt: "Hello!")).to eq(fixture)
+    end    
+  end
+end

spec/google_palm_api_spec.rb

@@ -4,8 +4,4 @@
   it "has a version number" do
     expect(GooglePalmApi::VERSION).not_to be nil
   end
-
-  it "does something useful" do
-    expect(false).to eq(true)
-  end
 end

spec/spec_helper.rb

@@ -1,6 +1,10 @@
 # frozen_string_literal: true
 
 require "google_palm_api"
+require "faraday"
+require "faraday_middleware"
+require "ostruct"
+require "json"
 
 RSpec.configure do |config|
   # Enable flags like --only-failures and --next-failure