refactor: move embedding cache to context module

Move embedding cache functionality from client to context module to improve
separation of concerns and code organization. The cache is now maintained in
the context module which is more appropriate since it deals with context
management.

This change simplifies the client module by removing the embedding cache
responsibility while maintaining the same caching functionality in a more
logical location.

Signed-off-by: Tomas Slusny <[email protected]>

Author: deathbeam

lua/CopilotChat/client.lua

@@ -244,7 +244,6 @@ end
 ---@field history table<CopilotChat.Provider.input>
 ---@field providers table<string, CopilotChat.Provider>
 ---@field provider_cache table<string, table>
----@field embedding_cache table<string, CopilotChat.context.embed>
 ---@field models table<string, CopilotChat.Client.model>?
 ---@field agents table<string, CopilotChat.Client.agent>?
 ---@field current_job string?
@@ -253,7 +252,6 @@ local Client = class(function(self)
   self.history = {}
   self.providers = {}
   self.provider_cache = {}
-  self.embedding_cache = {}
   self.models = nil
   self.agents = nil
   self.current_job = nil
@@ -749,25 +747,10 @@ function Client:embed(inputs, model)
   notify.publish(notify.STATUS, 'Generating embeddings for ' .. #inputs .. ' inputs')
 
   -- Initialize essentials
-  local to_process = {}
+  local to_process = inputs
   local results = {}
   local initial_chunk_size = 10
 
-  -- Process each input, using cache when possible
-  for _, input in ipairs(inputs) do
-    input.filename = input.filename or 'unknown'
-    input.filetype = input.filetype or 'text'
-
-    if input.content then
-      local cache_key = input.filename .. utils.quick_hash(input.content)
-      if self.embedding_cache[cache_key] then
-        table.insert(results, self.embedding_cache[cache_key])
-      else
-        table.insert(to_process, input)
-      end
-    end
-  end
-
   -- Process inputs in batches with adaptive chunk size
   while #to_process > 0 do
     local chunk_size = initial_chunk_size -- Reset chunk size for each new batch
@@ -814,9 +797,6 @@ function Client:embed(inputs, model)
         for _, embedding in ipairs(data) do
           local result = vim.tbl_extend('force', batch[embedding.index + 1], embedding)
           table.insert(results, result)
-
-          local cache_key = result.filename .. utils.quick_hash(result.content)
-          self.embedding_cache[cache_key] = result
         end
       end
     end
@@ -845,7 +825,6 @@ end
 function Client:reset()
   local stopped = self:stop()
   self.history = {}
-  self.embedding_cache = {}
   return stopped
 end
 

lua/CopilotChat/context.lua

@@ -23,6 +23,7 @@ local notify = require('CopilotChat.notify')
 local utils = require('CopilotChat.utils')
 local file_cache = {}
 local url_cache = {}
+local embedding_cache = {}
 
 local M = {}
 
@@ -661,15 +662,21 @@ function M.filter_embeddings(prompt, model, headless, embeddings)
   notify.publish(notify.STATUS, 'Ranking embeddings')
 
   -- Build query from history and prompt
-  local query = ''
+  local query = prompt
   if not headless then
-    for _, message in ipairs(client.history) do
-      if message.role == 'user' then
-        query = query .. '\n' .. message.content
-      end
-    end
+    query = table.concat(
+      vim
+        .iter(client.history)
+        :filter(function(m)
+          return m.role == 'user'
+        end)
+        :map(function(m)
+          return vim.trim(m.content)
+        end)
+        :totable(),
+      '\n'
+    ) .. '\n' .. prompt
   end
-  query = query .. '\n' .. prompt
 
   -- Rank embeddings by symbols
   embeddings = data_ranked_by_symbols(query, embeddings, MIN_SYMBOL_SIMILARITY)
@@ -678,26 +685,46 @@ function M.filter_embeddings(prompt, model, headless, embeddings)
     log.debug(string.format('%s: %s - %s', i, item.score, item.filename))
   end
 
-  -- Embed the query
-  table.insert(embeddings, {
+  -- Prepare embeddings for processing
+  local to_process = {}
+  local results = {}
+  for _, input in ipairs(embeddings) do
+    input.filename = input.filename or 'unknown'
+    input.filetype = input.filetype or 'text'
+    if input.content then
+      local cache_key = input.filename .. utils.quick_hash(input.content)
+      if embedding_cache[cache_key] then
+        table.insert(results, embedding_cache[cache_key])
+      else
+        table.insert(to_process, input)
+      end
+    end
+  end
+  table.insert(to_process, {
     content = query,
     filename = 'query',
     filetype = 'raw',
   })
 
-  -- Get embeddings from all items
-  embeddings = client:embed(embeddings, model)
+  -- Embed the data and process the results
+  for _, input in ipairs(client:embed(to_process, model)) do
+    if input.filetype ~= 'raw' then
+      local cache_key = input.filename .. utils.quick_hash(input.content)
+      embedding_cache[cache_key] = input
+    end
+    table.insert(results, input)
+  end
 
   -- Rate embeddings by relatedness to the query
-  local embedded_query = table.remove(embeddings, #embeddings)
+  local embedded_query = table.remove(results, #results)
   log.debug('Embedded query:', embedded_query.content)
-  embeddings = data_ranked_by_relatedness(embedded_query, embeddings, MIN_SEMANTIC_SIMILARITY)
-  log.debug('Ranked embeddings:', #embeddings)
-  for i, item in ipairs(embeddings) do
+  results = data_ranked_by_relatedness(embedded_query, results, MIN_SEMANTIC_SIMILARITY)
+  log.debug('Ranked embeddings:', #results)
+  for i, item in ipairs(results) do
     log.debug(string.format('%s: %s - %s', i, item.score, item.filename))
   end
 
-  return embeddings
+  return results
 end
 
 return M