langchainrb 0.5.5 → 0.5.7

data/lib/langchain/llm/google_palm.rb

@@ -22,14 +22,19 @@ module Langchain::LLM
 
     DEFAULTS = {
       temperature: 0.0,
-      dimension: 768 # This is what the `embedding-gecko-001` model generates
+      dimension: 768, # This is what the `embedding-gecko-001` model generates
+      completion_model_name: "text-bison-001",
+      chat_completion_model_name: "chat-bison-001",
+      embeddings_model_name: "embedding-gecko-001"
     }.freeze
+    LENGTH_VALIDATOR = Langchain::Utils::TokenLength::GooglePalmValidator
 
-    def initialize(api_key:)
+    def initialize(api_key:, default_options: {})
       depends_on "google_palm_api"
       require "google_palm_api"
 
       @client = ::GooglePalmApi::Client.new(api_key: api_key)
+      @defaults = DEFAULTS.merge(default_options)
     end
 
     #
@@ -55,7 +60,8 @@ module Langchain::LLM
     def complete(prompt:, **params)
       default_params = {
         prompt: prompt,
-        temperature: DEFAULTS[:temperature]
+        temperature: @defaults[:temperature],
+        completion_model_name: @defaults[:completion_model_name]
       }
 
       if params[:stop_sequences]
@@ -84,13 +90,15 @@ module Langchain::LLM
       raise ArgumentError.new(":prompt or :messages argument is expected") if prompt.empty? && messages.empty?
 
       default_params = {
-        temperature: DEFAULTS[:temperature],
+        temperature: @defaults[:temperature],
+        chat_completion_model_name: @defaults[:chat_completion_model_name],
         context: context,
         messages: compose_chat_messages(prompt: prompt, messages: messages),
         examples: compose_examples(examples)
       }
 
-      Langchain::Utils::TokenLength::GooglePalmValidator.validate_max_tokens!(self, default_params[:messages], "chat-bison-001")
+      # chat-bison-001 is the only model that currently supports countMessageTokens functions
+      LENGTH_VALIDATOR.validate_max_tokens!(default_params[:messages], "chat-bison-001", llm: self)
 
       if options[:stop_sequences]
         default_params[:stop] = options.delete(:stop_sequences)
@@ -116,13 +124,13 @@ module Langchain::LLM
     #
     def summarize(text:)
       prompt_template = Langchain::Prompt.load_from_path(
-        file_path: Langchain.root.join("langchain/llm/prompts/summarize_template.json")
+        file_path: Langchain.root.join("langchain/llm/prompts/summarize_template.yaml")
       )
       prompt = prompt_template.format(text: text)
 
       complete(
         prompt: prompt,
-        temperature: DEFAULTS[:temperature],
+        temperature: @defaults[:temperature],
         # Most models have a context length of 2048 tokens (except for the newest models, which support 4096).
         max_tokens: 2048
       )

data/lib/langchain/llm/openai.rb

@@ -17,12 +17,14 @@ module Langchain::LLM
       embeddings_model_name: "text-embedding-ada-002",
       dimension: 1536
     }.freeze
+    LENGTH_VALIDATOR = Langchain::Utils::TokenLength::OpenAIValidator
 
-    def initialize(api_key:, llm_options: {})
+    def initialize(api_key:, llm_options: {}, default_options: {})
       depends_on "ruby-openai"
       require "openai"
 
       @client = ::OpenAI::Client.new(access_token: api_key, **llm_options)
+      @defaults = DEFAULTS.merge(default_options)
     end
 
     #
@@ -33,9 +35,9 @@ module Langchain::LLM
     # @return [Array] The embedding
     #
     def embed(text:, **params)
-      parameters = {model: DEFAULTS[:embeddings_model_name], input: text}
+      parameters = {model: @defaults[:embeddings_model_name], input: text}
 
-      Langchain::Utils::TokenLength::OpenAIValidator.validate_max_tokens!(text, parameters[:model])
+      validate_max_tokens(text, parameters[:model])
 
       response = client.embeddings(parameters: parameters.merge(params))
       response.dig("data").first.dig("embedding")
@@ -49,37 +51,85 @@ module Langchain::LLM
     # @return [String] The completion
     #
     def complete(prompt:, **params)
-      parameters = compose_parameters DEFAULTS[:completion_model_name], params
+      parameters = compose_parameters @defaults[:completion_model_name], params
 
       parameters[:prompt] = prompt
-      parameters[:max_tokens] = Langchain::Utils::TokenLength::OpenAIValidator.validate_max_tokens!(prompt, parameters[:model])
+      parameters[:max_tokens] = validate_max_tokens(prompt, parameters[:model])
 
       response = client.completions(parameters: parameters)
       response.dig("choices", 0, "text")
     end
 
     #
-    # Generate a chat completion for a given prompt
+    # Generate a chat completion for a given prompt or messages.
+    #
+    # == Examples
+    #
+    #     # simplest case, just give a prompt
+    #     openai.chat prompt: "When was Ruby first released?"
+    #
+    #     # prompt plus some context about how to respond
+    #     openai.chat context: "You are RubyGPT, a helpful chat bot for helping people learn Ruby", prompt: "Does Ruby have a REPL like IPython?"
+    #
+    #     # full control over messages that get sent, equivilent to the above
+    #     openai.chat messages: [
+    #       {
+    #         role: "system",
+    #         content: "You are RubyGPT, a helpful chat bot for helping people learn Ruby", prompt: "Does Ruby have a REPL like IPython?"
+    #       },
+    #       {
+    #         role: "user",
+    #         content: "When was Ruby first released?"
+    #       }
+    #     ]
+    #
+    #     # few-short prompting with examples
+    #     openai.chat prompt: "When was factory_bot released?",
+    #       examples: [
+    #         {
+    #           role: "user",
+    #           content: "When was Ruby on Rails released?"
+    #         }
+    #         {
+    #           role: "assistant",
+    #           content: "2004"
+    #         },
+    #       ]
     #
     # @param prompt [String] The prompt to generate a chat completion for
-    # @param messages [Array] The messages that have been sent in the conversation
-    # @param context [String] The context of the conversation
-    # @param examples [Array] Examples of messages provide model with
-    # @param options extra parameters passed to OpenAI::Client#chat
+    # @param messages [Array<Hash>] The messages that have been sent in the conversation
+    #   Each message should be a Hash with the following keys:
+    #   - :content [String] The content of the message
+    #   - :role [String] The role of the sender (system, user, assistant, or function)
+    # @param context [String] An initial context to provide as a system message, ie "You are RubyGPT, a helpful chat bot for helping people learn Ruby"
+    # @param examples [Array<Hash>] Examples of messages to provide to the model. Useful for Few-Shot Prompting
+    #   Each message should be a Hash with the following keys:
+    #   - :content [String] The content of the message
+    #   - :role [String] The role of the sender (system, user, assistant, or function)
+    # @param options <Hash> extra parameters passed to OpenAI::Client#chat
+    # @yield [String] Stream responses back one String at a time
     # @return [String] The chat completion
     #
     def chat(prompt: "", messages: [], context: "", examples: [], **options)
       raise ArgumentError.new(":prompt or :messages argument is expected") if prompt.empty? && messages.empty?
 
-      parameters = compose_parameters DEFAULTS[:chat_completion_model_name], options
+      parameters = compose_parameters @defaults[:chat_completion_model_name], options
       parameters[:messages] = compose_chat_messages(prompt: prompt, messages: messages, context: context, examples: examples)
       parameters[:max_tokens] = validate_max_tokens(parameters[:messages], parameters[:model])
 
+      if (streaming = block_given?)
+        parameters[:stream] = proc do |chunk, _bytesize|
+          yield chunk.dig("choices", 0, "delta", "content")
+        end
+      end
+
       response = client.chat(parameters: parameters)
 
-      raise "Chat completion failed: #{response}" if response.dig("error")
+      raise "Chat completion failed: #{response}" if !response.empty? && response.dig("error")
 
-      response.dig("choices", 0, "message", "content")
+      unless streaming
+        response.dig("choices", 0, "message", "content")
+      end
     end
 
     #
@@ -90,17 +140,17 @@ module Langchain::LLM
     #
     def summarize(text:)
       prompt_template = Langchain::Prompt.load_from_path(
-        file_path: Langchain.root.join("langchain/llm/prompts/summarize_template.json")
+        file_path: Langchain.root.join("langchain/llm/prompts/summarize_template.yaml")
       )
       prompt = prompt_template.format(text: text)
 
-      complete(prompt: prompt, temperature: DEFAULTS[:temperature])
+      complete(prompt: prompt, temperature: @defaults[:temperature])
     end
 
     private
 
     def compose_parameters(model, params)
-      default_params = {model: model, temperature: DEFAULTS[:temperature]}
+      default_params = {model: model, temperature: @defaults[:temperature]}
 
       default_params[:stop] = params.delete(:stop_sequences) if params[:stop_sequences]
 
@@ -140,7 +190,7 @@ module Langchain::LLM
     end
 
     def validate_max_tokens(messages, model)
-      Langchain::Utils::TokenLength::OpenAIValidator.validate_max_tokens!(messages, model)
+      LENGTH_VALIDATOR.validate_max_tokens!(messages, model)
     end
   end
 end

data/lib/langchain/llm/prompts/summarize_template.yaml

@@ -0,0 +1,9 @@
+_type: prompt
+input_variables:
+  - text
+template: |
+  Write a concise summary of the following:
+
+  {text}
+
+  CONCISE SUMMARY:

data/lib/langchain/llm/replicate.rb

@@ -32,7 +32,7 @@ module Langchain::LLM
     #
     # @param api_key [String] The API key to use
     #
-    def initialize(api_key:)
+    def initialize(api_key:, default_options: {})
       depends_on "replicate-ruby"
       require "replicate"
 
@@ -41,6 +41,7 @@ module Langchain::LLM
       end
 
       @client = ::Replicate.client
+      @defaults = DEFAULTS.merge(default_options)
     end
 
     #
@@ -94,13 +95,13 @@ module Langchain::LLM
     #
     def summarize(text:)
       prompt_template = Langchain::Prompt.load_from_path(
-        file_path: Langchain.root.join("langchain/llm/prompts/summarize_template.json")
+        file_path: Langchain.root.join("langchain/llm/prompts/summarize_template.yaml")
       )
       prompt = prompt_template.format(text: text)
 
       complete(
         prompt: prompt,
-        temperature: DEFAULTS[:temperature],
+        temperature: @defaults[:temperature],
         # Most models have a context length of 2048 tokens (except for the newest models, which support 4096).
         max_tokens: 2048
       )
@@ -111,11 +112,11 @@ module Langchain::LLM
     private
 
     def completion_model
-      @completion_model ||= client.retrieve_model(DEFAULTS[:completion_model_name]).latest_version
+      @completion_model ||= client.retrieve_model(@defaults[:completion_model_name]).latest_version
     end
 
     def embeddings_model
-      @embeddings_model ||= client.retrieve_model(DEFAULTS[:embeddings_model_name]).latest_version
+      @embeddings_model ||= client.retrieve_model(@defaults[:embeddings_model_name]).latest_version
     end
   end
 end

data/lib/langchain/prompt/base.rb

@@ -45,11 +45,11 @@ module Langchain::Prompt
     end
 
     #
-    # Save the object to a file in JSON format.
+    # Save the object to a file in JSON or YAML format.
     #
     # @param file_path [String, Pathname] The path to the file to save the object to
     #
-    # @raise [ArgumentError] If file_path doesn't end with .json
+    # @raise [ArgumentError] If file_path doesn't end with .json or .yaml or .yml
     #
     # @return [void]
     #

data/lib/langchain/tool/base.rb

@@ -9,7 +9,7 @@ module Langchain::Tool
   #
   # - {Langchain::Tool::Calculator}: Calculate the result of a math expression
   # - {Langchain::Tool::RubyCodeInterpretor}: Runs ruby code
-  # - {Langchain::Tool::Search}: search on Google (via SerpAPI)
+  # - {Langchain::Tool::GoogleSearch}: search on Google (via SerpAPI)
   # - {Langchain::Tool::Wikipedia}: search on Wikipedia
   #
   # == Usage
@@ -30,13 +30,13 @@ module Langchain::Tool
   #     agent = Langchain::Agent::ChainOfThoughtAgent.new(
   #       llm: :openai, # or :cohere, :hugging_face, :google_palm or :replicate
   #       llm_api_key: ENV["OPENAI_API_KEY"],
-  #       tools: ["search", "calculator", "wikipedia"]
+  #       tools: ["google_search", "calculator", "wikipedia"]
   #     )
   #
   # 4. Confirm that the Agent is using the Tools you passed in:
   #
   #     agent.tools
-  #     # => ["search", "calculator", "wikipedia"]
+  #     # => ["google_search", "calculator", "wikipedia"]
   #
   # == Adding Tools
   #
@@ -57,6 +57,12 @@ module Langchain::Tool
       self.class.const_get(:NAME)
     end
 
+    def self.logger_options
+      {
+        color: :light_blue
+      }
+    end
+
     #
     # Returns the DESCRIPTION constant of the tool
     #

data/lib/langchain/tool/calculator.rb

@@ -16,6 +16,11 @@ module Langchain::Tool
       Useful for getting the result of a math expression.
 
       The input to this tool should be a valid mathematical expression that could be executed by a simple calculator.
+      Usage:
+        Action Input: 1 + 1
+        Action Input: 3 * 2 / 4
+        Action Input: 9 - 7
+        Action Input: (4.1 + 2.3) / (2.0 - 5.6) * 3
     DESC
 
     def initialize
@@ -28,18 +33,11 @@ module Langchain::Tool
     # @param input [String] math expression
     # @return [String] Answer
     def execute(input:)
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Executing \"#{input}\"")
+      Langchain.logger.info("Executing \"#{input}\"", for: self.class)
 
       Eqn::Calculator.calc(input)
     rescue Eqn::ParseError, Eqn::NoVariableValueError
-      # Sometimes the input is not a pure math expression, e.g: "12F in Celsius"
-      # We can use the google answer box to evaluate this expression
-      # TODO: Figure out to find a better way to evaluate these language expressions.
-      hash_results = Langchain::Tool::SerpApi
-        .new(api_key: ENV["SERPAPI_API_KEY"])
-        .execute_search(input: input)
-      hash_results.dig(:answer_box, :to) ||
-        hash_results.dig(:answer_box, :result)
+      "\"#{input}\" is an invalid mathematical expression"
     end
   end
 end

data/lib/langchain/tool/database.rb

@@ -14,15 +14,18 @@ module Langchain::Tool
       The input to this tool should be valid SQL.
     DESC
 
-    attr_reader :db
+    attr_reader :db, :requested_tables, :except_tables
 
     #
     # Establish a database connection
     #
     # @param connection_string [String] Database connection info, e.g. 'postgres://user:password@localhost:5432/db_name'
+    # @param tables [Array<Symbol>] The tables to use. Will use all if empty.
+    # @param except_tables [Array<Symbol>] The tables to exclude. Will exclude none if empty.
+
     # @return [Database] Database object
     #
-    def initialize(connection_string:)
+    def initialize(connection_string:, tables: [], except_tables: [])
       depends_on "sequel"
       require "sequel"
       require "sequel/extensions/schema_dumper"
@@ -30,7 +33,8 @@ module Langchain::Tool
       raise StandardError, "connection_string parameter cannot be blank" if connection_string.empty?
 
       @db = Sequel.connect(connection_string)
-      @db.extension :schema_dumper
+      @requested_tables = tables
+      @except_tables = except_tables
     end
 
     #
@@ -38,9 +42,26 @@ module Langchain::Tool
     #
     # @return [String] schema
     #
-    def schema
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Dumping schema")
-      db.dump_schema_migration(same_db: true, indexes: false) unless db.adapter_scheme == :mock
+    def dump_schema
+      Langchain.logger.info("Dumping schema tables and keys", for: self.class)
+      schema = ""
+      db.tables.each do |table|
+        next if except_tables.include?(table)
+        next unless requested_tables.empty? || requested_tables.include?(table)
+
+        schema << "CREATE TABLE #{table}(\n"
+        db.schema(table).each do |column|
+          schema << "#{column[0]} #{column[1][:type]}"
+          schema << " PRIMARY KEY" if column[1][:primary_key] == true
+          schema << "," unless column == db.schema(table).last
+          schema << "\n"
+        end
+        schema << ");\n"
+        db.foreign_key_list(table).each do |fk|
+          schema << "ALTER TABLE #{table} ADD FOREIGN KEY (#{fk[:columns][0]}) REFERENCES #{fk[:table]}(#{fk[:key][0]});\n"
+        end
+      end
+      schema
     end
 
     #
@@ -50,11 +71,11 @@ module Langchain::Tool
     # @return [Array] results
     #
     def execute(input:)
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Executing \"#{input}\"")
+      Langchain.logger.info("Executing \"#{input}\"", for: self.class)
 
       db[input].to_a
     rescue Sequel::DatabaseError => e
-      Langchain.logger.error("[#{self.class.name}]".light_red + ": #{e.message}")
+      Langchain.logger.error(e.message, for: self.class)
     end
   end
 end

data/lib/langchain/tool/{serp_api.rb → google_search.rb}

@@ -1,18 +1,18 @@
 # frozen_string_literal: true
 
 module Langchain::Tool
-  class SerpApi < Base
+  class GoogleSearch < Base
     #
-    # Wrapper around SerpAPI
+    # Wrapper around Google Serp SPI
     #
     # Gem requirements: gem "google_search_results", "~> 2.0.0"
     #
     # Usage:
-    # search = Langchain::Tool::SerpApi.new(api_key: "YOUR_API_KEY")
+    # search = Langchain::Tool::GoogleSearch.new(api_key: "YOUR_API_KEY")
     # search.execute(input: "What is the capital of France?")
     #
 
-    NAME = "search"
+    NAME = "google_search"
 
     description <<~DESC
       A wrapper around Google Search.
@@ -26,10 +26,10 @@ module Langchain::Tool
     attr_reader :api_key
 
     #
-    # Initializes the SerpAPI tool
+    # Initializes the Google Search tool
     #
-    # @param api_key [String] SerpAPI API key
-    # @return [Langchain::Tool::SerpApi] SerpAPI tool
+    # @param api_key [String] Search API key
+    # @return [Langchain::Tool::GoogleSearch] Google search tool
     #
     def initialize(api_key:)
       depends_on "google_search_results"
@@ -54,7 +54,7 @@ module Langchain::Tool
     # @return [String] Answer
     #
     def execute(input:)
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Executing \"#{input}\"")
+      Langchain.logger.info("Executing \"#{input}\"", for: self.class)
 
       hash_results = execute_search(input: input)
 
@@ -72,7 +72,7 @@ module Langchain::Tool
     # @return [Hash] hash_results JSON
     #
     def execute_search(input:)
-      GoogleSearch
+      ::GoogleSearch
         .new(q: input, serp_api_key: api_key)
         .get_hash
     end

data/lib/langchain/tool/ruby_code_interpreter.rb

@@ -21,7 +21,7 @@ module Langchain::Tool
     # @param input [String] ruby code expression
     # @return [String] Answer
     def execute(input:)
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Executing \"#{input}\"")
+      Langchain.logger.info("Executing \"#{input}\"", for: self.class)
 
       safe_eval(input)
     end

data/lib/langchain/tool/weather.rb

@@ -21,7 +21,7 @@ module Langchain::Tool
 
     description <<~DESC
       Useful for getting current weather data
-    
+
       The input to this tool should be a city name followed by the units (imperial, metric, or standard)
       Usage:
         Action Input: St Louis, Missouri; metric
@@ -54,7 +54,7 @@ module Langchain::Tool
     # @param input [String] comma separated city and unit (optional: imperial, metric, or standard)
     # @return [String] Answer
     def execute(input:)
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Executing for \"#{input}\"")
+      Langchain.logger.info("Executing for \"#{input}\"", for: self.class)
 
       input_array = input.split(";")
       city, units = *input_array.map(&:strip)

data/lib/langchain/tool/wikipedia.rb

@@ -26,7 +26,7 @@ module Langchain::Tool
     # @param input [String] search query
     # @return [String] Answer
     def execute(input:)
-      Langchain.logger.info("[#{self.class.name}]".light_blue + ": Executing \"#{input}\"")
+      Langchain.logger.info("Executing \"#{input}\"", for: self.class)
 
       page = ::Wikipedia.find(input)
       # It would be nice to figure out a way to provide page.content but the LLM token limit is an issue

data/lib/langchain/utils/token_length/base_validator.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Langchain
+  module Utils
+    module TokenLength
+      #
+      # Calculate the `max_tokens:` parameter to be set by calculating the context length of the text minus the prompt length
+      #
+      # @param content [String | Array<String>] The text or array of texts to validate
+      # @param model_name [String] The model name to validate against
+      # @return [Integer] Whether the text is valid or not
+      # @raise [TokenLimitExceeded] If the text is too long
+      #
+      class BaseValidator
+        def self.validate_max_tokens!(content, model_name, options = {})
+          text_token_length = if content.is_a?(Array)
+            content.sum { |item| token_length(item.to_json, model_name, options) }
+          else
+            token_length(content, model_name, options)
+          end
+
+          leftover_tokens = token_limit(model_name) - text_token_length
+
+          # Raise an error even if whole prompt is equal to the model's token limit (leftover_tokens == 0)
+          if leftover_tokens <= 0
+            raise limit_exceeded_exception(token_limit(model_name), text_token_length)
+          end
+
+          leftover_tokens
+        end
+
+        def self.limit_exceeded_exception(limit, length)
+          TokenLimitExceeded.new("This model's maximum context length is #{limit} tokens, but the given text is #{length} tokens long.", length - limit)
+        end
+      end
+    end
+  end
+end

data/lib/langchain/utils/token_length/google_palm_validator.rb

@@ -7,7 +7,7 @@ module Langchain
       # This class is meant to validate the length of the text passed in to Google Palm's API.
       # It is used to validate the token length before the API call is made
       #
-      class GooglePalmValidator
+      class GooglePalmValidator < BaseValidator
         TOKEN_LIMITS = {
           # Source:
           # This data can be pulled when `list_models()` method is called: https://github.com/andreibondarev/google_palm_api#usage
@@ -26,43 +26,23 @@ module Langchain
           # }
         }.freeze
 
-        #
-        # Validate the context length of the text
-        #
-        # @param content [String | Array<String>] The text or array of texts to validate
-        # @param model_name [String] The model name to validate against
-        # @return [Integer] Whether the text is valid or not
-        # @raise [TokenLimitExceeded] If the text is too long
-        #
-        def self.validate_max_tokens!(google_palm_llm, content, model_name)
-          text_token_length = if content.is_a?(Array)
-            content.sum { |item| token_length(google_palm_llm, item.to_json, model_name) }
-          else
-            token_length(google_palm_llm, content, model_name)
-          end
-
-          leftover_tokens = TOKEN_LIMITS.dig(model_name, "input_token_limit") - text_token_length
-
-          # Raise an error even if whole prompt is equal to the model's token limit (leftover_tokens == 0)
-          if leftover_tokens <= 0
-            raise TokenLimitExceeded, "This model's maximum context length is #{TOKEN_LIMITS.dig(model_name, "input_token_limit")} tokens, but the given text is #{text_token_length} tokens long."
-          end
-
-          leftover_tokens
-        end
-
         #
         # Calculate token length for a given text and model name
         #
-        # @param llm [Langchain::LLM:GooglePalm] The Langchain::LLM:GooglePalm instance
         # @param text [String] The text to calculate the token length for
         # @param model_name [String] The model name to validate against
+        # @param options [Hash] the options to create a message with
+        # @option options [Langchain::LLM:GooglePalm] :llm The Langchain::LLM:GooglePalm instance
         # @return [Integer] The token length of the text
         #
-        def self.token_length(llm, text, model_name = "chat-bison-001")
-          response = llm.client.count_message_tokens(model: model_name, prompt: text)
+        def self.token_length(text, model_name = "chat-bison-001", options)
+          response = options[:llm].client.count_message_tokens(model: model_name, prompt: text)
           response.dig("tokenCount")
         end
+
+        def self.token_limit(model_name)
+          TOKEN_LIMITS.dig(model_name, "input_token_limit")
+        end
       end
     end
   end