langchainrb 0.6.17 → 0.6.18

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 3b9bca59bfb5909f6ac24ebf6dba6074f5faf3d2cdadab1a3b3a8a0f75f98adc
-  data.tar.gz: a202726d383d2dc691cb4146e9b36cb7ea6f8ac35382a3df67f6e11d35b3562e
+  metadata.gz: 437c6387ded139ed1a513414bfb7242cdbadf1ba6526c7a89346aa2fa9490fc2
+  data.tar.gz: dd6f437a4bbc4807a16631dd790f66c9de4e9456011b2c4f84302fe3fab1377b
 SHA512:
-  metadata.gz: b4eaf631f22236035c9e29b3618a70d14487cc9e39b6885e44497ebad2a98670ce88997fdb25144b6467e0caa69a04ce7e625c9e10bc88322131181c2254a570
-  data.tar.gz: 981199fe2a0123e46ac3af54946c03d5eaa827473eae02f2e60accd0c680a0bbd40741800e05b79e890038523a1b910502a6cf4ed1f4ebf77845f4b2a2dbc5d9
+  metadata.gz: 24748539de50dfa816fdb71173ef00a6b04f9737f32926fca919865a49b9812dd9f1fdb286c361c98e33cc994f67e8988ab688bfdf6bf3020d954eb0c791177c
+  data.tar.gz: 283b10460187cada7485e08a19c89e7485925ab2f73a5ad51b06a72e8fd9ee1600ddac9d000f13c0c1af13f6defece9fdcc272489d0df803f94da96fe1c76cfd

data/CHANGELOG.md

@@ -1,5 +1,16 @@
 ## [Unreleased]
 
+## [0.6.18] - 2023-10-16
+- Introduce `Langchain::LLM::Response`` object
+- Introduce `Langchain::Chunk` object
+- Add the ask() method to the Langchain::ActiveRecord::Hooks
+
+## [0.6.17] - 2023-10-10
+- Bump weaviate and chroma-db deps
+- `Langchain::Chunker::Semantic` chunker
+- Re-structure Conversations class
+- Bug fixes
+
 ## [0.6.16] - 2023-10-02
 - HyDE-style similarity search
 - `Langchain::Chunker::Sentence` chunker

data/README.md

@@ -128,6 +128,21 @@ class Product < ActiveRecord::Base
 end
 ```
 
+### Exposed ActiveRecord methods
+```ruby
+# Retrieve similar products based on the query string passed in
+Product.similarity_search(
+    query:,
+    k:       # number of results to be retrieved
+)
+```
+```ruby
+# Q&A-style querying based on the question passed in
+Product.ask(
+    question:
+)
+```
+
 Additional info [here](https://github.com/andreibondarev/langchainrb/blob/main/lib/langchain/active_record/hooks.rb#L10-L38).
 
 ### Using Standalone LLMs 🗣️

data/lib/langchain/active_record/hooks.rb

@@ -92,6 +92,20 @@ module Langchain
           ids = records.map { |record| record.dig("id") || record.dig("__id") }
           where(id: ids)
         end
+
+        # Ask a question and return the answer
+        #
+        # @param question [String] The question to ask
+        # @param k [Integer] The number of results to have in context
+        # @yield [String] Stream responses back one String at a time
+        # @return [String] The answer to the question
+        def ask(question:, k: 4, &block)
+          class_variable_get(:@@provider).ask(
+            question: question,
+            k: k,
+            &block
+          )
+        end
       end
     end
   end

data/lib/langchain/agent/react_agent.rb

@@ -58,7 +58,7 @@ module Langchain::Agent
       max_iterations.times do
         Langchain.logger.info("Sending the prompt to the #{llm.class} LLM", for: self.class)
 
-        response = llm.complete(prompt: prompt, stop_sequences: ["Observation:"])
+        response = llm.complete(prompt: prompt, stop_sequences: ["Observation:"]).completion
 
         # Append the response to the prompt
         prompt += response

data/lib/langchain/agent/sql_query_agent.rb

@@ -27,7 +27,7 @@ module Langchain::Agent
 
       # Get the SQL string to execute
       Langchain.logger.info("Passing the inital prompt to the #{llm.class} LLM", for: self.class)
-      sql_string = llm.complete(prompt: prompt)
+      sql_string = llm.complete(prompt: prompt).completion
 
       # Execute the SQL string and collect the results
       Langchain.logger.info("Passing the SQL to the Database: #{sql_string}", for: self.class)
@@ -36,7 +36,7 @@ module Langchain::Agent
       # Pass the results and get the LLM to synthesize the answer to the question
       Langchain.logger.info("Passing the synthesize prompt to the #{llm.class} LLM with results: #{results}", for: self.class)
       prompt2 = create_prompt_for_answer(question: question, sql_query: sql_string, results: results)
-      llm.complete(prompt: prompt2)
+      llm.complete(prompt: prompt2).completion
     end
 
     private

data/lib/langchain/chunk.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module Langchain
+  class Chunk
+    # The chunking process is the process of splitting a document into smaller chunks and creating instances of Langchain::Chunk
+
+    attr_reader :text
+
+    # Initialize a new chunk
+    # @param [String] text
+    # @return [Langchain::Chunk]
+    def initialize(text:)
+      @text = text
+    end
+  end
+end

data/lib/langchain/chunker/base.rb

@@ -13,6 +13,10 @@ module Langchain
     # - {Langchain::Chunker::Semantic}
     # - {Langchain::Chunker::Sentence}
     class Base
+      # @return [Array<Langchain::Chunk>]
+      def chunks
+        raise NotImplementedError
+      end
     end
   end
 end

data/lib/langchain/chunker/recursive_text.rb

@@ -24,14 +24,17 @@ module Langchain
         @separators = separators
       end
 
-      # @return [Array<String>]
+      # @return [Array<Langchain::Chunk>]
       def chunks
         splitter = Baran::RecursiveCharacterTextSplitter.new(
           chunk_size: chunk_size,
           chunk_overlap: chunk_overlap,
           separators: separators
         )
-        splitter.chunks(text)
+
+        splitter.chunks(text).map do |chunk|
+          Langchain::Chunk.new(text: chunk[:text])
+        end
       end
     end
   end

data/lib/langchain/chunker/semantic.rb

@@ -23,7 +23,7 @@ module Langchain
         @prompt_template = prompt_template || default_prompt_template
       end
 
-      # @return [Array<String>]
+      # @return [Array<Langchain::Chunk>]
       def chunks
         prompt = prompt_template.format(text: text)
 
@@ -34,6 +34,9 @@ module Langchain
           .split("---")
           .map(&:strip)
           .reject(&:empty?)
+          .map do |chunk|
+            Langchain::Chunk.new(text: chunk)
+          end
       end
 
       private

data/lib/langchain/chunker/sentence.rb

@@ -19,10 +19,12 @@ module Langchain
         @text = text
       end
 
-      # @return [Array<String>]
+      # @return [Array<Langchain::Chunk>]
       def chunks
         ps = PragmaticSegmenter::Segmenter.new(text: text)
-        ps.segment
+        ps.segment.map do |chunk|
+          Langchain::Chunk.new(text: chunk)
+        end
       end
     end
   end

data/lib/langchain/chunker/text.rb

@@ -24,14 +24,17 @@ module Langchain
         @separator = separator
       end
 
-      # @return [Array<String>]
+      # @return [Array<Langchain::Chunk>]
       def chunks
         splitter = Baran::CharacterTextSplitter.new(
           chunk_size: chunk_size,
           chunk_overlap: chunk_overlap,
           separator: separator
         )
-        splitter.chunks(text)
+
+        splitter.chunks(text).map do |chunk|
+          Langchain::Chunk.new(text: chunk[:text])
+        end
       end
     end
   end

data/lib/langchain/conversation.rb

@@ -58,7 +58,7 @@ module Langchain
     # @return [Response] The response from the model
     def message(message)
       @memory.append_message ::Langchain::Conversation::Prompt.new(message)
-      ai_message = ::Langchain::Conversation::Response.new(llm_response)
+      ai_message = ::Langchain::Conversation::Response.new(llm_response.chat_completion)
       @memory.append_message(ai_message)
       ai_message
     end

data/lib/langchain/llm/ai21.rb

@@ -8,7 +8,7 @@ module Langchain::LLM
   #   gem "ai21", "~> 0.2.1"
   #
   # Usage:
-  #     ai21 = Langchain::LLM::AI21.new(api_key:)
+  #     ai21 = Langchain::LLM::AI21.new(api_key: ENV["AI21_API_KEY"])
   #
   class AI21 < Base
     DEFAULTS = {
@@ -30,7 +30,7 @@ module Langchain::LLM
     #
     # @param prompt [String] The prompt to generate a completion for
     # @param params [Hash] The parameters to pass to the API
-    # @return [String] The completion
+    # @return [Langchain::LLM::AI21Response] The completion
     #
     def complete(prompt:, **params)
       parameters = complete_parameters params
@@ -38,7 +38,7 @@ module Langchain::LLM
       parameters[:maxTokens] = LENGTH_VALIDATOR.validate_max_tokens!(prompt, parameters[:model], client)
 
       response = client.complete(prompt, parameters)
-      response.dig(:completions, 0, :data, :text)
+      Langchain::LLM::AI21Response.new response, model: parameters[:model]
     end
 
     #
@@ -51,6 +51,7 @@ module Langchain::LLM
     def summarize(text:, **params)
       response = client.summarize(text, "TEXT", params)
       response.dig(:summary)
+      # Should we update this to also return a Langchain::LLM::AI21Response?
     end
 
     private

data/lib/langchain/llm/anthropic.rb

@@ -8,7 +8,7 @@ module Langchain::LLM
   #   gem "anthropic", "~> 0.1.0"
   #
   # Usage:
-  #     anthorpic = Langchain::LLM::Anthropic.new(api_key:)
+  #     anthorpic = Langchain::LLM::Anthropic.new(api_key: ENV["ANTHROPIC_API_KEY"])
   #
   class Anthropic < Base
     DEFAULTS = {
@@ -32,7 +32,7 @@ module Langchain::LLM
     #
     # @param prompt [String] The prompt to generate a completion for
     # @param params [Hash] extra parameters passed to Anthropic::Client#complete
-    # @return [String] The completion
+    # @return [Langchain::LLM::AnthropicResponse] The completion
     #
     def complete(prompt:, **params)
       parameters = compose_parameters @defaults[:completion_model_name], params
@@ -43,7 +43,7 @@ module Langchain::LLM
       # parameters[:max_tokens_to_sample] = validate_max_tokens(prompt, parameters[:completion_model_name])
 
       response = client.complete(parameters: parameters)
-      response.dig("completion")
+      Langchain::LLM::AnthropicResponse.new(response)
     end
 
     private

data/lib/langchain/llm/cohere.rb

@@ -8,7 +8,7 @@ module Langchain::LLM
   #     gem "cohere-ruby", "~> 0.9.6"
   #
   # Usage:
-  #     cohere = Langchain::LLM::Cohere.new(api_key: "YOUR_API_KEY")
+  #     cohere = Langchain::LLM::Cohere.new(api_key: ENV["COHERE_API_KEY"])
   #
   class Cohere < Base
     DEFAULTS = {
@@ -30,14 +30,15 @@ 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::CohereResponse] Response object
     #
     def embed(text:)
       response = client.embed(
         texts: [text],
         model: @defaults[:embeddings_model_name]
       )
-      response.dig("embeddings").first
+
+      Langchain::LLM::CohereResponse.new response, model: @defaults[:embeddings_model_name]
     end
 
     #
@@ -45,7 +46,7 @@ module Langchain::LLM
     #
     # @param prompt [String] The prompt to generate a completion for
     # @param params[:stop_sequences]
-    # @return [Hash] The completion
+    # @return [Langchain::LLM::CohereResponse] Response object
     #
     def complete(prompt:, **params)
       default_params = {
@@ -64,7 +65,7 @@ module Langchain::LLM
       default_params[:max_tokens] = Langchain::Utils::TokenLength::CohereValidator.validate_max_tokens!(prompt, default_params[:model], client)
 
       response = client.generate(**default_params)
-      response.dig("generations").first.dig("text")
+      Langchain::LLM::CohereResponse.new response, model: @defaults[:completion_model_name]
     end
 
     # Cohere does not have a dedicated chat endpoint, so instead we call `complete()`

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 = {
@@ -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,7 +68,9 @@ 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
 
     #
@@ -79,7 +81,7 @@ module Langchain::LLM
     # @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 [String] 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?
@@ -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")
 
-      response.dig("candidates", 0, "content")
+      Langchain::LLM::GooglePalmResponse.new response,
+        model: default_params[:model]
+      # TODO: Pass in prompt_tokens: prompt_tokens
     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

@@ -42,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}
@@ -53,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
 
     #
@@ -61,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
@@ -75,7 +75,7 @@ module Langchain::LLM
         client.chat(parameters: parameters)
       end
 
-      response.dig("choices", 0, "message", "content")
+      Langchain::LLM::OpenAIResponse.new(response)
     end
 
     #
@@ -120,7 +120,7 @@ module Langchain::LLM
     # @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 [Hash] Stream responses back one token at a time
-    # @return [String|Array<String>] The chat completion
+    # @return [Langchain::LLM::OpenAIResponse] Response object
     #
     def chat(prompt: "", messages: [], context: "", examples: [], **options, &block)
       raise ArgumentError.new(":prompt or :messages argument is expected") if prompt.empty? && messages.empty?
@@ -138,7 +138,7 @@ module Langchain::LLM
 
       return if block
 
-      extract_response response
+      Langchain::LLM::OpenAIResponse.new(response)
     end
 
     #
@@ -154,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

data/lib/langchain/llm/replicate.rb

@@ -47,38 +47,34 @@ 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()`

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

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

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class OpenAIResponse < BaseResponse
+    def model
+      raw_response["model"]
+    end
+
+    def created_at
+      if raw_response.dig("created")
+        Time.at(raw_response.dig("created"))
+      end
+    end
+
+    def completion
+      completions&.dig(0, "message", "content")
+    end
+
+    def chat_completion
+      completion
+    end
+
+    def embedding
+      embeddings&.first
+    end
+
+    def completions
+      raw_response.dig("choices")
+    end
+
+    def chat_completions
+      raw_response.dig("choices")
+    end
+
+    def embeddings
+      raw_response.dig("data")&.map { |datum| datum.dig("embedding") }
+    end
+
+    def prompt_tokens
+      raw_response.dig("usage", "prompt_tokens")
+    end
+
+    def completion_tokens
+      raw_response.dig("usage", "completion_tokens")
+    end
+
+    def total_tokens
+      raw_response.dig("usage", "total_tokens")
+    end
+  end
+end

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

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Langchain::LLM
+  class ReplicateResponse < BaseResponse
+    def completions
+      # 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
+      raw_response.output[0] += " "
+      [raw_response.output.join]
+    end
+
+    def completion
+      completions.first
+    end
+
+    def created_at
+      Time.parse(raw_response.created_at)
+    end
+
+    def embedding
+      embeddings.first
+    end
+
+    def embeddings
+      [raw_response.output]
+    end
+  end
+end

data/lib/langchain/vectorsearch/base.rb

@@ -184,7 +184,7 @@ module Langchain::Vectorsearch
         .flatten
         .map do |path|
           data = Langchain::Loader.new(path)&.load&.chunks
-          data.map { |chunk| chunk[:text] }
+          data.map { |chunk| chunk.text }
         end
 
       texts.flatten!

data/lib/langchain/vectorsearch/chroma.rb

@@ -28,13 +28,13 @@ module Langchain::Vectorsearch
     end
 
     # Add a list of texts to the index
-    # @param texts [Array] The list of texts to add
+    # @param texts [Array<String>] The list of texts to add
     # @return [Hash] The response from the server
     def add_texts(texts:, ids: [])
       embeddings = Array(texts).map.with_index do |text, i|
         ::Chroma::Resources::Embedding.new(
           id: ids[i] ? ids[i].to_s : SecureRandom.uuid,
-          embedding: llm.embed(text: text),
+          embedding: llm.embed(text: text).embedding,
           # TODO: Add support for passing metadata
           metadata: {}, # metadatas[index],
           document: text # Do we actually need to store the whole original document?
@@ -49,7 +49,7 @@ module Langchain::Vectorsearch
       embeddings = Array(texts).map.with_index do |text, i|
         ::Chroma::Resources::Embedding.new(
           id: ids[i].to_s,
-          embedding: llm.embed(text: text),
+          embedding: llm.embed(text: text).embedding,
           # TODO: Add support for passing metadata
           metadata: [], # metadatas[index],
           document: text # Do we actually need to store the whole original document?
@@ -85,7 +85,7 @@ module Langchain::Vectorsearch
       query:,
       k: 4
     )
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(
         embedding: embedding,
@@ -94,7 +94,7 @@ module Langchain::Vectorsearch
     end
 
     # Search for similar texts by embedding
-    # @param embedding [Array] The embedding to search for
+    # @param embedding [Array<Float>] The embedding to search for
     # @param k [Integer] The number of results to return
     # @return [Chroma::Resources::Embedding] The response from the server
     def similarity_search_by_vector(

data/lib/langchain/vectorsearch/hnswlib.rb

@@ -35,15 +35,15 @@ module Langchain::Vectorsearch
     #
     # Add a list of texts and corresponding IDs to the index
     #
-    # @param texts [Array] The list of texts to add
-    # @param ids [Array] The list of corresponding IDs (integers) to the texts
+    # @param texts [Array<String>] The list of texts to add
+    # @param ids [Array<Integer>] The list of corresponding IDs (integers) to the texts
     # @return [Boolean] The response from the HNSW library
     #
     def add_texts(texts:, ids:)
       resize_index(texts.size)
 
       Array(texts).each_with_index do |text, i|
-        embedding = llm.embed(text: text)
+        embedding = llm.embed(text: text).embedding
 
         client.add_point(embedding, ids[i])
       end
@@ -64,7 +64,7 @@ module Langchain::Vectorsearch
       query:,
       k: 4
     )
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(
         embedding: embedding,
@@ -75,7 +75,7 @@ module Langchain::Vectorsearch
     #
     # Search for the K nearest neighbors of a given vector
     #
-    # @param embedding [Array] The embedding to search for
+    # @param embedding [Array<Float>] The embedding to search for
     # @param k [Integer] The number of results to return
     # @return [Array] Results in the format `[[id1, distance3], [id2, distance2]]`
     #

data/lib/langchain/vectorsearch/milvus.rb

@@ -32,7 +32,7 @@ module Langchain::Vectorsearch
           }, {
             field_name: "vectors",
             type: ::Milvus::DATA_TYPES["float_vector"],
-            field: Array(texts).map { |text| llm.embed(text: text) }
+            field: Array(texts).map { |text| llm.embed(text: text).embedding }
           }
         ]
       )
@@ -111,7 +111,7 @@ module Langchain::Vectorsearch
     end
 
     def similarity_search(query:, k: 4)
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(
         embedding: embedding,

data/lib/langchain/vectorsearch/pgvector.rb

@@ -52,7 +52,7 @@ module Langchain::Vectorsearch
     # the added or updated texts.
     def upsert_texts(texts:, ids:)
       data = texts.zip(ids).flat_map do |(text, id)|
-        {id: id, content: text, vectors: llm.embed(text: text).to_s, namespace: namespace}
+        {id: id, content: text, vectors: llm.embed(text: text).embedding.to_s, namespace: namespace}
       end
       # @db[table_name.to_sym].multi_insert(data, return: :primary_key)
       @db[table_name.to_sym]
@@ -70,7 +70,7 @@ module Langchain::Vectorsearch
     def add_texts(texts:, ids: nil)
       if ids.nil? || ids.empty?
         data = texts.map do |text|
-          {content: text, vectors: llm.embed(text: text).to_s, namespace: namespace}
+          {content: text, vectors: llm.embed(text: text).embedding.to_s, namespace: namespace}
         end
 
         @db[table_name.to_sym].multi_insert(data, return: :primary_key)
@@ -110,7 +110,7 @@ module Langchain::Vectorsearch
     # @param k [Integer] The number of top results to return
     # @return [Array<Hash>] The results of the search
     def similarity_search(query:, k: 4)
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(
         embedding: embedding,

data/lib/langchain/vectorsearch/pinecone.rb

@@ -31,7 +31,7 @@ module Langchain::Vectorsearch
     end
 
     # Find records by ids
-    # @param ids [Array] The ids to find
+    # @param ids [Array<Integer>] The ids to find
     # @param namespace String The namespace to search through
     # @return [Hash] The response from the server
     def find(ids: [], namespace: "")
@@ -44,8 +44,8 @@ module Langchain::Vectorsearch
     end
 
     # Add a list of texts to the index
-    # @param texts [Array] The list of texts to add
-    # @param ids [Array] The list of IDs to add
+    # @param texts [Array<String>] The list of texts to add
+    # @param ids [Array<Integer>] The list of IDs to add
     # @param namespace [String] The namespace to add the texts to
     # @param metadata [Hash] The metadata to use for the texts
     # @return [Hash] The response from the server
@@ -54,7 +54,7 @@ module Langchain::Vectorsearch
         {
           id: ids[i] ? ids[i].to_s : SecureRandom.uuid,
           metadata: metadata || {content: text},
-          values: llm.embed(text: text)
+          values: llm.embed(text: text).embedding
         }
       end
 
@@ -70,7 +70,7 @@ module Langchain::Vectorsearch
         .flatten
         .map do |path|
           data = Langchain::Loader.new(path)&.load&.chunks
-          data.map { |chunk| chunk[:text] }
+          data.map { |chunk| chunk.text }
         end
 
       texts.flatten!
@@ -79,8 +79,8 @@ module Langchain::Vectorsearch
     end
 
     # Update a list of texts in the index
-    # @param texts [Array] The list of texts to update
-    # @param ids [Array] The list of IDs to update
+    # @param texts [Array<String>] The list of texts to update
+    # @param ids [Array<Integer>] The list of IDs to update
     # @param namespace [String] The namespace to update the texts in
     # @param metadata [Hash] The metadata to use for the texts
     # @return [Array] The response from the server
@@ -90,7 +90,7 @@ module Langchain::Vectorsearch
         index.update(
           namespace: namespace,
           id: ids[i].to_s,
-          values: llm.embed(text: text),
+          values: llm.embed(text: text).embedding,
           set_metadata: metadata
         )
       end
@@ -130,7 +130,7 @@ module Langchain::Vectorsearch
       namespace: "",
       filter: nil
     )
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(
         embedding: embedding,
@@ -141,7 +141,7 @@ module Langchain::Vectorsearch
     end
 
     # Search for similar texts by embedding
-    # @param embedding [Array] The embedding to search for
+    # @param embedding [Array<Float>] The embedding to search for
     # @param k [Integer] The number of results to return
     # @param namespace [String] The namespace to search in
     # @param filter [String] The filter to use

data/lib/langchain/vectorsearch/qdrant.rb

@@ -29,7 +29,7 @@ module Langchain::Vectorsearch
     end
 
     # Find records by ids
-    # @param ids [Array] The ids to find
+    # @param ids [Array<Integer>] The ids to find
     # @return [Hash] The response from the server
     def find(ids: [])
       client.points.get_all(
@@ -41,7 +41,7 @@ module Langchain::Vectorsearch
     end
 
     # Add a list of texts to the index
-    # @param texts [Array] The list of texts to add
+    # @param texts [Array<String>] The list of texts to add
     # @return [Hash] The response from the server
     def add_texts(texts:, ids: [])
       batch = {ids: [], vectors: [], payloads: []}
@@ -49,7 +49,7 @@ module Langchain::Vectorsearch
       Array(texts).each_with_index do |text, i|
         id = ids[i] || SecureRandom.uuid
         batch[:ids].push(id)
-        batch[:vectors].push(llm.embed(text: text))
+        batch[:vectors].push(llm.embed(text: text).embedding)
         batch[:payloads].push({content: text})
       end
 
@@ -95,7 +95,7 @@ module Langchain::Vectorsearch
       query:,
       k: 4
     )
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(
         embedding: embedding,
@@ -104,7 +104,7 @@ module Langchain::Vectorsearch
     end
 
     # Search for similar texts by embedding
-    # @param embedding [Array] The embedding to search for
+    # @param embedding [Array<Float>] The embedding to search for
     # @param k [Integer] The number of results to return
     # @return [Hash] The response from the server
     def similarity_search_by_vector(

data/lib/langchain/vectorsearch/weaviate.rb

@@ -32,7 +32,7 @@ module Langchain::Vectorsearch
     end
 
     # Add a list of texts to the index
-    # @param texts [Array] The list of texts to add
+    # @param texts [Array<String>] The list of texts to add
     # @return [Hash] The response from the server
     def add_texts(texts:, ids: [])
       client.objects.batch_create(
@@ -41,7 +41,7 @@ module Langchain::Vectorsearch
     end
 
     # Update a list of texts in the index
-    # @param texts [Array] The list of texts to update
+    # @param texts [Array<String>] The list of texts to update
     # @return [Hash] The response from the server
     def update_texts(texts:, ids:)
       uuids = []
@@ -65,7 +65,7 @@ module Langchain::Vectorsearch
             __id: ids[i].to_s,
             content: text
           },
-          vector: llm.embed(text: text)
+          vector: llm.embed(text: text).embedding
         )
       end
     end
@@ -101,13 +101,13 @@ module Langchain::Vectorsearch
     # @param k [Integer|String] The number of results to return
     # @return [Hash] The search results
     def similarity_search(query:, k: 4)
-      embedding = llm.embed(text: query)
+      embedding = llm.embed(text: query).embedding
 
       similarity_search_by_vector(embedding: embedding, k: k)
     end
 
     # Return documents similar to the vector
-    # @param embedding [Array] The vector to search for
+    # @param embedding [Array<Float>] The vector to search for
     # @param k [Integer|String] The number of results to return
     # @return [Hash] The search results
     def similarity_search_by_vector(embedding:, k: 4)
@@ -154,7 +154,7 @@ module Langchain::Vectorsearch
           __id: id.to_s,
           content: text
         },
-        vector: llm.embed(text: text)
+        vector: llm.embed(text: text).embedding
       }
     end
   end

data/lib/langchain/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Langchain
-  VERSION = "0.6.17"
+  VERSION = "0.6.18"
 end

data/lib/langchain.rb

@@ -8,6 +8,7 @@ loader = Zeitwerk::Loader.for_gem
 loader.ignore("#{__dir__}/langchainrb.rb")
 loader.inflector.inflect(
   "ai21" => "AI21",
+  "ai21_response" => "AI21Response",
   "ai21_validator" => "AI21Validator",
   "csv" => "CSV",
   "html" => "HTML",
@@ -16,10 +17,12 @@ loader.inflector.inflect(
   "llm" => "LLM",
   "openai" => "OpenAI",
   "openai_validator" => "OpenAIValidator",
+  "openai_response" => "OpenAIResponse",
   "pdf" => "PDF",
   "react_agent" => "ReActAgent",
   "sql_query_agent" => "SQLQueryAgent"
 )
+loader.collapse("#{__dir__}/langchain/llm/response")
 loader.setup
 
 # Langchain.rb a is library for building LLM-backed Ruby applications. It is an abstraction layer that sits on top of the emerging AI-related tools that makes it easy for developers to consume and string those services together.

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: langchainrb
 version: !ruby/object:Gem::Version
-  version: 0.6.17
+  version: 0.6.18
 platform: ruby
 authors:
 - Andrei Bondarev
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-10-10 00:00:00.000000000 Z
+date: 2023-10-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: baran
@@ -532,6 +532,7 @@ files:
 - lib/langchain/agent/sql_query_agent.rb
 - lib/langchain/agent/sql_query_agent/sql_query_agent_answer_prompt.yaml
 - lib/langchain/agent/sql_query_agent/sql_query_agent_sql_prompt.yaml
+- lib/langchain/chunk.rb
 - lib/langchain/chunker/base.rb
 - lib/langchain/chunker/prompts/semantic_prompt_template.yml
 - lib/langchain/chunker/recursive_text.rb
@@ -558,6 +559,15 @@ files:
 - lib/langchain/llm/openai.rb
 - lib/langchain/llm/prompts/summarize_template.yaml
 - lib/langchain/llm/replicate.rb
+- lib/langchain/llm/response/ai21_response.rb
+- lib/langchain/llm/response/anthropic_response.rb
+- lib/langchain/llm/response/base_response.rb
+- lib/langchain/llm/response/cohere_response.rb
+- lib/langchain/llm/response/google_palm_response.rb
+- lib/langchain/llm/response/hugging_face_response.rb
+- lib/langchain/llm/response/ollama_response.rb
+- lib/langchain/llm/response/openai_response.rb
+- lib/langchain/llm/response/replicate_response.rb
 - lib/langchain/loader.rb
 - lib/langchain/output_parsers/base.rb
 - lib/langchain/output_parsers/output_fixing_parser.rb