Langchain::Assistant support for Google Gemini LLM

.env.example

@@ -12,6 +12,7 @@ GOOGLE_PALM_API_KEY=
 GOOGLE_VERTEX_AI_PROJECT_ID=
 # Automagical name which is picked up by Google Cloud Ruby auth module. If set, takes auth token from that key file.
 GOOGLE_CLOUD_CREDENTIALS=
+GOOGLE_GEMINI_API_KEY=
 HUGGING_FACE_API_KEY=
 LLAMACPP_MODEL_PATH=
 LLAMACPP_N_THREADS=

Gemfile.lock

@@ -167,9 +167,9 @@ GEM
       faraday (~> 2.0)
       typhoeus (~> 1.4)
     ffi (1.16.3)
-    google-apis-aiplatform_v1 (0.13.0)
-      google-apis-core (>= 0.12.0, < 2.a)
-    google-apis-core (0.13.0)
+    google-apis-aiplatform_v1 (0.15.0)
+      google-apis-core (>= 0.14.0, < 2.a)
+    google-apis-core (0.14.0)
       addressable (~> 2.5, >= 2.5.1)
       googleauth (~> 1.9)
       httpclient (>= 2.8.1, < 3.a)
@@ -182,7 +182,7 @@ GEM
     google_palm_api (0.1.3)
       faraday (>= 2.0.1, < 3.0)
     google_search_results (2.0.1)
-    googleauth (1.10.0)
+    googleauth (1.11.0)
       faraday (>= 1.0, < 3.a)
       google-cloud-env (~> 2.1)
       jwt (>= 1.4, < 3.0)
@@ -217,7 +217,8 @@ GEM
     json (2.7.1)
     json-schema (4.0.0)
       addressable (>= 2.8)
-    jwt (2.7.1)
+    jwt (2.8.1)
+      base64
     language_server-protocol (3.17.0.3)
     lint_roller (1.1.0)
     llama_cpp (0.9.5)
@@ -486,7 +487,7 @@ DEPENDENCIES
   epsilla-ruby (~> 0.0.4)
   eqn (~> 1.6.5)
   faraday
-  google-apis-aiplatform_v1 (~> 0.7)
+  google-apis-aiplatform_v1 (~> 0.8)
   google_palm_api (~> 0.1.3)
   google_search_results (~> 2.0.0)
   hnswlib (~> 0.8.1)

README.md

@@ -445,14 +445,14 @@ assistant = Langchain::Assistant.new(
   thread: thread,
   instructions: "You are a Meteorologist Assistant that is able to pull the weather for any location",
   tools: [
-    Langchain::Tool::GoogleSearch.new(api_key: ENV["SERPAPI_API_KEY"])
+    Langchain::Tool::Weather.new(api_key: ENV["OPEN_WEATHER_API_KEY"])
   ]
 )
 ```
 ### Using an Assistant
 You can now add your message to an Assistant.
 ```ruby
-assistant.add_message content: "What's the weather in New York City?"
+assistant.add_message content: "What's the weather in New York, New York?"
 ```
 
 Run the Assistant to generate a response. 

langchain.gemspec

@@ -55,7 +55,7 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "elasticsearch", "~> 8.2.0"
   spec.add_development_dependency "epsilla-ruby", "~> 0.0.4"
   spec.add_development_dependency "eqn", "~> 1.6.5"
-  spec.add_development_dependency "google-apis-aiplatform_v1", "~> 0.7"
+  spec.add_development_dependency "google-apis-aiplatform_v1", "~> 0.8"
   spec.add_development_dependency "google_palm_api", "~> 0.1.3"
   spec.add_development_dependency "google_search_results", "~> 2.0.0"
   spec.add_development_dependency "hnswlib", "~> 0.8.1"

lib/langchain.rb

@@ -21,6 +21,7 @@
   "openai" => "OpenAI",
   "openai_validator" => "OpenAIValidator",
   "openai_response" => "OpenAIResponse",
+  "openai_message" => "OpenAIMessage",
   "pdf" => "PDF"
 )
 loader.collapse("#{__dir__}/langchain/llm/response")

lib/langchain/assistants/assistant.rb

@@ -19,7 +19,7 @@ def initialize(
       tools: [],
       instructions: nil
     )
-      raise ArgumentError, "Invalid LLM; currently only Langchain::LLM::OpenAI is supported" unless llm.instance_of?(Langchain::LLM::OpenAI)
+      raise ArgumentError, "Invalid LLM; currently only Langchain::LLM::OpenAI and Langchain::LLM::GoogleGemini are supported" unless [Langchain::LLM::OpenAI, Langchain::LLM::GoogleGemini].include?(llm.class)
       raise ArgumentError, "Thread must be an instance of Langchain::Thread" unless thread.is_a?(Langchain::Thread)
       raise ArgumentError, "Tools must be an array of Langchain::Tool::Base instance(s)" unless tools.is_a?(Array) && tools.all? { |tool| tool.is_a?(Langchain::Tool::Base) }
 
@@ -30,7 +30,10 @@ def initialize(
 
       # The first message in the thread should be the system instructions
       # TODO: What if the user added old messages and the system instructions are already in there? Should this overwrite the existing instructions?
-      add_message(role: "system", content: instructions) if instructions
+      if llm.is_a?(Langchain::LLM::OpenAI)
+        add_message(role: "system", content: instructions) if instructions
+      end
+      # For Google Gemini, system instructions are added to the `system:` param in the `chat` method
     end
 
     # Add a user message to the thread
@@ -59,11 +62,12 @@ def run(auto_tool_execution: false)
 
       while running
         # TODO: I think we need to look at all messages and not just the last one.
-        case (last_message = thread.messages.last).role
-        when "system"
+        last_message = thread.messages.last
+
+        if last_message.system?
           # Do nothing
           running = false
-        when "assistant"
+        elsif last_message.llm?
           if last_message.tool_calls.any?
             if auto_tool_execution
               run_tools(last_message.tool_calls)
@@ -76,11 +80,11 @@ def run(auto_tool_execution: false)
             # Do nothing
             running = false
           end
-        when "user"
+        elsif last_message.user?
           # Run it!
           response = chat_with_llm
 
-          if response.tool_calls
+          if response.tool_calls.any?
             # Re-run the while(running) loop to process the tool calls
             running = true
             add_message(role: response.role, tool_calls: response.tool_calls)
@@ -89,12 +93,12 @@ def run(auto_tool_execution: false)
             running = false
             add_message(role: response.role, content: response.chat_completion)
           end
-        when "tool"
+        elsif last_message.tool?
           # Run it!
           response = chat_with_llm
           running = true
 
-          if response.tool_calls
+          if response.tool_calls.any?
             add_message(role: response.role, tool_calls: response.tool_calls)
           elsif response.chat_completion
             add_message(role: response.role, content: response.chat_completion)
@@ -121,8 +125,14 @@ def add_message_and_run(content:, auto_tool_execution: false)
     # @param output [String] The output of the tool
     # @return [Array<Langchain::Message>] The messages in the thread
     def submit_tool_output(tool_call_id:, output:)
+      tool_role = if llm.is_a?(Langchain::LLM::OpenAI)
+        Langchain::Messages::OpenAIMessage::TOOL_ROLE
+      elsif llm.is_a?(Langchain::LLM::GoogleGemini)
+        Langchain::Messages::GoogleGeminiMessage::TOOL_ROLE
+      end
+
       # TODO: Validate that `tool_call_id` is valid
-      add_message(role: "tool", content: output, tool_call_id: tool_call_id)
+      add_message(role: tool_role, content: output, tool_call_id: tool_call_id)
     end
 
     # Delete all messages in the thread
@@ -156,10 +166,15 @@ def instructions=(new_instructions)
     def chat_with_llm
       Langchain.logger.info("Sending a call to #{llm.class}", for: self.class)
 
-      params = {messages: thread.openai_messages}
+      params = {messages: thread.array_of_message_hashes}
 
       if tools.any?
-        params[:tools] = tools.map(&:to_openai_tools).flatten
+        if llm.is_a?(Langchain::LLM::OpenAI)
+          params[:tools] = tools.map(&:to_openai_tools).flatten
+        elsif llm.is_a?(Langchain::LLM::GoogleGemini)
+          params[:tools] = tools.map(&:to_google_gemini_tools).flatten
+          params[:system] = instructions if instructions
+        end
         # TODO: Not sure that tool_choice should always be "auto"; Maybe we can let the user toggle it.
         params[:tool_choice] = "auto"
       end
@@ -173,11 +188,11 @@ def chat_with_llm
     def run_tools(tool_calls)
       # Iterate over each function invocation and submit tool output
       tool_calls.each do |tool_call|
-        tool_call_id = tool_call.dig("id")
-
-        function_name = tool_call.dig("function", "name")
-        tool_name, method_name = function_name.split("-")
-        tool_arguments = JSON.parse(tool_call.dig("function", "arguments"), symbolize_names: true)
+        tool_call_id, tool_name, method_name, tool_arguments = if llm.is_a?(Langchain::LLM::OpenAI)
+          extract_openai_tool_call(tool_call: tool_call)
+        elsif llm.is_a?(Langchain::LLM::GoogleGemini)
+          extract_google_gemini_tool_call(tool_call: tool_call)
+        end
 
         tool_instance = tools.find do |t|
           t.name == tool_name
@@ -190,13 +205,41 @@ def run_tools(tool_calls)
 
       response = chat_with_llm
 
-      if response.tool_calls
+      if response.tool_calls.any?
         add_message(role: response.role, tool_calls: response.tool_calls)
       elsif response.chat_completion
         add_message(role: response.role, content: response.chat_completion)
       end
     end
 
+    # Extract the tool call information from the OpenAI tool call hash
+    #
+    # @param tool_call [Hash] The tool call hash
+    # @return [Array] The tool call information
+    def extract_openai_tool_call(tool_call:)
+      tool_call_id = tool_call.dig("id")
+
+      function_name = tool_call.dig("function", "name")
+      tool_name, method_name = function_name.split("__")
+      tool_arguments = JSON.parse(tool_call.dig("function", "arguments"), symbolize_names: true)
+
+      [tool_call_id, tool_name, method_name, tool_arguments]
+    end
+
+    # Extract the tool call information from the Google Gemini tool call hash
+    #
+    # @param tool_call [Hash] The tool call hash, format: {"functionCall"=>{"name"=>"weather__execute", "args"=>{"input"=>"NYC"}}}
+    # @return [Array] The tool call information
+    def extract_google_gemini_tool_call(tool_call:)
+      tool_call_id = tool_call.dig("functionCall", "name")
+
+      function_name = tool_call.dig("functionCall", "name")
+      tool_name, method_name = function_name.split("__")
+      tool_arguments = tool_call.dig("functionCall", "args").transform_keys(&:to_sym)
+
+      [tool_call_id, tool_name, method_name, tool_arguments]
+    end
+
     # Build a message
     #
     # @param role [String] The role of the message
@@ -205,7 +248,11 @@ def run_tools(tool_calls)
     # @param tool_call_id [String] The ID of the tool call to include in the message
     # @return [Langchain::Message] The Message object
     def build_message(role:, content: nil, tool_calls: [], tool_call_id: nil)
-      Message.new(role: role, content: content, tool_calls: tool_calls, tool_call_id: tool_call_id)
+      if llm.is_a?(Langchain::LLM::OpenAI)
+        Langchain::Messages::OpenAIMessage.new(role: role, content: content, tool_calls: tool_calls, tool_call_id: tool_call_id)
+      elsif llm.is_a?(Langchain::LLM::GoogleGemini)
+        Langchain::Messages::GoogleGeminiMessage.new(role: role, content: content, tool_calls: tool_calls, tool_call_id: tool_call_id)
+      end
     end
 
     # TODO: Fix the message truncation when context window is exceeded

lib/langchain/assistants/messages/base.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module Langchain
+  module Messages
+    class Base
+      attr_reader :role, :content, :tool_calls, :tool_call_id
+
+      # Check if the message came from a user
+      #
+      # @param [Boolean] true/false whether the message came from a user
+      def user?
+        role == "user"
+      end
+    end
+  end
+end