refactor: `.sources rag` show document ids (#975)

1 files changed, 17 insertions, 9 deletions

diff --git a/src/rag/mod.rs b/src/rag/mod.rs
index 1fd0376..4ddb0da 100644
--- a/src/rag/mod.rs
+++ b/src/rag/mod.rs
@@ -214,18 +214,26 @@ impl Rag {
     }
 
     pub fn set_last_sources(&self, ids: &[DocumentId]) {
-        let sources: IndexSet<_> = ids
-            .iter()
-            .filter_map(|id| {
-                let (file_index, _) = id.split();
-                let file = self.data.files.get(&file_index)?;
-                Some(file.path.clone())
-            })
-            .collect();
+        let mut sources: IndexMap<String, Vec<String>> = IndexMap::new();
+        for id in ids {
+            let (file_index, _) = id.split();
+            if let Some(file) = self.data.files.get(&file_index) {
+                sources
+                    .entry(file.path.clone())
+                    .or_default()
+                    .push(format!("{id:?}"));
+            }
+        }
         let sources = if sources.is_empty() {
             None
         } else {
-            Some(sources.into_iter().collect::<Vec<_>>().join("\n"))
+            Some(
+                sources
+                    .into_iter()
+                    .map(|(path, ids)| format!("{path} ({})", ids.join(",")))
+                    .collect::<Vec<_>>()
+                    .join("\n"),
+            )
         };
         *self.last_sources.write() = sources;
     }