smart_rag 0.1.0 → 0.2.1

data/examples/common.rb

@@ -29,7 +29,7 @@ module Examples
         llm: {
           provider: ENV["SMARTRAG_LLM_PROVIDER"] || ENV["LLM_PROVIDER"] || "openai",
           api_key: ENV["OPENAI_API_KEY"] || ENV["LLM_API_KEY"] || "ollama-local",
-          endpoint: ENV["LLM_ENDPOINT"] || "http://localhost:11434/v1/chat/completions",
+          endpoint: ENV["LLM_ENDPOINT"] || "http://192.168.1.48:11434/v1/chat/completions",
           model: ENV["LLM_MODEL"] || "qwen3",
         },
       }

data/exe/smart-rag-db

@@ -0,0 +1,163 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# SmartRAG database bootstrap CLI.
+#
+# Usage:
+#   smart-rag-db create    # create the database
+#   smart-rag-db drop      # drop the database
+#   smart-rag-db migrate   # enable pg extensions + run migrations
+#   smart-rag-db seed      # load full-text search seed data
+#   smart-rag-db reset     # drop + create + migrate + seed
+
+require "optparse"
+
+COMMANDS = %w[create drop migrate seed reset].freeze
+
+command = ARGV.shift
+abort "Usage: smart-rag-db {#{COMMANDS.join('|')}}" unless COMMANDS.include?(command)
+
+require "smart_rag"
+
+# --- helpers ---------------------------------------------------------------
+
+def load_db_config
+  SmartRAG::Config.load[:database]
+rescue StandardError => e
+  warn "Warning: could not load config (#{e.message}); falling back to SMARTRAG_DB_* env vars."
+  {
+    adapter: "postgresql",
+    host: ENV["SMARTRAG_DB_HOST"] || "localhost",
+    port: ENV["SMARTRAG_DB_PORT"] || 5432,
+    database: ENV["SMARTRAG_DB_NAME"] || "smart_rag_development",
+    username: ENV["SMARTRAG_DB_USER"] || "postgres",
+    password: ENV["SMARTRAG_DB_PASSWORD"],
+    pool: 5,
+    encoding: "UTF8",
+    timeout: 5000,
+  }
+end
+
+def windows?
+  /mswin|mingw|cygwin/i.match?(RbConfig::CONFIG["host_os"])
+end
+
+def running_as_root?
+  Process.respond_to?(:uid) && Process.uid == 0
+end
+
+def connect(config, database: nil)
+  require "sequel"
+  cfg = config.reject { |key, _| key == :extensions || key.to_s == "extensions" }
+  cfg = cfg.merge(database: database) if database
+  Sequel.connect(cfg)
+end
+
+def migrations_dir
+  File.expand_path("../db/migrations", __dir__)
+end
+
+def seeds_file
+  File.expand_path("../db/seeds/text_search_configs.sql", __dir__)
+end
+
+def print_hints(command)
+  warn "You may need to:"
+  warn "  1. Ensure PostgreSQL is running"
+  warn "  2. Set SMARTRAG_DB_* environment variables (or create a .env file)"
+  warn "  3. Use `sudo -u postgres smart-rag-db #{command}` when running as root"
+end
+
+# --- commands --------------------------------------------------------------
+
+def create_database(config)
+  database = config[:database].to_s
+  host = (config[:host] || "localhost").to_s
+  port = (config[:port] || 5432).to_s
+  username = (config[:username] || "postgres").to_s
+
+  created = false
+  if host == "localhost" && running_as_root? && !windows?
+    created = system("sudo", "-u", username, "createdb", "-h", host, "-p", port, database)
+  end
+
+  return true if created
+
+  db = connect(config, database: "postgres")
+  db.execute("CREATE DATABASE #{database}")
+  db.disconnect
+  true
+end
+
+def drop_database(config)
+  database = config[:database].to_s
+  host = (config[:host] || "localhost").to_s
+  port = (config[:port] || 5432).to_s
+  username = (config[:username] || "postgres").to_s
+
+  dropped = false
+  if host == "localhost" && running_as_root? && !windows?
+    dropped = system("sudo", "-u", username, "dropdb", "-h", host, "-p", port, "--if-exists", database)
+  end
+
+  return true if dropped
+
+  db = connect(config, database: "postgres")
+  db.execute("DROP DATABASE IF EXISTS #{database}")
+  db.disconnect
+  true
+end
+
+def migrate_database(config)
+  require "sequel"
+  Sequel.extension :migration
+
+  db = connect(config)
+  db.run "CREATE EXTENSION IF NOT EXISTS vector"
+  begin
+    db.run "CREATE EXTENSION IF NOT EXISTS pg_jieba"
+  rescue Sequel::DatabaseError => e
+    puts "Warning: pg_jieba extension not available (#{e.message}); Chinese full-text will fall back to 'simple'"
+  end
+
+  Sequel::Migrator.run(db, migrations_dir)
+  puts "Migrations completed"
+  db.disconnect
+end
+
+def seed_database(config)
+  abort "Seeds file not found: #{seeds_file}" unless File.exist?(seeds_file)
+
+  db = connect(config)
+  db.run(File.read(seeds_file))
+  puts "Database seeded"
+  db.disconnect
+end
+
+# --- dispatch ---------------------------------------------------------------
+
+begin
+  config = load_db_config
+
+  case command
+  when "create"
+    create_database(config)
+    puts "Database #{config[:database]} created"
+  when "drop"
+    drop_database(config)
+    puts "Database #{config[:database]} dropped"
+  when "migrate"
+    migrate_database(config)
+  when "seed"
+    seed_database(config)
+  when "reset"
+    drop_database(config)
+    create_database(config)
+    migrate_database(config)
+    seed_database(config)
+  end
+rescue StandardError => e
+  warn "smart-rag-db #{command} failed: #{e.message}"
+  print_hints(command)
+  exit 1
+end

data/exe/smart-rag-media-worker

@@ -0,0 +1,34 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require 'optparse'
+require 'smart_rag'
+
+options = { batch_size: 10, interval: 2.0, once: false }
+OptionParser.new do |parser|
+  parser.banner = 'Usage: smart-rag-media-worker [options]'
+  parser.on('--once', 'Process one batch and exit') { options[:once] = true }
+  parser.on('--batch-size N', Integer, 'Jobs per batch (default: 10)') { |value| options[:batch_size] = value }
+  parser.on('--interval SECONDS', Float, 'Idle polling interval (default: 2)') { |value| options[:interval] = value }
+end.parse!
+
+config = SmartRAG::Config.load(ENV['SMARTRAG_CONFIG_PATH'])
+rag = SmartRAG::SmartRAG.new(config)
+rag.recover_media_jobs
+last_pruned_at = Time.at(0)
+last_recovered_at = Time.now
+
+loop do
+  jobs = rag.run_media_jobs(limit: options[:batch_size])
+  if Time.now - last_recovered_at >= 60
+    rag.recover_media_jobs
+    last_recovered_at = Time.now
+  end
+  if Time.now - last_pruned_at >= 3600
+    rag.prune_media_jobs
+    rag.garbage_collect_media_objects
+    last_pruned_at = Time.now
+  end
+  break if options[:once]
+  sleep(options[:interval]) if jobs.empty?
+end

data/lib/smart_rag/config.rb

@@ -8,6 +8,8 @@ module SmartRAG
         # If file_path is a Hash, return it directly (already a config hash)
         return symbolize_keys(file_path) if file_path.is_a?(Hash)
 
+        load_dotenv
+
         file_path ||= default_config_path
 
         unless File.exist?(file_path)
@@ -73,6 +75,16 @@ module SmartRAG
         symbolize_keys(config) if config.is_a?(Hash)
       end
 
+      # Load .env from the current directory so ERB `<%= ENV['...'] %>`
+      # templates and SMARTRAG_* variables work consistently. Falls back to
+      # real environment variables when the dotenv gem is unavailable.
+      def load_dotenv
+        require "dotenv"
+        Dotenv.load
+      rescue LoadError
+        # dotenv is optional; plain environment variables still work
+      end
+
       private
 
       def default_config_path

data/lib/smart_rag/core/document_processor.rb

@@ -3,6 +3,8 @@ require 'net/http'
 require 'fileutils'
 require 'tempfile'
 require 'digest'
+require 'timeout'
+require_relative 'media_safety_policy'
 require_relative '../../smart_rag'
 require_relative '../models'
 require_relative '../chunker/markdown_chunker'
@@ -149,7 +151,15 @@ module SmartRAG
       # @param [String] url Source URL
       # @param [Hash] options Download options
       # @return [String] Path to downloaded file
-      def download_from_url(url, options = {})
+      def download_from_url(url, options = {}, redirect_count = 0)
+        max_redirects = options.fetch(:max_redirects, 5).to_i
+        raise ArgumentError, "too many redirects (max #{max_redirects})" if redirect_count > max_redirects
+
+        MediaSafetyPolicy.new(
+          max_bytes: options[:max_file_size] || 50 * 1024 * 1024,
+          allow_private_urls: options.fetch(:allow_private_urls, false),
+          allowed_hosts: options[:allowed_hosts]
+        ).validate_url!(url)
         uri = URI.parse(url)
         @logger.info "Downloading from URL: #{url}"
 
@@ -161,34 +171,59 @@ module SmartRAG
         temp_file.close
 
         # Download the file
-        Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == 'https') do |http|
+        timeout = options.fetch(:download_timeout, 30).to_i
+        max_bytes = options.fetch(:max_file_size, 50 * 1024 * 1024).to_i
+        redirect_url = nil
+        Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == 'https',
+                        open_timeout: timeout, read_timeout: timeout) do |http|
           request = Net::HTTP::Get.new(uri)
           # Set user agent to avoid being blocked
           request['User-Agent'] = 'SmartRAG Document Processor/1.0'
 
-          response = http.request(request)
-
-          case response.code
-          when '200'
-            File.write(temp_path, response.body)
-          when '301', '302', '303', '307', '308'
-            # Follow redirect
-            redirect_url = response['Location']
-            @logger.info "Redirecting to: #{redirect_url}"
-            return download_from_url(redirect_url, options)
-          else
-            raise "HTTP Error: #{response.code} - #{response.message}"
+          http.request(request) do |response|
+            case response.code
+            when '200'
+              content_length = response['Content-Length'].to_i
+              raise ArgumentError, "download exceeds #{max_bytes} bytes" if content_length > max_bytes
+
+              written = 0
+              File.open(temp_path, 'wb') do |file|
+                response.read_body do |chunk|
+                  written += chunk.bytesize
+                  raise ArgumentError, "download exceeds #{max_bytes} bytes" if written > max_bytes
+                  file.write(chunk)
+                end
+              end
+            when '301', '302', '303', '307', '308'
+              redirect_url = URI.join(url, response['Location']).to_s
+              @logger.info "Redirecting to: #{redirect_url}"
+            else
+              raise "HTTP Error: #{response.code} - #{response.message}"
+            end
           end
         end
 
+        if redirect_url
+          File.delete(temp_path) if File.exist?(temp_path)
+          return download_from_url(redirect_url, options, redirect_count + 1)
+        end
+
         @downloaded_file = temp_path
         @logger.info "Downloaded file to: #{temp_path}"
         temp_path
       rescue StandardError => e
+        File.delete(temp_path) if defined?(temp_path) && temp_path && File.exist?(temp_path)
         @logger.error "Download failed: #{e.message}"
         raise e
       end
 
+      def cleanup_downloaded_file
+        return unless @downloaded_file && File.exist?(@downloaded_file)
+
+        File.delete(@downloaded_file)
+        @downloaded_file = nil
+      end
+
       # Extract metadata from file
       # @param [String] file_path Path to file
       # @param [Hash] options Metadata options
@@ -235,7 +270,12 @@ module SmartRAG
         ext = File.extname(file_path).downcase
         if ['.md', '.markdown'].include?(ext)
           @logger.info "Detected markdown source; skipping conversion"
-          return File.read(file_path)
+          return read_text_file(file_path)
+        end
+
+        if ['.txt', '.text'].include?(ext)
+          @logger.info "Detected plain text source; skipping markitdown conversion"
+          return read_text_file(file_path)
         end
 
         # Use markitdown bridge for conversion
@@ -256,7 +296,7 @@ module SmartRAG
 
         retries = 0
         begin
-          markdown = bridge.convert(file_path)
+          markdown = normalize_text_content(bridge.convert(file_path))
 
           raise 'Conversion failed: empty result' if markdown.nil? || markdown.strip.empty?
 
@@ -281,6 +321,27 @@ module SmartRAG
         raise e
       end
 
+      def read_text_file(file_path)
+        normalize_text_content(File.binread(file_path))
+      rescue StandardError => e
+        @logger.error "Failed to read text file #{file_path}: #{e.message}"
+        raise e
+      end
+
+      def normalize_text_content(content)
+        return '' if content.nil?
+
+        normalized = content.is_a?(String) ? content.dup : content.to_s
+
+        begin
+          normalized = normalized.encode(Encoding::UTF_8, invalid: :replace, undef: :replace, replace: '')
+        rescue Encoding::UndefinedConversionError, Encoding::InvalidByteSequenceError
+          normalized = normalized.force_encoding(Encoding::UTF_8).scrub
+        end
+
+        normalized.scrub
+      end
+
       # Create or update document record
       # @param [String] source Original source
       # @param [Hash] metadata Document metadata
@@ -302,6 +363,7 @@ module SmartRAG
           source_type: source_type,
           source_uri: normalized_source_uri,
           content_hash: content_hash,
+          principal: options[:principal] || metadata[:principal] || 'system',
           metadata: metadata.to_json
         }
 
@@ -355,11 +417,13 @@ module SmartRAG
       # @option options [Boolean] :generate_tags Whether to generate tags for sections
       def save_sections(document, chunks, options = {})
         sections = chunks.each_with_index.map do |chunk, index|
+          chunk_metadata = chunk[:metadata] || {}
           {
             document_id: document.id,
             section_title: chunk[:title],
             section_number: index + 1,
             content: chunk[:content],
+            metadata: chunk_metadata.to_json,
             created_at: Time.now,
             updated_at: Time.now
           }

data/lib/smart_rag/core/local_content_store.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require 'digest'
+require 'fileutils'
+
+module SmartRAG
+  module Core
+    class LocalContentStore
+      def initialize(root:)
+        @root = File.expand_path(root)
+        FileUtils.mkdir_p(@root)
+      end
+
+      def put(file_path)
+        digest = Digest::SHA256.file(file_path).hexdigest
+        extension = File.extname(file_path).downcase
+        destination = File.join(root, digest[0, 2], digest[2, 2], "#{digest}#{extension}")
+        FileUtils.mkdir_p(File.dirname(destination))
+        copy_atomically(file_path, destination) unless File.exist?(destination)
+        { content_hash: digest, storage_uri: "file://#{destination}", stored_path: destination }
+      end
+
+      def delete(storage_uri)
+        path = storage_uri.to_s.delete_prefix('file://')
+        return false unless path.start_with?("#{root}#{File::SEPARATOR}") && File.file?(path)
+
+        File.delete(path)
+        true
+      end
+
+      def materialize(storage_uri)
+        path = storage_uri.to_s.delete_prefix('file://')
+        raise ArgumentError, 'storage URI does not belong to this local store' unless path.start_with?("#{root}#{File::SEPARATOR}")
+        raise ArgumentError, 'stored media object is missing' unless File.file?(path)
+        [path, false]
+      end
+
+      private
+
+      attr_reader :root
+
+      def copy_atomically(source, destination)
+        temporary = "#{destination}.#{Process.pid}.tmp"
+        FileUtils.cp(source, temporary)
+        File.rename(temporary, destination)
+      ensure
+        File.delete(temporary) if defined?(temporary) && File.exist?(temporary)
+      end
+    end
+  end
+end

data/lib/smart_rag/core/media_extractors.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+require 'base64'
+require 'json'
+require 'net/http'
+require 'open3'
+require 'timeout'
+require 'uri'
+
+module SmartRAG
+  module Core
+    module MediaExtractors
+      class OpenAICompatible
+        def initialize(base_url:, api_key:, vision_model: nil, transcription_model: nil,
+                       timeout_seconds: 120, language: nil)
+          @base_url = base_url.to_s.sub(%r{/+$}, '')
+          @api_key = api_key
+          @vision_model = vision_model
+          @transcription_model = transcription_model
+          @timeout_seconds = timeout_seconds.to_i
+          @language = language
+        end
+
+        def describe(file_path, _timestamp_ms = nil)
+          raise ArgumentError, 'vision_model is not configured' if @vision_model.to_s.empty?
+
+          mime = mime_type(file_path)
+          body = {
+            model: @vision_model,
+            messages: [{ role: 'user', content: [
+              { type: 'text', text: 'Describe this image precisely for retrieval. Include visible text and important objects.' },
+              { type: 'image_url', image_url: { url: "data:#{mime};base64,#{Base64.strict_encode64(File.binread(file_path))}" } }
+            ] }],
+            temperature: 0.1
+          }
+          parsed = post_json('/v1/chat/completions', body)
+          parsed.dig('choices', 0, 'message', 'content').to_s.strip
+        end
+
+        alias extract describe
+
+        def transcribe(file_path)
+          raise ArgumentError, 'transcription_model is not configured' if @transcription_model.to_s.empty?
+
+          boundary = "SmartRAG#{rand(1_000_000_000)}"
+          fields = { 'model' => @transcription_model, 'response_format' => 'verbose_json',
+                     'timestamp_granularities[]' => 'segment' }
+          fields['language'] = @language unless @language.to_s.empty?
+          body = multipart_body(boundary, fields, file_path)
+          response = request('/v1/audio/transcriptions', body, "multipart/form-data; boundary=#{boundary}")
+          JSON.parse(response.body)
+        end
+
+        private
+
+        def post_json(path, body)
+          response = request(path, JSON.generate(body), 'application/json')
+          JSON.parse(response.body)
+        end
+
+        def request(path, body, content_type)
+          uri = URI.parse("#{@base_url}#{path}")
+          request = Net::HTTP::Post.new(uri)
+          request['Authorization'] = "Bearer #{@api_key}" unless @api_key.to_s.empty?
+          request['Content-Type'] = content_type
+          request.body = body
+          response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == 'https',
+                                     open_timeout: @timeout_seconds, read_timeout: @timeout_seconds) do |http|
+            http.request(request)
+          end
+          raise "media model HTTP #{response.code}: #{response.body.to_s[0, 300]}" unless response.is_a?(Net::HTTPSuccess)
+          response
+        end
+
+        def multipart_body(boundary, fields, file_path)
+          body = +''
+          fields.each do |key, value|
+            body << "--#{boundary}\r\nContent-Disposition: form-data; name=\"#{key}\"\r\n\r\n#{value}\r\n"
+          end
+          body << "--#{boundary}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"#{File.basename(file_path)}\"\r\n"
+          body << "Content-Type: #{mime_type(file_path)}\r\n\r\n"
+          body << File.binread(file_path)
+          body << "\r\n--#{boundary}--\r\n"
+          body
+        end
+
+        def mime_type(path)
+          { '.jpg' => 'image/jpeg', '.jpeg' => 'image/jpeg', '.png' => 'image/png',
+            '.webp' => 'image/webp', '.wav' => 'audio/wav', '.mp3' => 'audio/mpeg',
+            '.m4a' => 'audio/mp4' }.fetch(File.extname(path).downcase, 'application/octet-stream')
+        end
+      end
+
+      class Tesseract
+        def initialize(language: nil, timeout_seconds: 60)
+          @language = language
+          @timeout_seconds = timeout_seconds.to_i
+        end
+
+        def available?
+          system('which', 'tesseract', out: File::NULL, err: File::NULL)
+        end
+
+        def extract(file_path)
+          raise 'tesseract is not installed' unless available?
+
+          args = ['tesseract', file_path, 'stdout']
+          args.concat(['-l', @language]) unless @language.to_s.empty?
+          output = nil
+          status = nil
+          Timeout.timeout(@timeout_seconds) { output, status = Open3.capture2e(*args) }
+          raise "tesseract failed: #{output.to_s.strip}" unless status.success?
+
+          output.to_s.strip
+        end
+      end
+
+      module Factory
+        module_function
+
+        def build(config)
+          media = config || {}
+          extractors = {}
+          if media[:openai]
+            adapter = OpenAICompatible.new(**media[:openai])
+            extractors[:image_describer] = adapter.method(:describe) if media.dig(:openai, :vision_model)
+            if media.dig(:openai, :transcription_model)
+              extractors[:audio_transcriber] = adapter.method(:transcribe)
+              extractors[:video_transcriber] = adapter.method(:transcribe)
+            end
+          end
+          if media[:ocr]&.fetch(:provider, nil).to_s == 'tesseract'
+            extractors[:ocr_extractor] = Tesseract.new(**media[:ocr].reject { |key, _| key == :provider })
+          end
+          extractors
+        end
+      end
+    end
+  end
+end