Module: OllamaChat::RAGHandling

Included in:
Chat
Defined in:
lib/ollama_chat/rag_handling.rb

Overview

A module that handles Retrieval-Augmented Generation (RAG) operations for the OllamaChat application.

This module provides functionality for managing collections, clearing and changing collections, listing collections, and renaming collections within the RAG system.

Instance Method Summary collapse

Instance Method Details

#all_collections ⇒ Sequel::Dataset (private)

Retrieves all document collections registered in the database.

Returns:

  • (Sequel::Dataset) —

    a dataset of all collections ordered by name.



152
153
154
 
# File 'lib/ollama_chat/rag_handling.rb', line 152

def all_collections
  models::Collection.order(:name)
end

#choose_collection(current_collection) ⇒ Object (private)

The choose_collection method presents a menu to select or create a document collection. It displays existing collections along with options to create a new one or exit. The method prompts the user for input and updates the document collection accordingly.

Parameters:

  • current_collection (String, nil) —

    the name of the currently active collection



93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
 
# File 'lib/ollama_chat/rag_handling.rb', line 93

def choose_collection(current_collection)
  collections = [ current_collection ] + all_collections.pluck(:name)
  collections = collections.filter_map(&:to_s).uniq.sort
  collections.unshift('[EXIT]').unshift('[NEW]')
  collection = choose_entry(
    collections,
    prompt: 'Which archive of knowledge shall we delve into? %s'
  ) || current_collection
  case collection = collection&.to_s
  when '[NEW]'
    if name = create_collection
      @documents.collection = name
    end
  when nil, '[EXIT]'
    STDOUT.puts "Exiting chooser."
  when /./
    @documents.collection = collection
  end
ensure
  if collection
    @session.update(current_collection: collection)
    log(:info, "Collection switched", data: { collection: })
    STDOUT.puts "Using collection #{bold{collection}}."
  end
  info
end

#clear_collection ⇒ Object (private)

Clears documents from the collection through an interactive user interface.

This method allows users to selectively clear documents by choosing specific tags from the current collection or to clear all documents in the collection. It provides a loop for multiple deletions until the user exits or completes a clear operation.



50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
 
# File 'lib/ollama_chat/rag_handling.rb', line 50

def clear_collection
  choose_with_state do
    loop do
      tags = @documents.tags.to_a.unshift('[ALL]').unshift('[EXIT]')
      tag = choose_entry(
        tags,
        prompt: 'What obsolete records are to be excised from the annals? %s'
      )
      case tag
      when nil, '[EXIT]'
        STDOUT.puts "Exiting chooser."
        break
      when '[ALL]'
        if confirm?(prompt: '🔔 Are you sure? (y/n) ', yes: /\Ay/i)
          @documents.clear
          log(:info, "Collection cleared", data: { collection: })
          STDOUT.puts "Cleared collection #{bold{collection}}."
          break
        end
      when /./
        @documents.clear(tags: [ tag ])
        log(:info, "Tag cleared from collection", data: { collection:, tag: })
        STDOUT.puts "Cleared tag #{tag} from collection #{bold{collection}}."
      end
    end
  end
end

#collection ⇒ String, Symbol (private)

Returns the name of the currently active document collection.

Returns:

  • (String, Symbol) —

    the name of the current collection



40
41
42
 
# File 'lib/ollama_chat/rag_handling.rb', line 40

def collection
  @documents.collection
end

#create_collection ⇒ Object (private)

Interactively create a new collection record.



234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
 
# File 'lib/ollama_chat/rag_handling.rb', line 234

def create_collection
  name = switch_history(:collection_name) do
    ask?(prompt: "📚 Name of the new collection: ")
  end
  unless name.full?
    STDERR.puts "❌ Cancelled creation of collection."
    return
  end

  if database_collection?(name)
    STDERR.puts "❌ Collection #{name.inspect} already exists."
    return
  end

  description = switch_history(:collection_description) do
    ask?(prompt: "📝 Description: ")
  end
  unless description.full?
    STDERR.puts "❌ Cancelled creation of collection #{name.inspect}."
    return
  end
  patterns_str = switch_history(:patterns) do
    ask?(prompt: "🔍 Patterns (space-separated globs, e.g., lib/**/*.rb): ")
  end
  patterns = extract_patterns(patterns_str)

  begin
    models::Collection.create(
      name: name.to_s,
      description: description.to_s,
      patterns:
    )
    if patterns.full?
      update_collection(name.to_s)
    end
    STDOUT.puts "✅ Created collection '#{name}'."
    log(:info, "Collection created", data: { name: name, description:, patterns: })
  rescue Sequel::UniqueConstraintViolation
    STDERR.puts "❌ Collection #{name.inspect} already exists."
  rescue Sequel::Error => e
    STDERR.puts "❌ Database error: #{e.message}"
  end
  name.to_s
end

#database_collection?(collection) ⇒ OllamaChat::Database::Models::Collection? (private)

Looks up a collection in the database by name.

Parameters:

  • collection (String, Symbol, #to_s) —

    the collection name to look up

Returns:



33
34
35
 
# File 'lib/ollama_chat/rag_handling.rb', line 33

def database_collection?(collection)
  models::Collection[name: collection.to_s]
end

#delete_collection ⇒ Object (private)

Permanently remove a collection from DB and purge from Documentrix.



322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
 
# File 'lib/ollama_chat/rag_handling.rb', line 322

def delete_collection
  choose_with_state do
    loop do
      collections = models::Collection.order(:name).pluck(:name)
      collections.unshift('[CANCEL]')
      target_name = choose_entry(
        collections,
        prompt: '🗑️ Which collection to delete? %s'
      )
      break if target_name.nil? || target_name == '[CANCEL]'

      col = database_collection?(target_name)
      unless col
        STDERR.puts "❌ Collection #{target_name.inspect} not found in database."
        next
      end

      if confirm?(prompt: "⚠️ Are you sure you want to permanently delete #{target_name.inspect}? (y/n) ", yes: /\Ay/i)
        begin
          col.chat = self
          col.destroy
          STDOUT.puts "✅ Deleted collection #{target_name.inspect}."
          log(:info, "Collection deleted", data: { name: target_name })
        rescue Sequel::Error => e
          STDERR.puts "❌ Database error: #{e.message}"
        rescue => e
          STDERR.puts "❌ Error removing from Documentrix: #{e.message}"
        end
      else
        STDOUT.puts "🚫 Deletion cancelled."
      end
    end
  end
end

#edit_collection ⇒ Object (private)

Interactively update attributes of an existing collection.



280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
 
# File 'lib/ollama_chat/rag_handling.rb', line 280

def edit_collection
  collections = models::Collection.order(:name).pluck(:name)
  collections.unshift('[CANCEL]')
  target_name = choose_entry(
    collections,
    prompt: '📚 Which collection to edit? %s'
  )
  return if target_name.nil? || target_name == '[CANCEL]'

  col = database_collection?(target_name)
  unless col
    STDERR.puts "❌ Collection #{target_name.inspect} not found in database."
    return
  end

  new_description = switch_history :collection_description do
    ask?(
      prompt: "📝 New description (leave blank to keep): ",
      prefill: col.description
    )
  end
  patterns = switch_history :patterns do
    patterns_str = ask?(
      prompt: "🔍 New patterns (space-separated, leave blank to keep): ",
      prefill: col.patterns.full?(:join, ' ')
    )
    extract_patterns(patterns_str)
  end

  col.description = new_description.full? ? new_description.to_s : col.description
  col.patterns    = patterns

  begin
    col.save
    STDOUT.puts "✅ Updated collection '#{col.name}'."
    log(:info, "Collection updated", data: { name: col.name })
  rescue Sequel::Error => e
    STDERR.puts "❌ Database error: #{e.message}"
  end
end

#extract_patterns(patterns_str) ⇒ Array<String> (private)

Extracts and normalizes file patterns from a shell-style string.

Splits the input string using shell semantics, correctly handling spaces enclosed in single/double quotes or escaped with backslashes. Each resulting pattern is expanded to an absolute path.

Parameters:

  • patterns_str (String, nil) —

    the shell-style glob patterns

Returns:

  • (Array<String>) —

    an array of expanded absolute path patterns



228
229
230
231
 
# File 'lib/ollama_chat/rag_handling.rb', line 228

def extract_patterns(patterns_str)
  patterns = Shellwords.split(patterns_str.to_s)
  patterns.map { File.expand_path(_1) }
end

#list_collections ⇒ Object (private)

Displays the list of available collections in the terminal.

This method retrieves the current collection and the full list of available collections and their descriptions from the internal document handler, highlighting the active one.



161
162
163
164
165
166
167
168
169
170
171
 
# File 'lib/ollama_chat/rag_handling.rb', line 161

def list_collections
  current_collection = collection.to_s
  collections = all_collections.select(:name, :description)
  use_pager do |output|
    collections.each { |c|
      collection_name = current_collection == c.name ? bold { c.name } : c.name
      collection_description = c.description
      output.puts '%s: %s' % [ collection_name, collection_description ]
    }
  end
end

#rename_collection(current_collection) ⇒ Object (private)

Rename an existing collection to a new, user‑supplied name.

This helper prompts the user to provide a new name for the collection identified by current_collection. It then renames the current collection to have the new_name and switches to it.

Parameters:

  • current_collection (Symbol) —

    the current collection name



127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
 
# File 'lib/ollama_chat/rag_handling.rb', line 127

def rename_collection(current_collection)
  prompt = 'Rename collection %s to: ' % current_collection
  new_collection = switch_history :collection_name do
    ask?(prompt:, prefill: current_collection).full?(:to_sym)
  end
  if new_collection
    begin
      @documents.rename_collection(new_collection)
      col = database_collection?(current_collection)
      col&.update(name: new_collection.to_s)
      log(:info, "Collection renamed", data: { old_name: current_collection, new_name: new_collection })
      STDOUT.puts "Renamed current collection #{current_collection} to #{new_collection}."
    rescue Sequel::UniqueConstraintViolation
      STDERR.puts "❌ Renaming to #{new_collection} failed, it already exists in database."
    rescue => e
      STDERR.puts "❌ Renaming to #{new_collection} failed: #{e.message}"
    end
  else
    STDOUT.puts "Renaming cancelled."
  end
end

#set_current_collection(collection) ⇒ String, Symbol (private)

Sets the current document collection.

Parameters:

  • collection (String, Symbol) —

    the name of the collection to set

Returns:

  • (String, Symbol) —

    the newly set collection name



82
83
84
 
# File 'lib/ollama_chat/rag_handling.rb', line 82

def set_current_collection(collection)
  @documents.collection = collection
end

#switch_collection(other_collection = nil) { ... } ⇒ Object

Temporarily switches the RAG collection to the specified collection.

The current collection is stored and restored after the block is executed, ensuring the state remains consistent regardless of whether the block completes successfully or raises an error.

Parameters:

  • other_collection (String, nil) (defaults to: nil) —

    the collection to switch to. If nil, the current collection is used.

Yields:

  • The code to execute within the context of the switched collection.

Returns:

  • (Object) —

    the result of the block.



18
19
20
21
22
23
24
 
# File 'lib/ollama_chat/rag_handling.rb', line 18

def switch_collection(other_collection = nil)
  other_collection ||= collection
  old_collection, @documents.collection = collection, other_collection
  yield other_collection
ensure
  @documents.collection = old_collection
end

#update_collection(collection) ⇒ String (private)

Updates the documents in the current collection by re-embedding any sources that have been modified since they were first added, and embedding any new files matching the collection's patterns.

This method iterates through all records in the active collection and identifies unique sources. For each modified source, it preserves the existing tags, removes the stale records, and re-embeds the current version of the source. It then scans for new files matching the collection's patterns and embeds them.

Returns:

  • (String) —

    a newline-separated string of embedding result messages.



184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
 
# File 'lib/ollama_chat/rag_handling.rb', line 184

def update_collection(collection)
  switch_collection(collection) do
    unless col = database_collection?(collection)
      STDERR.puts "❌ Collection #{collection.inspect} not found in database."
      return
    end
    results = []
    seen = {}
    @documents.each_record do |record|
      source = @documents.normalize_source(record.source) or next
      seen.key?(source) and next
      seen[source] = true
      unless @documents.source_modified?(source)
        infobar.puts "Source #{source.to_s.inspect} is unmodified. => Skipping."
        next
      end
      tags = record.tags_set
      @documents.source_remove(source)
      r = embed(source, tags:) or next
      results << r
    end

    if patterns = col.patterns.full?
      all_file_set(patterns).each do |file|
        seen[file.to_s] and next
        seen[file.to_s] = true
        r = embed(file.to_s, tags: []) or next
        results << r
      end
    end

    log(:info, "Collection updated", data: { collection:, sources_updated: results.size })
    results * "\n"
  end
end