ragrig_gpui/rag/
controller.rs

1//! The GPUI-side controller: commands in, Ragrig events out.
2//!
3//! [`RagController`] owns the worker thread, mirrors the state widgets need
4//! (sessions, library, component status), and re-emits every [`RagEvent`].
5//! It never runs Ragrig futures on GPUI's executor.
6
7use std::{
8    path::PathBuf,
9    thread,
10    time::{Duration, SystemTime},
11};
12
13use async_channel::{Sender, TrySendError};
14use gpui_kit::{Context, EventEmitter};
15
16use crate::model::{AgentComponent, ComponentKind, ComponentStatus, SessionEntry};
17#[cfg(feature = "grobid")]
18use crate::references::{ReferenceEntry, ReferenceOptions, ReferenceProgress, ReferenceReport};
19
20use super::{
21    AgentInfo, IngestProgress, LibrarySnapshot, PendingAttachment, RagCommand, RagEvent, RagSetup,
22    RequestId, SessionInfo, actor, count_label,
23};
24
25/// Sends commands to the Ragrig worker and mirrors its state for widgets.
26pub struct RagController {
27    commands: Sender<RagCommand>,
28    worker: Option<thread::JoinHandle<()>>,
29    sessions: Vec<SessionEntry>,
30    current_session: Option<String>,
31    library: LibrarySnapshot,
32    agent: Option<AgentInfo>,
33    ingest: Option<IngestProgress>,
34    sync_request: Option<RequestId>,
35    busy: bool,
36    last_error: Option<String>,
37    last_notice: Option<String>,
38    next_request: u64,
39    #[cfg(feature = "grobid")]
40    references: Vec<ReferenceEntry>,
41    #[cfg(feature = "grobid")]
42    reference_bibliography: Option<PathBuf>,
43    #[cfg(feature = "grobid")]
44    reference_report: Option<ReferenceReport>,
45    #[cfg(feature = "grobid")]
46    reference_progress: Option<ReferenceProgress>,
47    #[cfg(feature = "grobid")]
48    reference_request: Option<RequestId>,
49    #[cfg(feature = "grobid")]
50    reference_folder: Option<PathBuf>,
51}
52
53impl RagController {
54    /// Spawn the worker and start mirroring its events.
55    ///
56    /// Returns an error when the worker thread cannot be created; an agent
57    /// that fails to build is reported later through
58    /// [`RagEvent::AgentFailed`].
59    pub fn new(setup: RagSetup, cx: &mut Context<Self>) -> Result<Self, String> {
60        let (commands, command_rx) = async_channel::bounded(32);
61        let (events_tx, events) = async_channel::bounded(256);
62        let worker = actor::start(setup, None, command_rx, events_tx)
63            .map_err(|error| format!("cannot start the RAG worker: {error}"))?;
64
65        cx.spawn(async move |this, cx| {
66            while let Ok(event) = events.recv().await {
67                if this
68                    .update(cx, |controller, cx| controller.on_event(event, cx))
69                    .is_err()
70                {
71                    break;
72                }
73            }
74        })
75        .detach();
76
77        let controller = Self {
78            commands,
79            worker: Some(worker),
80            sessions: Vec::new(),
81            current_session: None,
82            library: LibrarySnapshot::default(),
83            agent: None,
84            ingest: None,
85            sync_request: None,
86            busy: false,
87            last_error: None,
88            last_notice: None,
89            next_request: 0,
90            #[cfg(feature = "grobid")]
91            references: Vec::new(),
92            #[cfg(feature = "grobid")]
93            reference_bibliography: None,
94            #[cfg(feature = "grobid")]
95            reference_report: None,
96            #[cfg(feature = "grobid")]
97            reference_progress: None,
98            #[cfg(feature = "grobid")]
99            reference_request: None,
100            #[cfg(feature = "grobid")]
101            reference_folder: None,
102        };
103        // The worker announces ready, the library, and the session list as
104        // soon as its agent is built.
105        Ok(controller)
106    }
107
108    /// Ask one question and return the request identity.
109    pub fn ask(
110        &mut self,
111        prompt: impl Into<String>,
112        attachments: Vec<PendingAttachment>,
113        cx: &mut Context<Self>,
114    ) -> Result<RequestId, String> {
115        let prompt = prompt.into();
116        let prompt = prompt.trim().to_owned();
117        if prompt.is_empty() {
118            return Err("enter a question first".to_owned());
119        }
120        let id = self.request_id();
121        self.send(
122            RagCommand::Ask {
123                id,
124                prompt,
125                attachments,
126            },
127            cx,
128        )?;
129        Ok(id)
130    }
131
132    /// Stop one running generation.
133    pub fn stop(&mut self, id: RequestId, cx: &mut Context<Self>) {
134        self.send_quiet(RagCommand::Stop { id }, cx);
135    }
136
137    /// Start an empty conversation.
138    pub fn new_session(&mut self, cx: &mut Context<Self>) {
139        self.send_quiet(RagCommand::NewSession, cx);
140    }
141
142    /// Continue a saved conversation.
143    pub fn load_session(&mut self, id: impl Into<String>, cx: &mut Context<Self>) {
144        self.send_quiet(RagCommand::LoadSession { id: id.into() }, cx);
145    }
146
147    /// Re-read the session store.
148    pub fn refresh_sessions(&mut self, cx: &mut Context<Self>) {
149        self.send_quiet(RagCommand::RefreshSessions, cx);
150    }
151
152    /// Re-read the vector store.
153    pub fn refresh_library(&mut self, cx: &mut Context<Self>) {
154        self.send_quiet(RagCommand::RefreshLibrary, cx);
155    }
156
157    /// Index a folder as a named corpus.
158    pub fn sync_corpus(
159        &mut self,
160        name: impl Into<String>,
161        folder: impl Into<PathBuf>,
162        cx: &mut Context<Self>,
163    ) -> Result<RequestId, String> {
164        let id = self.request_id();
165        self.send(
166            RagCommand::SyncCorpus {
167                id,
168                name: name.into(),
169                folder: folder.into(),
170            },
171            cx,
172        )?;
173        Ok(id)
174    }
175
176    /// Stop the running ingestion.
177    pub fn stop_sync(&mut self, cx: &mut Context<Self>) {
178        if let Some(id) = self.sync_request {
179            self.send_quiet(RagCommand::StopSync { id }, cx);
180        }
181    }
182
183    /// Parse `folder` with GROBID and merge fresh records into the workspace
184    /// bibliography.
185    #[cfg(feature = "grobid")]
186    pub fn extract_references(
187        &mut self,
188        folder: impl Into<PathBuf>,
189        options: ReferenceOptions,
190        cx: &mut Context<Self>,
191    ) -> Result<RequestId, String> {
192        let folder = folder.into();
193        let id = self.request_id();
194        self.reference_folder = Some(folder.clone());
195        self.send(
196            RagCommand::ExtractReferences {
197                id,
198                folder,
199                options,
200            },
201            cx,
202        )?;
203        Ok(id)
204    }
205
206    /// Stop the running reference extraction.
207    #[cfg(feature = "grobid")]
208    pub fn stop_references(&mut self, cx: &mut Context<Self>) {
209        if let Some(id) = self.reference_request {
210            self.send_quiet(RagCommand::StopReferences { id }, cx);
211        }
212    }
213
214    /// Load the workspace bibliography without contacting GROBID.
215    #[cfg(feature = "grobid")]
216    pub fn refresh_references(&mut self, folder: impl Into<PathBuf>, cx: &mut Context<Self>) {
217        let folder = folder.into();
218        self.reference_folder = Some(folder.clone());
219        self.send_quiet(RagCommand::RefreshReferences { folder }, cx);
220    }
221
222    /// Rebuild the agent with new configuration.
223    pub fn apply(&mut self, setup: RagSetup, cx: &mut Context<Self>) {
224        self.agent = None;
225        self.busy = false;
226        self.ingest = None;
227        self.sync_request = None;
228        self.current_session = None;
229        self.library = LibrarySnapshot::default();
230        self.last_error = None;
231        self.last_notice = None;
232        #[cfg(feature = "grobid")]
233        let reference_folder = self.reference_folder.clone();
234        #[cfg(feature = "grobid")]
235        {
236            self.references.clear();
237            self.reference_bibliography = None;
238            self.reference_report = None;
239            self.reference_progress = None;
240            self.reference_request = None;
241        }
242        self.send_quiet(
243            RagCommand::Apply {
244                setup: Box::new(setup),
245            },
246            cx,
247        );
248        #[cfg(feature = "grobid")]
249        if let Some(folder) = reference_folder {
250            self.send_quiet(RagCommand::RefreshReferences { folder }, cx);
251        }
252    }
253
254    /// Saved conversations, newest first.
255    pub fn sessions(&self) -> &[SessionEntry] {
256        &self.sessions
257    }
258
259    /// The active conversation id, when known.
260    pub fn current_session(&self) -> Option<&str> {
261        self.current_session.as_deref()
262    }
263
264    /// The last library snapshot.
265    pub fn library(&self) -> &LibrarySnapshot {
266        &self.library
267    }
268
269    /// The active agent, once it built successfully.
270    pub fn agent(&self) -> Option<&AgentInfo> {
271        self.agent.as_ref()
272    }
273
274    /// The running ingestion, if any.
275    pub fn ingest_progress(&self) -> Option<&IngestProgress> {
276        self.ingest.as_ref()
277    }
278
279    /// Whether an answer is streaming.
280    pub fn is_busy(&self) -> bool {
281        self.busy
282    }
283
284    /// Whether a corpus is being indexed.
285    pub fn is_ingesting(&self) -> bool {
286        self.ingest.is_some()
287    }
288
289    /// Reference rows loaded from the workspace bibliography.
290    #[cfg(feature = "grobid")]
291    pub fn references(&self) -> &[ReferenceEntry] {
292        &self.references
293    }
294
295    /// The `.bib` file backing the reference rows.
296    #[cfg(feature = "grobid")]
297    pub fn reference_bibliography(&self) -> Option<&std::path::Path> {
298        self.reference_bibliography.as_deref()
299    }
300
301    /// What the last reference extraction did.
302    #[cfg(feature = "grobid")]
303    pub fn reference_report(&self) -> Option<&ReferenceReport> {
304        self.reference_report.as_ref()
305    }
306
307    /// Progress of the running reference extraction.
308    #[cfg(feature = "grobid")]
309    pub fn reference_progress(&self) -> Option<&ReferenceProgress> {
310        self.reference_progress.as_ref()
311    }
312
313    /// Whether GROBID is currently parsing a corpus.
314    #[cfg(feature = "grobid")]
315    pub fn is_extracting_references(&self) -> bool {
316        self.reference_progress.is_some()
317    }
318
319    /// The last fatal error.
320    pub fn error(&self) -> Option<&str> {
321        self.last_error.as_deref()
322    }
323
324    /// One line for a status bar.
325    pub fn status_line(&self) -> String {
326        if let Some(error) = &self.last_error {
327            return error.clone();
328        }
329        if let Some(progress) = &self.ingest {
330            return progress.summary();
331        }
332        #[cfg(feature = "grobid")]
333        if let Some(progress) = &self.reference_progress {
334            return progress.summary();
335        }
336        if self.busy {
337            return "Generating…".to_owned();
338        }
339        match &self.agent {
340            Some(agent) => format!(
341                "{} · {} · {}",
342                agent.chat_model,
343                count_label(agent.documents, "document", "documents"),
344                count_label(agent.chunks, "chunk", "chunks")
345            ),
346            None => self
347                .last_notice
348                .clone()
349                .unwrap_or_else(|| "Starting the RAG worker…".to_owned()),
350        }
351    }
352
353    /// Pipeline components for [`AgentBar`](crate::AgentBar).
354    pub fn agent_components(&self) -> Vec<AgentComponent> {
355        let Some(agent) = &self.agent else {
356            let status = if self.last_error.is_some() {
357                ComponentStatus::Error
358            } else {
359                ComponentStatus::Inactive
360            };
361            let components: Vec<AgentComponent> = [
362                (ComponentKind::Chat, "Chat model"),
363                (ComponentKind::Embedder, "Embedder"),
364                (ComponentKind::Store, "Vector store"),
365                (ComponentKind::Parser, "PDF parser"),
366                (ComponentKind::Chunker, "Chunker"),
367                (ComponentKind::Retriever, "Retriever"),
368            ]
369            .into_iter()
370            .map(|(kind, label)| AgentComponent::new(kind_label(kind), kind, label).status(status))
371            .collect();
372            #[cfg(feature = "grobid")]
373            let components = {
374                let mut components = components;
375                components.push(
376                    AgentComponent::new("references", ComponentKind::References, "GROBID")
377                        .status(status),
378                );
379                components
380            };
381            return components;
382        };
383
384        let pipeline = if self.ingest.is_some() {
385            ComponentStatus::Busy
386        } else {
387            ComponentStatus::Ready
388        };
389        let chat = if self.busy {
390            ComponentStatus::Busy
391        } else {
392            ComponentStatus::Ready
393        };
394
395        let components = vec![
396            AgentComponent::new("chat", ComponentKind::Chat, agent.chat_model.clone())
397                .status(chat)
398                .detail(format!("{} · {} memory", agent.chat_backend, agent.memory)),
399            AgentComponent::new("embed", ComponentKind::Embedder, agent.embed_model.clone())
400                .status(pipeline)
401                .detail(format!("{} · top {}", agent.embed_backend, agent.top_k)),
402            AgentComponent::new("store", ComponentKind::Store, "Vector store")
403                .status(pipeline)
404                .detail(format!(
405                    "{} · {}",
406                    count_label(agent.documents, "document", "documents"),
407                    count_label(agent.chunks, "chunk", "chunks")
408                )),
409            AgentComponent::new("parser", ComponentKind::Parser, agent.pdf_backend.clone())
410                .status(pipeline)
411                .detail("PDF and Markdown"),
412            AgentComponent::new("chunker", ComponentKind::Chunker, "Markdown")
413                .status(pipeline)
414                .detail(format!(
415                    "{} tokens · {} overlap",
416                    agent.chunk_size, agent.chunk_overlap
417                )),
418            AgentComponent::new("retriever", ComponentKind::Retriever, "Hybrid retrieval")
419                .status(pipeline)
420                .detail(format!("top {}", agent.top_k)),
421        ];
422        #[cfg(feature = "grobid")]
423        let components = {
424            let mut components = components;
425            let status = if self.reference_progress.is_some() {
426                ComponentStatus::Busy
427            } else if self.references.is_empty() {
428                ComponentStatus::Inactive
429            } else {
430                ComponentStatus::Ready
431            };
432            let detail = match &self.reference_bibliography {
433                Some(path) => format!(
434                    "{} · {}",
435                    count_label(self.references.len(), "reference", "references"),
436                    path.file_name()
437                        .and_then(|name| name.to_str())
438                        .unwrap_or("references.bib")
439                ),
440                None => count_label(self.references.len(), "reference", "references"),
441            };
442            components.push(
443                AgentComponent::new("references", ComponentKind::References, "GROBID")
444                    .status(status)
445                    .detail(detail),
446            );
447            components
448        };
449        components
450    }
451
452    fn request_id(&mut self) -> RequestId {
453        self.next_request = self.next_request.wrapping_add(1);
454        RequestId::new(self.next_request)
455    }
456
457    fn send(&self, command: RagCommand, cx: &mut Context<Self>) -> Result<(), String> {
458        self.commands
459            .try_send(command)
460            .map_err(|error| match error {
461                TrySendError::Full(_) => "the RAG worker is busy; try again",
462                TrySendError::Closed(_) => "the RAG worker stopped",
463            })?;
464        cx.notify();
465        Ok(())
466    }
467
468    fn send_quiet(&mut self, command: RagCommand, cx: &mut Context<Self>) {
469        if let Err(message) = self.send(command, cx) {
470            self.last_error = Some(message);
471            cx.notify();
472        }
473    }
474
475    fn on_event(&mut self, event: RagEvent, cx: &mut Context<Self>) {
476        match &event {
477            RagEvent::AgentReady(info) => {
478                self.agent = Some(info.clone());
479                self.last_error = None;
480            }
481            RagEvent::AgentFailed { message } => {
482                self.agent = None;
483                self.busy = false;
484                self.last_error = Some(message.clone());
485            }
486            RagEvent::SessionStarted { id } | RagEvent::SessionLoaded { id, .. } => {
487                self.current_session = Some(id.clone());
488            }
489            RagEvent::Sessions(sessions) => {
490                self.sessions = session_entries(sessions, SystemTime::now());
491            }
492            RagEvent::Started { .. } => {
493                self.busy = true;
494                self.last_notice = None;
495                self.last_error = None;
496            }
497            RagEvent::Chunk { .. } => {}
498            RagEvent::Finished { .. } | RagEvent::Stopped { .. } => {
499                self.busy = false;
500            }
501            RagEvent::Failed { message, .. } => {
502                self.busy = false;
503                self.last_error = Some(message.clone());
504            }
505            RagEvent::Library(snapshot) => {
506                self.library = snapshot.clone();
507                if let Some(agent) = &mut self.agent {
508                    agent.chunks = snapshot.chunks;
509                    agent.documents = snapshot.documents.len();
510                }
511            }
512            #[cfg(feature = "grobid")]
513            RagEvent::Citations { .. } => {}
514            #[cfg(feature = "grobid")]
515            RagEvent::ReferencesStarted { id, folder } => {
516                self.reference_request = Some(*id);
517                self.reference_folder = Some(folder.clone());
518                self.reference_progress = Some(ReferenceProgress::default());
519                self.last_error = None;
520            }
521            #[cfg(feature = "grobid")]
522            RagEvent::ReferencesProgress { progress, .. } => {
523                self.reference_progress = Some(progress.clone());
524            }
525            #[cfg(feature = "grobid")]
526            RagEvent::ReferencesLoaded {
527                entries,
528                bibliography,
529            } => {
530                self.references = entries.clone();
531                self.reference_bibliography = Some(bibliography.clone());
532            }
533            #[cfg(feature = "grobid")]
534            RagEvent::ReferencesFinished { report, .. } => {
535                self.reference_report = Some(report.clone());
536                self.reference_progress = None;
537                self.reference_request = None;
538            }
539            #[cfg(feature = "grobid")]
540            RagEvent::ReferencesStopped { .. } => {
541                self.reference_progress = None;
542                self.reference_request = None;
543            }
544            #[cfg(feature = "grobid")]
545            RagEvent::ReferencesFailed { message, .. } => {
546                self.reference_progress = None;
547                self.reference_request = None;
548                self.last_error = Some(message.clone());
549            }
550            RagEvent::SyncStarted { id, corpus } => {
551                self.sync_request = Some(*id);
552                self.last_error = None;
553                self.ingest = Some(IngestProgress {
554                    corpus: corpus.clone(),
555                    ..IngestProgress::default()
556                });
557            }
558            RagEvent::SyncProgress { progress, .. } => {
559                self.ingest = Some(progress.clone());
560            }
561            RagEvent::SyncFinished { .. } | RagEvent::SyncStopped { .. } => {
562                self.ingest = None;
563                self.sync_request = None;
564            }
565            RagEvent::SyncFailed { message, .. } => {
566                self.ingest = None;
567                self.sync_request = None;
568                self.last_error = Some(message.clone());
569            }
570            RagEvent::Notice(message) => {
571                self.last_notice = Some(message.clone());
572            }
573        }
574
575        cx.notify();
576        cx.emit(event);
577    }
578}
579
580impl EventEmitter<RagEvent> for RagController {}
581
582impl Drop for RagController {
583    fn drop(&mut self) {
584        self.commands.close();
585        if let Some(worker) = self.worker.take() {
586            let _ = worker.join();
587        }
588    }
589}
590
591fn kind_label(kind: ComponentKind) -> &'static str {
592    match kind {
593        ComponentKind::Chat => "chat",
594        ComponentKind::Embedder => "embed",
595        ComponentKind::Store => "store",
596        ComponentKind::Parser => "parser",
597        ComponentKind::Chunker => "chunker",
598        ComponentKind::Retriever => "retriever",
599        ComponentKind::Vision => "vision",
600        ComponentKind::References => "references",
601        ComponentKind::Other => "component",
602    }
603}
604
605fn session_entries(sessions: &[SessionInfo], now: SystemTime) -> Vec<SessionEntry> {
606    let mut sessions = sessions.to_vec();
607    sessions.sort_by_key(|session| std::cmp::Reverse(session.updated));
608    sessions
609        .into_iter()
610        .map(|session| {
611            let label = session
612                .summary
613                .as_deref()
614                .map(str::trim)
615                .filter(|summary| !summary.is_empty())
616                .map(str::to_owned)
617                .unwrap_or_else(|| match session.turn_count {
618                    0 => "Empty conversation".to_owned(),
619                    1 => "1 turn".to_owned(),
620                    count => format!("{count} turns"),
621                });
622            SessionEntry::new(session.id, label).detail(relative_time(now, session.updated))
623        })
624        .collect()
625}
626
627fn relative_time(now: SystemTime, then: SystemTime) -> String {
628    let elapsed = now.duration_since(then).unwrap_or_default();
629    if elapsed < Duration::from_secs(60) {
630        "just now".to_owned()
631    } else if elapsed < Duration::from_secs(3_600) {
632        format!("{} min ago", elapsed.as_secs() / 60)
633    } else if elapsed < Duration::from_secs(86_400) {
634        format!("{} h ago", elapsed.as_secs() / 3_600)
635    } else {
636        format!("{} d ago", elapsed.as_secs() / 86_400)
637    }
638}
639
640#[cfg(test)]
641mod tests {
642    use super::*;
643
644    fn info(id: &str, turns: usize, summary: Option<&str>, updated: u64) -> SessionInfo {
645        SessionInfo {
646            id: id.to_owned(),
647            turn_count: turns,
648            summary: summary.map(str::to_owned),
649            updated: SystemTime::UNIX_EPOCH + Duration::from_secs(updated),
650        }
651    }
652
653    #[test]
654    fn sessions_are_sorted_and_labelled() {
655        let now = SystemTime::UNIX_EPOCH + Duration::from_secs(10_000);
656        let entries = session_entries(
657            &[
658                info("old", 1, None, 1_000),
659                info("new", 3, None, 9_000),
660                info("summary", 0, Some("Methods review"), 5_000),
661            ],
662            now,
663        );
664        assert_eq!(entries[0].id, "new");
665        assert_eq!(entries[0].label, "3 turns");
666        assert_eq!(entries[1].id, "summary");
667        assert_eq!(entries[1].label, "Methods review");
668        assert_eq!(entries[2].id, "old");
669        assert_eq!(entries[2].label, "1 turn");
670    }
671
672    #[test]
673    fn relative_times_are_human_readable() {
674        let now = SystemTime::UNIX_EPOCH + Duration::from_secs(100_000);
675        assert_eq!(relative_time(now, now - Duration::from_secs(5)), "just now");
676        assert_eq!(
677            relative_time(now, now - Duration::from_secs(120)),
678            "2 min ago"
679        );
680        assert_eq!(
681            relative_time(now, now - Duration::from_secs(7_200)),
682            "2 h ago"
683        );
684        assert_eq!(
685            relative_time(now, now - Duration::from_secs(172_800)),
686            "2 d ago"
687        );
688    }
689}
ragrig_gpui/rag/controller.rs

ragrig_gpui/rag/
controller.rs
