Merge pull request #1 from andreibondarev/fixes

fixes

Author: andreibondarev

Gemfile.lock

@@ -100,6 +100,7 @@ GEM
 
 PLATFORMS
   x86_64-darwin-19
+  x86_64-linux
 
 DEPENDENCIES
   google_palm_api!

README.md

@@ -6,20 +6,37 @@ Welcome to your new gem! In this directory, you'll find the files you need to be
 
 ## Installation
 
-TODO: Replace `UPDATE_WITH_YOUR_GEM_NAME_PRIOR_TO_RELEASE_TO_RUBYGEMS_ORG` with your gem name right after releasing it to RubyGems.org. Please do not do it earlier due to security reasons. Alternatively, replace this section with instructions to install your gem from git if you don't plan to release to RubyGems.org.
-
 Install the gem and add to the application's Gemfile by executing:
 
-    $ bundle add UPDATE_WITH_YOUR_GEM_NAME_PRIOR_TO_RELEASE_TO_RUBYGEMS_ORG
+    $ bundle add google_palm_api
 
 If bundler is not being used to manage dependencies, install the gem by executing:
 
-    $ gem install UPDATE_WITH_YOUR_GEM_NAME_PRIOR_TO_RELEASE_TO_RUBYGEMS_ORG
+    $ gem install google_palm_api
 
 ## Usage
 
-TODO: Write usage instructions here
-
+```ruby
+require 'google_palm_api'
+```
+```ruby
+client.generate_text(prompt:)
+```
+```ruby
+client.generate_chat_message(prompt:)
+```
+```ruby
+client.embed(text:)
+```
+```ruby
+client.get_model(model:)
+```
+```ruby
+client.list_models()
+```
+```ruby
+client.count_message_tokens(message:)
+```
 ## Development
 
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.

lib/google_palm_api.rb

@@ -4,6 +4,6 @@
 
 module GooglePalmApi
   class Error < StandardError; end
-  
+
   autoload :Client, "google_palm_api/client"
 end

lib/google_palm_api/client.rb

@@ -13,7 +13,7 @@ class Client
       temperature: 0.0,
       completion_model_name: "text-bison-001",
       chat_completion_model_name: "chat-bison-001",
-      embeddings_model_name: "embedding-gecko-001",
+      embeddings_model_name: "embedding-gecko-001"
     }
 
     def initialize(api_key:)
@@ -40,7 +40,6 @@ def initialize(api_key:)
     #
     def generate_text(
       prompt:,
-      model: nil,
       temperature: nil,
       candidate_count: nil,
       max_output_tokens: nil,
@@ -50,10 +49,10 @@ def generate_text(
       stop_sequences: nil,
       client: nil
     )
-      response = connection.post("/v1beta2/models/#{model || DEFAULTS[:completion_model_name]}:generateText") do |req|
+      response = connection.post("/v1beta2/models/#{DEFAULTS[:completion_model_name]}:generateText") do |req|
         req.params = {key: api_key}
 
-        req.body = {prompt: { text: prompt }}
+        req.body = {prompt: {text: prompt}}
         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
@@ -87,7 +86,6 @@ def generate_text(
     #
     def generate_chat_message(
       prompt:,
-      model: nil,
       context: nil,
       examples: nil,
       messages: nil,
@@ -98,11 +96,10 @@ def generate_chat_message(
       client: nil
     )
       # Overwrite the default ENDPOINT_URL for this method.
-      response = connection.post("/v1beta2/models/#{model || DEFAULTS[:chat_completion_model_name]}:generateMessage") do |req|
+      response = connection.post("/v1beta2/models/#{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 = {prompt: {messages: [{content: prompt}]}}
         req.body[:context] = context if context
         req.body[:examples] = examples if examples
         req.body[:messages] = messages if messages
@@ -113,7 +110,7 @@ def generate_chat_message(
         req.body[:client] = client if client
       end
       response.body
-    end    
+    end
 
     #
     # The embedding service in the PaLM API generates state-of-the-art embeddings for words, phrases, and sentences.
@@ -144,49 +141,49 @@ def embed(
 
     #
     # 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      
+      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}] }}
+        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 
+    end
 
     private
 

spec/google_palm_api/client_spec.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "cohere"
-
 RSpec.describe GooglePalmApi::Client do
   subject { described_class.new(api_key: "123") }
 
@@ -71,8 +69,8 @@
 
     before do
       allow_any_instance_of(Faraday::Connection).to receive(:post)
-      .with("/v1beta2/models/text-bison-001:generateText")
-      .and_return(response)
+        .with("/v1beta2/models/text-bison-001:generateText")
+        .and_return(response)
     end
 
     it "returns the generated text" do
@@ -84,15 +82,14 @@
     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)
+        .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
 end