langchainrb 0.6.16 → 0.6.18

data/lib/langchain/llm/google_palm.rb

@@ -8,7 +8,7 @@ module Langchain::LLM
   #     gem "google_palm_api", "~> 0.1.3"
   #
   # Usage:
-  #     google_palm = Langchain::LLM::GooglePalm.new(api_key: "YOUR_API_KEY")
+  #     google_palm = Langchain::LLM::GooglePalm.new(api_key: ENV["GOOGLE_PALM_API_KEY"])
   #
   class GooglePalm < Base
     DEFAULTS = {
@@ -20,7 +20,7 @@ module Langchain::LLM
     }.freeze
     LENGTH_VALIDATOR = Langchain::Utils::TokenLength::GooglePalmValidator
     ROLE_MAPPING = {
-      "human" => "user"
+      "assistant" => "ai"
     }
 
     def initialize(api_key:, default_options: {})
@@ -34,13 +34,13 @@ module Langchain::LLM
     # Generate an embedding for a given text
     #
     # @param text [String] The text to generate an embedding for
-    # @return [Array] The embedding
+    # @return [Langchain::LLM::GooglePalmResponse] Response object
     #
     def embed(text:)
-      response = client.embed(
-        text: text
-      )
-      response.dig("embedding", "value")
+      response = client.embed(text: text)
+
+      Langchain::LLM::GooglePalmResponse.new response,
+        model: @defaults[:embeddings_model_name]
     end
 
     #
@@ -48,7 +48,7 @@ module Langchain::LLM
     #
     # @param prompt [String] The prompt to generate a completion for
     # @param params extra parameters passed to GooglePalmAPI::Client#generate_text
-    # @return [String] The completion
+    # @return [Langchain::LLM::GooglePalmResponse] Response object
     #
     def complete(prompt:, **params)
       default_params = {
@@ -68,18 +68,20 @@ module Langchain::LLM
       default_params.merge!(params)
 
       response = client.generate_text(**default_params)
-      response.dig("candidates", 0, "output")
+
+      Langchain::LLM::GooglePalmResponse.new response,
+        model: default_params[:model]
     end
 
     #
     # Generate a chat completion for a given prompt
     #
-    # @param prompt [HumanMessage] The prompt to generate a chat completion for
-    # @param messages [Array<AIMessage|HumanMessage>] The messages that have been sent in the conversation
-    # @param context [SystemMessage] 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<AIMessage|HumanMessage>] Examples of messages to provide to the model. Useful for Few-Shot Prompting
+    # @param prompt [String] The prompt to generate a chat completion for
+    # @param messages [Array<Hash>] The messages that have been sent in the conversation
+    # @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
     # @param options [Hash] extra parameters passed to GooglePalmAPI::Client#generate_chat_message
-    # @return [AIMessage] The chat completion
+    # @return [Langchain::LLM::GooglePalmResponse] Response object
     #
     def chat(prompt: "", messages: [], context: "", examples: [], **options)
       raise ArgumentError.new(":prompt or :messages argument is expected") if prompt.empty? && messages.empty?
@@ -87,7 +89,7 @@ module Langchain::LLM
       default_params = {
         temperature: @defaults[:temperature],
         model: @defaults[:chat_completion_model_name],
-        context: context.to_s,
+        context: context,
         messages: compose_chat_messages(prompt: prompt, messages: messages),
         examples: compose_examples(examples)
       }
@@ -108,7 +110,9 @@ module Langchain::LLM
       response = client.generate_chat_message(**default_params)
       raise "GooglePalm API returned an error: #{response}" if response.dig("error")
 
-      Langchain::AIMessage.new(response.dig("candidates", 0, "content"))
+      Langchain::LLM::GooglePalmResponse.new response,
+        model: default_params[:model]
+      # TODO: Pass in prompt_tokens: prompt_tokens
     end
 
     #
@@ -150,8 +154,8 @@ module Langchain::LLM
     def compose_examples(examples)
       examples.each_slice(2).map do |example|
         {
-          input: {content: example.first.content},
-          output: {content: example.last.content}
+          input: {content: example.first[:content]},
+          output: {content: example.last[:content]}
         }
       end
     end
@@ -159,8 +163,8 @@ module Langchain::LLM
     def transform_messages(messages)
       messages.map do |message|
         {
-          author: ROLE_MAPPING.fetch(message.type, message.type),
-          content: message.content
+          author: ROLE_MAPPING.fetch(message[:role], message[:role]),
+          content: message[:content]
         }
       end
     end

data/lib/langchain/llm/hugging_face.rb

@@ -8,7 +8,7 @@ module Langchain::LLM
   #     gem "hugging-face", "~> 0.3.4"
   #
   # Usage:
-  #     hf = Langchain::LLM::HuggingFace.new(api_key: "YOUR_API_KEY")
+  #     hf = Langchain::LLM::HuggingFace.new(api_key: ENV["HUGGING_FACE_API_KEY"])
   #
   class HuggingFace < Base
     # The gem does not currently accept other models:
@@ -34,13 +34,14 @@ module Langchain::LLM
     # Generate an embedding for a given text
     #
     # @param text [String] The text to embed
-    # @return [Array] The embedding
+    # @return [Langchain::LLM::HuggingFaceResponse] Response object
     #
     def embed(text:)
-      client.embedding(
+      response = client.embedding(
         input: text,
         model: DEFAULTS[:embeddings_model_name]
       )
+      Langchain::LLM::HuggingFaceResponse.new(response, model: DEFAULTS[:embeddings_model_name])
     end
   end
 end

data/lib/langchain/llm/llama_cpp.rb

@@ -34,7 +34,7 @@ module Langchain::LLM
 
     # @param text [String] The text to embed
     # @param n_threads [Integer] The number of CPU threads to use
-    # @return [Array] The embedding
+    # @return [Array<Float>] The embedding
     def embed(text:, n_threads: nil)
       # contexts are kinda stateful when it comes to embeddings, so allocate one each time
       context = embedding_context

data/lib/langchain/llm/ollama.rb

@@ -22,18 +22,23 @@ module Langchain::LLM
       @url = url
     end
 
+    #
     # Generate the completion for a given prompt
+    #
     # @param prompt [String] The prompt to complete
     # @param model [String] The model to use
     # @param options [Hash] The options to use (https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
-    # @return [String] The completed prompt
+    # @return [Langchain::LLM::OllamaResponse] Response object
+    #
     def complete(prompt:, model: nil, **options)
       response = +""
 
+      model_name = model || DEFAULTS[:completion_model_name]
+
       client.post("api/generate") do |req|
         req.body = {}
         req.body["prompt"] = prompt
-        req.body["model"] = model || DEFAULTS[:completion_model_name]
+        req.body["model"] = model_name
 
         req.body["options"] = options if options.any?
 
@@ -47,27 +52,34 @@ module Langchain::LLM
         end
       end
 
-      response
+      Langchain::LLM::OllamaResponse.new(response, model: model_name)
     end
 
+    #
     # Generate an embedding for a given text
+    #
     # @param text [String] The text to generate an embedding for
     # @param model [String] The model to use
-    # @param options [Hash] The options to use (
+    # @param options [Hash] The options to use
+    # @return [Langchain::LLM::OllamaResponse] Response object
+    #
     def embed(text:, model: nil, **options)
+      model_name = model || DEFAULTS[:embeddings_model_name]
+
       response = client.post("api/embeddings") do |req|
         req.body = {}
         req.body["prompt"] = text
-        req.body["model"] = model || DEFAULTS[:embeddings_model_name]
+        req.body["model"] = model_name
 
         req.body["options"] = options if options.any?
       end
 
-      response.body.dig("embedding")
+      Langchain::LLM::OllamaResponse.new(response.body, model: model_name)
     end
 
     private
 
+    # @return [Faraday::Connection] Faraday client
     def client
       @client ||= Faraday.new(url: url) do |conn|
         conn.request :json

data/lib/langchain/llm/openai.rb

@@ -11,6 +11,7 @@ module Langchain::LLM
   #
   class OpenAI < Base
     DEFAULTS = {
+      n: 1,
       temperature: 0.0,
       completion_model_name: "gpt-3.5-turbo",
       chat_completion_model_name: "gpt-3.5-turbo",
@@ -26,10 +27,6 @@ module Langchain::LLM
     ].freeze
 
     LENGTH_VALIDATOR = Langchain::Utils::TokenLength::OpenAIValidator
-    ROLE_MAPPING = {
-      "ai" => "assistant",
-      "human" => "user"
-    }
 
     attr_accessor :functions
 
@@ -45,7 +42,7 @@ module Langchain::LLM
     #
     # @param text [String] The text to generate an embedding for
     # @param params extra parameters passed to OpenAI::Client#embeddings
-    # @return [Array] The embedding
+    # @return [Langchain::LLM::OpenAIResponse] Response object
     #
     def embed(text:, **params)
       parameters = {model: @defaults[:embeddings_model_name], input: text}
@@ -56,7 +53,7 @@ module Langchain::LLM
         client.embeddings(parameters: parameters.merge(params))
       end
 
-      response.dig("data").first.dig("embedding")
+      Langchain::LLM::OpenAIResponse.new(response)
     end
 
     #
@@ -64,7 +61,7 @@ module Langchain::LLM
     #
     # @param prompt [String] The prompt to generate a completion for
     # @param params  extra parameters passed to OpenAI::Client#complete
-    # @return [String] The completion
+    # @return [Langchain::LLM::Response::OpenaAI] Response object
     #
     def complete(prompt:, **params)
       parameters = compose_parameters @defaults[:completion_model_name], params
@@ -78,7 +75,7 @@ module Langchain::LLM
         client.chat(parameters: parameters)
       end
 
-      response.dig("choices", 0, "message", "content")
+      Langchain::LLM::OpenAIResponse.new(response)
     end
 
     #
@@ -117,18 +114,18 @@ module Langchain::LLM
     #         },
     #       ]
     #
-    # @param prompt [HumanMessage] The prompt to generate a chat completion for
-    # @param messages [Array<AIMessage|HumanMessage>] The messages that have been sent in the conversation
-    # @param context [SystemMessage] 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<AIMessage|HumanMessage>] Examples of messages to provide to the model. Useful for Few-Shot Prompting
+    # @param prompt [String] The prompt to generate a chat completion for
+    # @param messages [Array<Hash>] The messages that have been sent in the conversation
+    # @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
     # @param options [Hash] extra parameters passed to OpenAI::Client#chat
-    # @yield [AIMessage] Stream responses back one String at a time
-    # @return [AIMessage] The chat completion
+    # @yield [Hash] Stream responses back one token at a time
+    # @return [Langchain::LLM::OpenAIResponse] Response object
     #
-    def chat(prompt: "", messages: [], context: "", examples: [], **options)
+    def chat(prompt: "", messages: [], context: "", examples: [], **options, &block)
       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, &block
       parameters[:messages] = compose_chat_messages(prompt: prompt, messages: messages, context: context, examples: examples)
 
       if functions
@@ -137,25 +134,11 @@ module Langchain::LLM
         parameters[:max_tokens] = validate_max_tokens(parameters[:messages], parameters[:model])
       end
 
-      if (streaming = block_given?)
-        parameters[:stream] = proc do |chunk, _bytesize|
-          delta = chunk.dig("choices", 0, "delta")
-          content = delta["content"]
-          additional_kwargs = {function_call: delta["function_call"]}.compact
-          yield Langchain::AIMessage.new(content, additional_kwargs)
-        end
-      end
+      response = with_api_error_handling { client.chat(parameters: parameters) }
 
-      response = with_api_error_handling do
-        client.chat(parameters: parameters)
-      end
+      return if block
 
-      unless streaming
-        message = response.dig("choices", 0, "message")
-        content = message["content"]
-        additional_kwargs = {function_call: message["function_call"]}.compact
-        Langchain::AIMessage.new(content.to_s, additional_kwargs)
-      end
+      Langchain::LLM::OpenAIResponse.new(response)
     end
 
     #
@@ -171,6 +154,7 @@ module Langchain::LLM
       prompt = prompt_template.format(text: text)
 
       complete(prompt: prompt, temperature: @defaults[:temperature])
+      # Should this return a Langchain::LLM::OpenAIResponse as well?
     end
 
     private
@@ -191,12 +175,18 @@ module Langchain::LLM
       response.dig("choices", 0, "text")
     end
 
-    def compose_parameters(model, params)
-      default_params = {model: model, temperature: @defaults[:temperature]}
-
+    def compose_parameters(model, params, &block)
+      default_params = {model: model, temperature: @defaults[:temperature], n: @defaults[:n]}
       default_params[:stop] = params.delete(:stop_sequences) if params[:stop_sequences]
+      parameters = default_params.merge(params)
 
-      default_params.merge(params)
+      if block
+        parameters[:stream] = proc do |chunk, _bytesize|
+          yield chunk.dig("choices", 0)
+        end
+      end
+
+      parameters
     end
 
     def compose_chat_messages(prompt:, messages: [], context: "", examples: [])
@@ -206,9 +196,9 @@ module Langchain::LLM
 
       history.concat transform_messages(messages) unless messages.empty?
 
-      unless context.nil? || context.to_s.empty?
+      unless context.nil? || context.empty?
         history.reject! { |message| message[:role] == "system" }
-        history.prepend({role: "system", content: context.content})
+        history.prepend({role: "system", content: context})
       end
 
       unless prompt.empty?
@@ -225,14 +215,16 @@ module Langchain::LLM
     def transform_messages(messages)
       messages.map do |message|
         {
-          role: ROLE_MAPPING.fetch(message.type, message.type),
-          content: message.content
+          role: message[:role],
+          content: message[:content]
         }
       end
     end
 
     def with_api_error_handling
       response = yield
+      return if response.empty?
+
       raise Langchain::LLM::ApiError.new "OpenAI API error: #{response.dig("error", "message")}" if response&.dig("error")
 
       response
@@ -241,5 +233,10 @@ module Langchain::LLM
     def validate_max_tokens(messages, model)
       LENGTH_VALIDATOR.validate_max_tokens!(messages, model)
     end
+
+    def extract_response(response)
+      results = response.dig("choices").map { |choice| choice.dig("message", "content") }
+      (results.size == 1) ? results.first : results
+    end
   end
 end

data/lib/langchain/llm/replicate.rb

@@ -47,44 +47,40 @@ module Langchain::LLM
     # Generate an embedding for a given text
     #
     # @param text [String] The text to generate an embedding for
-    # @return [Hash] The embedding
+    # @return [Langchain::LLM::ReplicateResponse] Response object
     #
     def embed(text:)
       response = embeddings_model.predict(input: text)
 
       until response.finished?
         response.refetch
-        sleep(1)
+        sleep(0.1)
       end
 
-      response.output
+      Langchain::LLM::ReplicateResponse.new(response, model: @defaults[:embeddings_model_name])
     end
 
     #
     # Generate a completion for a given prompt
     #
     # @param prompt [String] The prompt to generate a completion for
-    # @return [Hash] The completion
+    # @return [Langchain::LLM::ReplicateResponse] Reponse object
     #
     def complete(prompt:, **params)
       response = completion_model.predict(prompt: prompt)
 
       until response.finished?
         response.refetch
-        sleep(1)
+        sleep(0.1)
       end
 
-      # Response comes back as an array of strings, e.g.: ["Hi", "how ", "are ", "you?"]
-      # The first array element is missing a space at the end, so we add it manually
-      response.output[0] += " "
-
-      response.output.join
+      Langchain::LLM::ReplicateResponse.new(response, model: @defaults[:completion_model_name])
     end
 
     # Cohere does not have a dedicated chat endpoint, so instead we call `complete()`
     def chat(...)
       response_text = complete(...)
-      Langchain::AIMessage.new(response_text)
+      ::Langchain::Conversation::Response.new(response_text)
     end
 
     #

data/lib/langchain/llm/response/ai21_response.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class AI21Response < BaseResponse
+    def completions
+      raw_response.dig(:completions)
+    end
+
+    def completion
+      completions.dig(0, :data, :text)
+    end
+  end
+end

data/lib/langchain/llm/response/anthropic_response.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class AnthropicResponse < BaseResponse
+    def model
+      raw_response.dig("model")
+    end
+
+    def completion
+      completions.first
+    end
+
+    def completions
+      [raw_response.dig("completion")]
+    end
+
+    def stop_reason
+      raw_response.dig("stop_reason")
+    end
+
+    def stop
+      raw_response.dig("stop")
+    end
+
+    def log_id
+      raw_response.dig("log_id")
+    end
+  end
+end

data/lib/langchain/llm/response/base_response.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Langchain
+  module LLM
+    class BaseResponse
+      attr_reader :raw_response, :model
+
+      def initialize(raw_response, model: nil)
+        @raw_response = raw_response
+        @model = model
+      end
+
+      # Returns the completion text
+      #
+      # @return [String]
+      #
+      def completion
+        raise NotImplementedError
+      end
+
+      # Returns the chat completion text
+      #
+      # @return [String]
+      #
+      def chat_completion
+        raise NotImplementedError
+      end
+
+      # Return the first embedding
+      #
+      # @return [Array<Float>]
+      def embedding
+        raise NotImplementedError
+      end
+
+      # Return the completion candidates
+      #
+      # @return [Array]
+      def completions
+        raise NotImplementedError
+      end
+
+      # Return the chat completion candidates
+      #
+      # @return [Array]
+      def chat_completions
+        raise NotImplementedError
+      end
+
+      # Return the embeddings
+      #
+      # @return [Array<Array>]
+      def embeddings
+        raise NotImplementedError
+      end
+
+      # Number of tokens utilized in the prompt
+      #
+      # @return [Integer]
+      def prompt_tokens
+        raise NotImplementedError
+      end
+
+      # Number of tokens utilized to generate the completion
+      #
+      # @return [Integer]
+      def completion_tokens
+        raise NotImplementedError
+      end
+
+      # Total number of tokens utilized
+      #
+      # @return [Integer]
+      def total_tokens
+        raise NotImplementedError
+      end
+    end
+  end
+end

data/lib/langchain/llm/response/cohere_response.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class CohereResponse < BaseResponse
+    def embedding
+      embeddings.first
+    end
+
+    def embeddings
+      raw_response.dig("embeddings")
+    end
+
+    def completions
+      raw_response.dig("generations")
+    end
+
+    def completion
+      completions&.dig(0, "text")
+    end
+  end
+end

data/lib/langchain/llm/response/google_palm_response.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class GooglePalmResponse < BaseResponse
+    attr_reader :prompt_tokens
+
+    def initialize(raw_response, model: nil, prompt_tokens: nil)
+      @prompt_tokens = prompt_tokens
+      super(raw_response, model: model)
+    end
+
+    def completion
+      completions&.dig(0, "output")
+    end
+
+    def embedding
+      embeddings.first
+    end
+
+    def completions
+      raw_response.dig("candidates")
+    end
+
+    def chat_completion
+      chat_completions&.dig(0, "content")
+    end
+
+    def chat_completions
+      raw_response.dig("candidates")
+    end
+
+    def embeddings
+      [raw_response.dig("embedding", "value")]
+    end
+  end
+end

data/lib/langchain/llm/response/hugging_face_response.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class HuggingFaceResponse < BaseResponse
+    def embeddings
+      [raw_response]
+    end
+
+    def embedding
+      embeddings.first
+    end
+  end
+end

data/lib/langchain/llm/response/ollama_response.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class OllamaResponse < BaseResponse
+    def initialize(raw_response, model: nil, prompt_tokens: nil)
+      @prompt_tokens = prompt_tokens
+      super(raw_response, model: model)
+    end
+
+    def completion
+      raw_response.first
+    end
+
+    def completions
+      raw_response.is_a?(String) ? [raw_response] : []
+    end
+
+    def embedding
+      embeddings.first
+    end
+
+    def embeddings
+      [raw_response&.dig("embedding")]
+    end
+  end
+end