Class: Rails::Nl2sql::Providers::GeminiProvider

Inherits:
Base
  • Object
show all
Defined in:
lib/rails/nl2sql/providers/gemini_provider.rb

Instance Method Summary collapse

Constructor Details

#initialize(api_key:, model: 'gemini-pro') ⇒ GeminiProvider

Returns a new instance of GeminiProvider.



7
8
9
10
11
12
13
14
15
 
# File 'lib/rails/nl2sql/providers/gemini_provider.rb', line 7

def initialize(api_key:, model: 'gemini-pro')
  # Authenticate with Google Cloud
  # You can use a service account or your user credentials.
  # For more information, see: https://cloud.google.com/docs/authentication
  Google::Auth.new.apply(scope: 'https://www.googleapis.com/auth/cloud-platform')

  @client = Google::Apis::AiplatformV1::AIPlatformService.new
  @model = model
end

Instance Method Details

#complete(prompt:, **params) ⇒ Object 



17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
 
# File 'lib/rails/nl2sql/providers/gemini_provider.rb', line 17

def complete(prompt:, **params)
  # The Gemini API uses a different request format than the other providers.
  # We need to convert the prompt into a format that the Gemini API can understand.
  request = Google::Apis::AiplatformV1::GoogleCloudAiplatformV1PredictRequest.new(
    instances: [
      {
        prompt: prompt
      }
    ],
    parameters: {
      temperature: params[:temperature] || 0.2,
      maxOutputTokens: params[:max_tokens] || 256,
      topP: params[:top_p] || 0.95,
      topK: params[:top_k] || 40
    }
  )

  # The endpoint for the Gemini API is different for each region.
  # We need to get the endpoint for the region that the user is in.
  # For more information, see: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models
  endpoint = 'us-central1-aiplatform.googleapis.com' # Or another regional endpoint

  response = @client.predict_project_location_publisher_model(endpoint, request)

  # The Gemini API returns a different response format than the other providers.
  # We need to convert the response into a format that the other providers can understand.
  response.predictions.first['content']
end