gemini_ai.rb in boxcars-0.6.9

- old
+ new

@@ -1,90 +1,80 @@
 # frozen_string_literal: true
 
-# Boxcars - a framework for running a series of tools to get an answer to a question.
+# Boxcars is a framework for running a series of tools to get an answer to a question.
 module Boxcars
-  # A engine that uses Gemini's API.
+  # A engine that uses GeminiAI's API.
   class GeminiAi < Engine
-    attr_reader :prompts, :llm_params, :model_kwargs, :batch_size
+    attr_reader :prompts, :llm_parmas, :model_kwargs, :batch_size
 
     # The default parameters to use when asking the engine.
     DEFAULT_PARAMS = {
-      model: "gemini-1.5-flash-latest"
+      model: "gemini-1.5-flash-latest",
+      temperature: 0.1
     }.freeze
 
     # the default name of the engine
-    DEFAULT_NAME = "Google Gemini AI engine"
+    DEFAULT_NAME = "GeminiAI engine"
     # the default description of the engine
-    DEFAULT_DESCRIPTION = "useful for when you need to use Google Gemini AI to answer questions. " \
+    DEFAULT_DESCRIPTION = "useful for when you need to use AI to answer questions. " \
                           "You should ask targeted questions"
 
-    # A engine is the driver for a single tool to run.
-    # @param name [String] The name of the engine. Defaults to "OpenAI engine".
+    # A engine is a container for a single tool to run.
+    # @param name [String] The name of the engine. Defaults to "GeminiAI engine".
     # @param description [String] A description of the engine. Defaults to:
     #        useful for when you need to use AI to answer questions. You should ask targeted questions".
     # @param prompts [Array<String>] The prompts to use when asking the engine. Defaults to [].
-    def initialize(name: DEFAULT_NAME, description: DEFAULT_DESCRIPTION, prompts: [], **kwargs)
-      @llm_params = DEFAULT_PARAMS.merge(kwargs)
+    # @param batch_size [Integer] The number of prompts to send to the engine at once. Defaults to 20.
+    def initialize(name: DEFAULT_NAME, description: DEFAULT_DESCRIPTION, prompts: [], batch_size: 20, **kwargs)
+      @llm_parmas = DEFAULT_PARAMS.merge(kwargs)
       @prompts = prompts
-      @batch_size = 20
+      @batch_size = batch_size
       super(description: description, name: name)
     end
 
+    # Get the OpenAI API client
+    # @param gemini_api_key [String] The access token to use when asking the engine.
+    #   Defaults to Boxcars.configuration.gemini_api_key
+    # @return [OpenAI::Client] The OpenAI API gem client.
+    def self.open_ai_client(gemini_api_key: nil)
+      access_token = Boxcars.configuration.gemini_api_key(gemini_api_key: gemini_api_key)
+      ::OpenAI::Client.new(access_token: access_token, uri_base: "https://generativelanguage.googleapis.com/v1beta/openai/")
+    end
+
     def conversation_model?(_model)
       true
     end
 
-    def chat(params, gemini_api_key)
-      raise Boxcars::ConfigurationError('Google AI API key not set') if gemini_api_key.blank?
-
-      model_string = params.delete(:model_string)
-      raise Boxcars::ConfigurationError('Google AI API key not set') if model_string.blank?
-
-      # Define the API endpoint and parameters
-      api_endpoint = "https://generativelanguage.googleapis.com/v1beta/models/#{model_string}:generateContent?key=#{gemini_api_key}"
-
-      connection = Faraday.new(api_endpoint) do |faraday|
-        faraday.request :url_encoded
-        faraday.headers['Content-Type'] = 'application/json'
-      end
-
-      # Make the API call
-      response = connection.post { |req| req.body = params.to_json }
-
-      JSON.parse(response.body, symbolize_names: true)
-    end
-
     # Get an answer from the engine.
     # @param prompt [String] The prompt to use when asking the engine.
-    # @param gemini_api_key [String] Optional api key to use when asking the engine.
+    # @param gemini_api_key [String] The access token to use when asking the engine.
     #   Defaults to Boxcars.configuration.gemini_api_key.
     # @param kwargs [Hash] Additional parameters to pass to the engine if wanted.
-    def client(prompt:, inputs: {}, **kwargs)
-      api_key = Boxcars.configuration.gemini_api_key(**kwargs)
-      option_params = llm_params.merge(kwargs)
-      model_string = option_params.delete(:model) || DEFAULT_PARAMS[:model]
-      convo = prompt.as_messages(inputs: inputs)
-      # Convert conversation to Google Gemini format
-      params = to_google_gemini_format(convo[:messages], option_params)
-      params[:model_string] = model_string
-      Boxcars.debug("Prompt after formatting:#{params[:message]}", :cyan) if Boxcars.configuration.log_prompts
-      chat(params, api_key)
+    def client(prompt:, inputs: {}, gemini_api_key: nil, **kwargs)
+      clnt = GeminiAi.open_ai_client(gemini_api_key: gemini_api_key)
+      params = llm_parmas.merge(kwargs)
+      prompt = prompt.first if prompt.is_a?(Array)
+      params = prompt.as_messages(inputs).merge(params)
+      if Boxcars.configuration.log_prompts
+        Boxcars.debug(params[:messages].last(2).map { |p| ">>>>>> Role: #{p[:role]} <<<<<<\n#{p[:content]}" }.join("\n"), :cyan)
+      end
+      clnt.chat(parameters: params)
+    rescue => e
+      Boxcars.error(e, :red)
+      raise
     end
 
     # get an answer from the engine for a question.
     # @param question [String] The question to ask the engine.
     # @param kwargs [Hash] Additional parameters to pass to the engine if wanted.
     def run(question, **kwargs)
       prompt = Prompt.new(template: question)
       response = client(prompt: prompt, **kwargs)
-
       raise Error, "GeminiAI: No response from API" unless response
-      raise Error, "GeminiAI: #{response[:error]}" if response[:error]
 
-      answer = response[:candidates].first[:content][:parts].first[:text]
-      Boxcars.debug(response, :yellow)
-      answer
+      check_response(response)
+      response["choices"].map { |c| c.dig("message", "content") || c["text"] }.join("\n").strip
     end
 
     # Get the default parameters for the engine.
     def default_params
       llm_params
@@ -93,68 +83,36 @@
     # make sure we got a valid response
     # @param response [Hash] The response to check.
     # @param must_haves [Array<String>] The keys that must be in the response. Defaults to %w[choices].
     # @raise [KeyError] if there is an issue with the access token.
     # @raise [ValueError] if the response is not valid.
-    def check_response(response, must_haves: %w[completion])
-      if response['error']
+    def check_response(response, must_haves: %w[choices])
+      if response['error'].is_a?(Hash)
         code = response.dig('error', 'code')
         msg = response.dig('error', 'message') || 'unknown error'
-        raise KeyError, "ANTHOPIC_API_KEY not valid" if code == 'invalid_api_key'
+        raise KeyError, "GEMINI_API_TOKEN not valid" if code == 'invalid_api_key'
 
-        raise ValueError, "Gemini error: #{msg}"
+        raise ValueError, "GeminiAI error: #{msg}"
       end
 
       must_haves.each do |key|
         raise ValueError, "Expecting key #{key} in response" unless response.key?(key)
       end
     end
 
     # the engine type
     def engine_type
-      "claude"
+      "gemini_ai"
     end
 
-    # lookup the context size for a model by name
-    # @param modelname [String] The name of the model to lookup.
-    def modelname_to_contextsize(_modelname)
-      100000
-    end
-
     # Calculate the maximum number of tokens possible to generate for a prompt.
     # @param prompt_text [String] The prompt text to use.
     # @return [Integer] the number of tokens possible to generate.
     def max_tokens_for_prompt(prompt_text)
       num_tokens = get_num_tokens(prompt_text)
 
       # get max context size for model by name
-      max_size = modelname_to_contextsize(model_name)
+      max_size = 8096
       max_size - num_tokens
-    end
-
-    def to_google_gemini_format(convo, option_params)
-      instructions = convo.shift.last if convo.first && convo.first[:role] == :system
-      system_instructions = instructions || "You are a helpful assistant."
-
-      # Convert conversation history to the format expected by Google
-      contents = convo.map { |message| { text: message[:content] } }
-
-      generation_config = {}
-      if option_params.length.positive?
-        generation_config.merge!(option_params)
-        generation_config[:stopSequences] = [generation_config.delete(:stop)] if generation_config[:stop].present?
-      end
-
-      rv = {
-        system_instruction: { parts: { text: system_instructions } }, # System instructions or context
-        contents: { parts: contents } # The chat messages
-      }
-
-      rv[:generationConfig] = generation_config if generation_config.length.positive?
-      rv
-    end
-
-    def default_prefixes
-      { system: "SYSTEM: ", user: "USER: ", assistant: "CHATBOT: ", history: :history }
     end
   end
 end