refactor(search): improve symbol search and references

Improves symbol search algorithm by using trigram-based similarity matching
instead of simple term matching. This provides better accuracy for fuzzy
matches and handles compound words better.

Key changes:
- Replace basic term matching with trigram-based Jaccard similarity
- Normalize and weight symbol matches separately from filename matches
- Add reference list display in chat UI
- Move authentication status messages to more logical locations
- Fix various edge cases and improve code organization

Signed-off-by: Tomas Slusny <slusnucky@gmail.com>

Author: deathbeam

lua/CopilotChat/client.lua

@@ -8,6 +8,7 @@
 ---@field temperature number?
 ---@field on_progress nil|fun(response: string):nil
 
+local async = require('plenary.async')
 local log = require('plenary.log')
 local tiktoken = require('CopilotChat.tiktoken')
 local notify = require('CopilotChat.notify')
@@ -187,21 +188,6 @@ local function generate_embedding_request(inputs, threshold)
   end, inputs)
 end
 
-local function generate_references(references)
-  local out = ''
-
-  for _, reference in ipairs(references) do
-    out = out .. '\n[' .. reference.name .. '](' .. reference.url .. ')'
-  end
-
-  if out == '' then
-    return out
-  end
-
-  out = '\n\n**`References:`**' .. out
-  return out
-end
-
 ---@class CopilotChat.Client : Class
 ---@field providers table<string, CopilotChat.Provider>
 ---@field provider_cache table<string, table>
@@ -235,7 +221,6 @@ function Client:authenticate(provider_name)
   if not headers or (expires_at and expires_at <= math.floor(os.time())) then
     local token
     if provider.get_token then
-      notify.publish(notify.STATUS, 'Authenticating to provider ' .. provider_name)
       token, expires_at = provider.get_token()
     end
 
@@ -260,8 +245,8 @@ function Client:fetch_models()
   for _, provider_name in ipairs(provider_order) do
     local provider = self.providers[provider_name]
     if not provider.disabled and provider.get_models then
-      local headers = self:authenticate(provider_name)
       notify.publish(notify.STATUS, 'Fetching models from ' .. provider_name)
+      local headers = self:authenticate(provider_name)
       local ok, provider_models = pcall(provider.get_models, headers)
       if ok then
         for _, model in ipairs(provider_models) do
@@ -296,8 +281,8 @@ function Client:fetch_agents()
   for _, provider_name in ipairs(provider_order) do
     local provider = self.providers[provider_name]
     if not provider.disabled and provider.get_agents then
-      local headers = self:authenticate(provider_name)
       notify.publish(notify.STATUS, 'Fetching agents from ' .. provider_name)
+      local headers = self:authenticate(provider_name)
       local ok, provider_agents = pcall(provider.get_agents, headers)
       if ok then
         for _, agent in ipairs(provider_agents) do
@@ -320,6 +305,7 @@ end
 --- Ask a question to Copilot
 ---@param prompt string: The prompt to send to Copilot
 ---@param opts CopilotChat.Client.ask: Options for the request
+---@return string, table, number, number
 function Client:ask(prompt, opts)
   opts = opts or {}
   prompt = vim.trim(prompt)
@@ -369,10 +355,13 @@ function Client:ask(prompt, opts)
   log.debug('Tokenizer: ', tokenizer)
   tiktoken.load(tokenizer)
 
+  notify.publish(notify.STATUS, 'Generating request')
+
+  async.util.scheduler()
   local references = {}
   for _, embed in ipairs(embeddings) do
     table.insert(references, {
-      name = embed.filename,
+      name = utils.filename(embed.filename),
       url = embed.filename,
     })
   end
@@ -540,6 +529,8 @@ function Client:ask(prompt, opts)
     parse_stream_line(line, job)
   end
 
+  notify.publish(notify.STATUS, 'Thinking')
+
   opts.agent = opts.agent and opts.agent:gsub(':' .. provider_name .. '$', '')
   opts.model = opts.model:gsub(':' .. provider_name .. '$', '')
   local headers = self:authenticate(provider_name)
@@ -560,12 +551,10 @@ function Client:ask(prompt, opts)
     args.stream = stream_func
   end
 
-  notify.publish(notify.STATUS, 'Thinking')
-
   local response, err = utils.curl_post(provider.get_url(opts), args)
 
   if self.current_job ~= job_id then
-    return nil, nil, nil
+    return
   end
 
   self.current_job = nil
@@ -628,20 +617,12 @@ function Client:ask(prompt, opts)
     return
   end
 
-  local full_references = generate_references(references)
-  log.info('References: ', full_references)
-  if full_references ~= '' then
-    full_response = full_response .. full_references
-    if on_progress then
-      on_progress(full_references)
-    end
-  end
-
   log.trace('Full response: ', full_response)
   log.debug('Last message: ', last_message)
 
   return full_response,
-    last_message and last_message.usage and last_message.usage.total_tokens,
+    references,
+    last_message and last_message.usage and last_message.usage.total_tokens or 0,
     max_tokens
 end
 

lua/CopilotChat/context.lua

@@ -14,7 +14,7 @@
 ---@field outline string?
 ---@field symbols table<string, CopilotChat.context.symbol>?
 ---@field embedding table<number>?
----@field score?
+---@field score number?
 
 local async = require('plenary.async')
 local log = require('plenary.log')
@@ -66,8 +66,8 @@ local OFF_SIDE_RULE_LANGUAGES = {
   'fsharp',
 }
 
-local MIN_SYMBOL_SIMILARITY = 0.4 -- Symbol-based matching can be more lenient
-local MIN_SEMANTIC_SIMILARITY = 0.3 -- Cosine similarity should be stricter for relevance
+local MIN_SYMBOL_SIMILARITY = 0.3
+local MIN_SEMANTIC_SIMILARITY = 0.4
 local MULTI_FILE_THRESHOLD = 5
 local MAX_FILES = 2500
 
@@ -115,65 +115,106 @@ local function data_ranked_by_relatedness(query, data, min_similarity)
   return results
 end
 
---- Rank data by symbols
----@param prompt string
----@param data table<CopilotChat.context.embed>
----@param min_similarity number
----@return table<CopilotChat.context.embed>
-local function data_ranked_by_symbols(prompt, data, min_similarity)
-  local query_terms = {}
-  for term in prompt:lower():gmatch('%w+') do
-    query_terms[term] = true
+-- Create trigrams from text (e.g., "hello" -> {"hel", "ell", "llo"})
+local function get_trigrams(text)
+  local trigrams = {}
+  text = text:lower()
+  for i = 1, #text - 2 do
+    trigrams[text:sub(i, i + 2)] = true
+  end
+  return trigrams
+end
+
+-- Calculate Jaccard similarity between two trigram sets
+local function trigram_similarity(set1, set2)
+  local intersection = 0
+  local union = 0
+
+  -- Count intersection and union
+  for trigram in pairs(set1) do
+    if set2[trigram] then
+      intersection = intersection + 1
+    end
+    union = union + 1
+  end
+
+  for trigram in pairs(set2) do
+    if not set1[trigram] then
+      union = union + 1
+    end
+  end
+
+  return intersection / union
+end
+
+local function data_ranked_by_symbols(query, data, min_similarity)
+  -- Get query trigrams including compound versions
+  local query_trigrams = {}
+
+  -- Add trigrams for each word
+  for term in query:lower():gmatch('%w+') do
+    for trigram in pairs(get_trigrams(term)) do
+      query_trigrams[trigram] = true
+    end
+  end
+
+  -- Add trigrams for compound query
+  local compound_query = query:lower():gsub('[^%w]', '')
+  for trigram in pairs(get_trigrams(compound_query)) do
+    query_trigrams[trigram] = true
   end
 
   local results = {}
-  for _, entry in ipairs(data) do
-    local total_terms = 0
-    local matched_terms = 0
-    local filename = entry.filename and entry.filename:lower() or ''
-
-    -- Calculate similarity score based on term matches
-    for term in pairs(query_terms) do
-      total_terms = total_terms + 1
-
-      -- Filename matches
-      if filename:find(term, 1, true) then
-        matched_terms = matched_terms + 1
-        if vim.fn.fnamemodify(filename, ':t'):gsub('%..*$', '') == term then
-          matched_terms = matched_terms + 0.5 -- Bonus for exact filename match
-        end
-      end
+  local max_score = 0
 
-      -- Symbol matches
-      if entry.symbols then
-        for _, symbol in ipairs(entry.symbols) do
-          if symbol.name and symbol.name:lower():find(term, 1, true) then
-            matched_terms = matched_terms + 1
-            if symbol.name:lower() == term then
-              matched_terms = matched_terms + 0.5 -- Bonus for exact symbol match
-            end
-          end
-          if symbol.signature and symbol.signature:lower():find(term, 1, true) then
-            matched_terms = matched_terms + 0.5 -- Partial credit for signature matches
-          end
+  for _, entry in ipairs(data) do
+    local score = 0
+    local basename = vim.fn.fnamemodify(entry.filename, ':t'):gsub('%..*$', '')
+
+    -- Get trigrams for basename and compound version
+    local file_trigrams = get_trigrams(basename)
+    local compound_trigrams = get_trigrams(basename:gsub('[^%w]', ''))
+
+    -- Calculate similarities
+    local name_sim = trigram_similarity(query_trigrams, file_trigrams)
+    local compound_sim = trigram_similarity(query_trigrams, compound_trigrams)
+
+    -- Take best match
+    score = math.max(name_sim, compound_sim)
+
+    -- Add symbol matches
+    if entry.symbols then
+      local symbol_score = 0
+      for _, symbol in ipairs(entry.symbols) do
+        if symbol.name then
+          local symbol_trigrams = get_trigrams(symbol.name)
+          local sym_sim = trigram_similarity(query_trigrams, symbol_trigrams)
+          symbol_score = math.max(symbol_score, sym_sim)
         end
       end
+      score = score + (symbol_score * 0.5) -- Weight symbol matches less
     end
 
-    -- Calculate similarity score (0 to 1 range)
-    local similarity = matched_terms / (total_terms * 2) -- Denominator accounts for potential bonuses
+    if score > 0 then
+      max_score = math.max(max_score, score)
+      table.insert(results, vim.tbl_extend('force', entry, { score = score }))
+    end
+  end
 
-    -- Only include results above similarity threshold
-    if similarity >= min_similarity then
-      table.insert(results, vim.tbl_extend('force', entry, { score = similarity }))
+  -- Normalize and filter results
+  local filtered_results = {}
+  for _, result in ipairs(results) do
+    result.score = result.score / max_score
+    if result.score >= min_similarity then
+      table.insert(filtered_results, result)
     end
   end
 
-  table.sort(results, function(a, b)
+  table.sort(filtered_results, function(a, b)
     return a.score > b.score
   end)
 
-  return results
+  return filtered_results
 end
 
 --- Get the full signature of a declaration
@@ -603,6 +644,8 @@ function M.filter_embeddings(prompt, model, embeddings)
     return embeddings
   end
 
+  notify.publish(notify.STATUS, 'Ranking embeddings')
+
   -- Rank embeddings by symbols
   embeddings = data_ranked_by_symbols(prompt, embeddings, MIN_SYMBOL_SIMILARITY)
   log.debug('Ranked data:', #embeddings)

lua/CopilotChat/init.lua

@@ -659,21 +659,22 @@ function M.ask(prompt, config)
       return
     end
 
-    local ask_ok, response, token_count, token_max_count = pcall(client.ask, client, prompt, {
-      history = history,
-      selection = selection,
-      embeddings = filtered_embeddings,
-      system_prompt = system_prompt,
-      model = selected_model,
-      agent = selected_agent,
-      temperature = config.temperature,
-      on_progress = vim.schedule_wrap(function(token)
-        if not config.headless then
-          state.chat:append(token)
-        end
-        has_output = true
-      end),
-    })
+    local ask_ok, response, references, token_count, token_max_count =
+      pcall(client.ask, client, prompt, {
+        history = history,
+        selection = selection,
+        embeddings = filtered_embeddings,
+        system_prompt = system_prompt,
+        model = selected_model,
+        agent = selected_agent,
+        temperature = config.temperature,
+        on_progress = vim.schedule_wrap(function(token)
+          if not config.headless then
+            state.chat:append(token)
+          end
+          has_output = true
+        end),
+      })
 
     async.util.scheduler()
 
@@ -691,6 +692,7 @@ function M.ask(prompt, config)
 
     if not config.headless then
       state.last_response = response
+      state.chat.references = references
       state.chat.token_count = token_count
       state.chat.token_max_count = token_max_count
     end

lua/CopilotChat/ui/chat.lua

@@ -65,6 +65,7 @@ end
 ---@field spinner CopilotChat.ui.Spinner
 ---@field sections table<CopilotChat.ui.Chat.Section>
 ---@field config CopilotChat.config.shared
+---@field references table
 ---@field token_count number?
 ---@field token_max_count number?
 local Chat = class(function(self, question_header, answer_header, separator, help, on_buf_create)
@@ -82,6 +83,7 @@ local Chat = class(function(self, question_header, answer_header, separator, hel
 
   -- Variables
   self.config = {}
+  self.references = {}
   self.token_count = nil
   self.token_max_count = nil
 end, Overlay)
@@ -250,6 +252,16 @@ function Chat:render()
       msg = msg .. self.token_count .. '/' .. self.token_max_count .. ' tokens used'
     end
 
+    if self.references and #self.references > 0 then
+      if msg ~= '' then
+        msg = msg .. '\n'
+      end
+      msg = msg .. '\nReferences:\n'
+      for _, ref in ipairs(self.references) do
+        msg = msg .. '  ' .. ref.name .. '\n'
+      end
+    end
+
     self:show_help(msg, last_section.start_line - last_section.end_line - 1)
   else
     self:clear_help()
@@ -438,6 +450,7 @@ end
 
 function Chat:clear()
   self:validate()
+  self.references = {}
   self.token_count = nil
   self.token_max_count = nil
   vim.bo[self.bufnr].modifiable = true