diff --git a/refresh_sources.lua b/refresh_sources.lua index 8b3c95e..f39b259 100644 --- a/refresh_sources.lua +++ b/refresh_sources.lua @@ -3,6 +3,8 @@ package.path = (arg[0]:match("@?(.*/)") or arg[0]:match("@?(.*\\)")) .. "lib" .. package.config:sub(1, 1) .. "?.lua;" .. package.path local utility = require "utility" +local json = utility.require("dkjson") +local text_processing = utility.require("text_processing") local timing = utility.require("timing") -- TODO make utility have a function for getting/setting defaults where locking is only used to set defaults if they aren't present @@ -17,7 +19,9 @@ if not config.models then utility.save_config() end -local refresht_file_list = function(data_source) +-- TODO recognize uninitialized source and initialize it +-- this would need to be in the config +local refresh_file_list = function(data_source) timing.mark("Assembling file list.") os.execute("cd " .. ("PRIVATE_DATA" .. utility.path_separator .. data_source.path):enquote() .. " && " .. data_source.refresh_command) @@ -32,9 +36,11 @@ local refresht_file_list = function(data_source) file_list[#file_list + 1] = file_name end) + timing.mark("Finished assembling file list.") return file_list end +-- returns nothing when too much text is sent local generate_embeddings = function(text) if #text > config.models.embedding.max_chunk_size then return nil, "generate_embeddings() must only be passed appropriately-sized chunks!" @@ -42,21 +48,68 @@ local generate_embeddings = function(text) local result = utility.llm_prompt(text, config.models.embedding.model) - -- skipping unnecessary decoding to an array and putting back into text for saving - -- this would cause problems if we intended to use the results immediately - return setmetatable({}, { + result = setmetatable({}, { __tojson = function() return result end, }) + + return json.decode(result) end -local process_file = function(file_name) +-- returns nothing for empty files +local process_file = function(data_source, file_name) local text = utility.read_file(file_name) - if w + + if data_source.strip_frontmatter then + text = text_processing.strip_frontmatter(text) + end + + if #text == 0 then + print(file_name .. "\n is empty and being skipped.") + return + end + + local chunk_size = config.models.embedding.max_chunk_size + local half_chunk_size = math.floor(chunk_size / 2) + local chunks = { text } + + while #text > chunk_size do + local first_chunk = text:sub(1, chunk_size) + local overlap_chunk = text:sub(half_chunk_size, chunk_size + half_chunk_size) + + chunks[#chunks + 1] = first_chunk + chunks[#chunks + 1] = overlap_chunk + + text = text:sub(chunk_size) + if #text > half_chunk_size and (not #text > chunk_size) then + -- last chunk would be skipped if we didn't handle this here + chunks[#chunks + 1] = text + end + end + + local embeddings = {} + for i = 1, #chunks do + embeddings[i] = generate_embeddings(chunks[i]) + end + + if embeddings[1] == nil then + -- TODO average all other embeddings to this one + end + + return chunks, embeddings end -local to_be_named_generate_file_embeddings = function(text) - if #text > config.models.embedding.max_chunk_size then - -- TODO, split, process each chunk, return both a total and each chunk +-- TODO somewhere near the top of this function chain needs to be the ability to recognize a file as unchanged and skip re-generating embeddings +-- this means running the same shasum command I will put below on that source file before trying to generate an embedding +-- something something merging new changes into the old data structure means we need to load that data structure first so we only replace modified chunks +local refresh_sources = function() + local sources = utility.load_data("PRIVATE_DATA/sources.json") + for _, data_source in pairs(sources) do + local file_list = refresh_file_list(data_source) + for _, file_name in ipairs(file_list) do + local chunks, embeddings = process_file(file_name) + -- TODO decide how these will be stored + -- for chunks, save to a specific local tmp file to shasum, then mv that file based on the sum + -- for embeddings, store sum = embedding ? (look at how generate_embeddings.lua worked) + end end - -- return just the file and its embedding end