simple_rag 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 3152b26c0832a438b14e23a4c871e93d946b40bc6ea8bf974a415d1ec47baa6a
+  data.tar.gz: 71caac3f3e0c7549fc4d99e4eb2ad13791004537ee66abb1efee72c82dc04d25
+SHA512:
+  metadata.gz: 807cad84683e4ac9079b07e2d610c788b3badcb7acbc9b7eabc96a293f0f44f76422f4f1ec3549ff568b02586c3f62942a421122add4034bed8ff355ec81d1bd
+  data.tar.gz: 9186c52f2f19ae30f4a9d1ec1be15ed40d558d584662b62199731cec18be257969fabb37a911e362cd1204f3d446cfb7eeb238b509b047a90d67da2e472778de

data/exe/simple_rag

@@ -0,0 +1,5 @@
+#!/usr/bin/env ruby
+
+require "simple_rag"
+
+SimpleRag::CLI.start(ARGV)

data/lib/simple_rag/cli.rb

@@ -0,0 +1,8 @@
+module SimpleRag
+  class CLI
+    def self.start(args)
+      puts "Hello from SimpleRag!"
+      SimpleRag::Engine.new.run
+    end
+  end
+end

data/lib/simple_rag/embed.rb

@@ -0,0 +1,13 @@
+module SimpleRag
+  class Embed
+    def self.embed_text(client, input)
+      embeddings_batch_response = client.embeddings({model: "mistral-embed", input: input})
+      embeddings_batch_response.dig("data", 0, "embedding")
+    end
+
+    def self.embed_chunks(client, chunks)
+      text_embeddings = chunks.map { |chunk| embed_text(client, chunk) }
+      Numo::DFloat[*text_embeddings]
+    end
+  end
+end

data/lib/simple_rag/generate.rb

@@ -0,0 +1,15 @@
+module SimpleRag
+  class Generate
+    def prompt(query, retrieved_chunks)
+      prompt = <<~PROMPT
+        Context information is below.
+        ---------------------
+        #{retrieved_chunks.join("\n---------------------\n")}
+        ---------------------
+        Given the context information and not prior knowledge, answer the query.
+        Query: #{query}
+        Answer:
+      PROMPT
+    end
+  end
+end

data/lib/simple_rag/index.rb

@@ -0,0 +1,32 @@
+module SimpleRag
+  class Index
+    def initialize(client)
+      @text = nil
+      @client = client
+    end
+
+    def load(url)
+      response = HTTParty.get(url)
+      text = response.body
+      File.write("data/essay.txt", text)
+      @text = text
+      text
+    end
+
+    def chunk(text)
+      chunk_size = 2048
+      text.chars.each_slice(chunk_size).map(&:join)
+    end
+
+    def embed_chunks(chunks)
+      SimpleRag::Embed.embed_chunks(@client, chunks)
+    end
+
+    def save(text_embeddings)
+      d = text_embeddings.shape[1]
+      index = Faiss::IndexFlatL2.new(d)
+      index.add(text_embeddings)
+      index
+    end
+  end
+end

data/lib/simple_rag/retrieve.rb

@@ -0,0 +1,29 @@
+module SimpleRag
+  class Retrieve
+    attr_accessor :chunks
+    def initialize(client)
+      @client = client
+      # @chunks = nil
+      @index = nil
+    end
+
+    def save_chunks(chunks)
+      @chunks = chunks
+    end
+
+    def save_index(index)
+      @index = index
+    end
+
+    def embed_query(query)
+      query_embedding = SimpleRag::Embed.embed_text(@client, query)
+      question_embeddings = Numo::DFloat[query_embedding]
+    end
+
+    def similarity_search(question_embeddings, k_neighbors_count)
+      distances, indices = @index.search(question_embeddings, k_neighbors_count)
+      index_array = indices.to_a[0]
+      index_array.map { |i| @chunks[i] }
+    end
+  end
+end

data/lib/simple_rag/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module SimpleRag
+  VERSION = "0.1.0"
+end

data/lib/simple_rag.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+require "httparty"
+require "numo/narray"
+require "faiss"
+require "matrix"
+require "io/console"
+require "mistral-ai"
+require "zeitwerk"
+require "dotenv/load"
+require "byebug"
+require_relative "simple_rag/version"
+require_relative "simple_rag/cli"
+
+loader = Zeitwerk::Loader.for_gem
+loader.setup
+
+module SimpleRag
+  class Error < StandardError; end
+
+  class Engine
+    DEFAULT_URL = "https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt"
+
+    def run_mistral(client, user_message, model: "mistral-medium-latest")
+      messages = [{role: "user", content: user_message}]
+      chat_response = client.chat_completions({model: model, messages: messages})
+      chat_response.dig("choices", 0, "message", "content")
+    end
+
+    def prompt_user_for_url
+      print "Specify a URL to an HTML document you would like to ask questions of (Default: What I Worked On by Paul Graham): "
+      input_url = gets.chomp
+      input_url.empty? ? DEFAULT_URL : input_url
+    end
+
+    def valid_url?(url)
+      uri = URI.parse(url)
+      uri.is_a?(URI::HTTP) || uri.is_a?(URI::HTTPS)
+    rescue URI::InvalidURIError
+      false
+    end
+
+    def get_url
+      url = prompt_user_for_url
+      until valid_url?(url)
+        puts "The URL provided is invalid. Please try again."
+        url = prompt_user_for_url
+      end
+      url
+    end
+
+    def run
+      url = get_url
+      puts "Document Downloaded"
+
+      # Setup LLM of choice
+      api_key = ENV["MISTRAL_AI_KEY"] || STDIN.getpass("Type your API Key: ")
+      raise "Missing API Key" unless api_key
+
+      client = Mistral.new(
+        credentials: {api_key: api_key},
+        options: {server_sent_events: true}
+      )
+
+      # Indexing
+      puts "Initialize indexing"
+      index_instance = SimpleRag::Index.new(client)
+      puts "Loading url"
+      text = index_instance.load(url)
+      puts "Chunk text"
+      chunks = index_instance.chunk(text)
+      puts "Embed chunks"
+      text_embeddings = index_instance.embed_chunks(chunks)
+      index = index_instance.save(text_embeddings)
+
+      retrieve_instance = SimpleRag::Retrieve.new(client)
+      retrieve_instance.save_index(index)
+      retrieve_instance.save_chunks(chunks)
+
+      loop do
+        print "Enter your query (or type 'exit' to quit): "
+        query = gets.chomp
+        break if query.downcase == "exit"
+        puts
+
+        # Retrieval/Search
+        question_embedding = retrieve_instance.embed_query(query)
+        retrieved_chunks = retrieve_instance.similarity_search(question_embedding, 2)
+
+        # Generation
+        prompt = SimpleRag::Generate.new.prompt(query, retrieved_chunks)
+
+        puts run_mistral(client, prompt)
+        puts
+      end
+    end
+  end
+end

metadata

@@ -0,0 +1,52 @@
+--- !ruby/object:Gem::Specification
+name: simple_rag
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Landon Gray
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2024-08-12 00:00:00.000000000 Z
+dependencies: []
+description:
+email:
+- landon.gray@hey.com
+executables:
+- simple_rag
+extensions: []
+extra_rdoc_files: []
+files:
+- exe/simple_rag
+- lib/simple_rag.rb
+- lib/simple_rag/cli.rb
+- lib/simple_rag/embed.rb
+- lib/simple_rag/generate.rb
+- lib/simple_rag/index.rb
+- lib/simple_rag/retrieve.rb
+- lib/simple_rag/version.rb
+homepage:
+licenses: []
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.3.7
+signing_key:
+specification_version: 4
+summary: Simple Rag is a lightweight library that transforms any Ruby project into
+  a simple RAG application.
+test_files: []