chatapp

backend/src/main.rs

@@ -2,39 +2,47 @@
 #[macro_use]
 extern crate rocket;
 
+use rocket::fairing::Fairing;
 use rocket::http::Method;
 use rocket::response::stream::{Event, EventStream};
 use rocket::serde::json::Json;
 use rocket::{Build, Rocket};
 use rocket_cors::{AllowedOrigins, CorsOptions};
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use tokio::sync::broadcast;
 
+use crate::llm::query_llm;
+
+pub mod auth;
+pub mod llm;
+
 /// ---------- shared broadcaster ----------
 struct ChatBroadcaster {
-    sender: broadcast::Sender<String>,
+    sender: broadcast::Sender<ChatMsg>,
 }
 
 impl ChatBroadcaster {
     fn new(buffer_size: usize) -> Self {
-        let (sender, _rx) = broadcast::channel::<String>(buffer_size);
+        let (sender, _rx) = broadcast::channel::<ChatMsg>(buffer_size);
         Self { sender }
     }
 
-    async fn publish(&self, msg: String) {
+    async fn publish(&self, msg: ChatMsg) {
         let _ = self.sender.send(msg);
     }
 
-    fn subscribe(&self) -> broadcast::Receiver<String> {
+    fn subscribe(&self) -> broadcast::Receiver<ChatMsg> {
         self.sender.subscribe()
     }
 }
 
 /// ---------- Rocket routes ----------
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 struct ChatMsg {
+    userid: usize,
     text: String,
+    timestamp: usize,
 }
 
 #[post("/chat", format = "json", data = "<msg>")]
@@ -42,8 +50,7 @@ async fn post_message(
     msg: Json<ChatMsg>,
     chat: &rocket::State<Arc<ChatBroadcaster>>,
 ) -> &'static str {
-    let text = msg.text.clone();
-    chat.publish(text).await;
+    chat.publish(msg.into_inner()).await;
     "Message sent"
 }
 
@@ -54,7 +61,7 @@ async fn event_stream(chat: &rocket::State<Arc<ChatBroadcaster>>) -> EventStream
     EventStream! {
         loop {
             match rx.recv().await {
-                Ok(msg) => yield Event::data(msg),
+                Ok(msg) => yield Event::json(&msg),
                 Err(broadcast::error::RecvError::Lagged(_)) => {
                     yield Event::comment("lagged");
                 }
@@ -64,6 +71,68 @@ async fn event_stream(chat: &rocket::State<Arc<ChatBroadcaster>>) -> EventStream
     }
 }
 
+// ---------- LLM worker ----------
+async fn start_llm_worker(chat: Arc<ChatBroadcaster>) {
+    let mut rx = chat.subscribe();
+
+    loop {
+        match rx.recv().await {
+            Ok(msg) => {
+                if msg.userid == 0 {
+                    // ignore bot messages
+                    continue;
+                }
+
+                let user_text = msg.text.clone();
+                let chat_clone = chat.clone();
+
+                rocket::tokio::spawn(async move {
+                    match query_llm(&user_text).await {
+                        Ok(reply) => {
+                            let bot_msg = ChatMsg {
+                                userid: 0,
+                                text: reply,
+                                timestamp: chrono::Local::now().timestamp() as usize,
+                            };
+                            chat_clone.publish(bot_msg).await;
+                        }
+                        Err(e) => eprintln!("LLM error: {}", e),
+                    }
+                });
+            }
+            Err(_) => break, // channel closed
+        }
+    }
+}
+
+pub struct LlmWorkerFairing;
+
+#[rocket::async_trait]
+impl Fairing for LlmWorkerFairing {
+    fn info(&self) -> rocket::fairing::Info {
+        rocket::fairing::Info {
+            name: "LLM background worker",
+            kind: rocket::fairing::Kind::Ignite,
+        }
+    }
+
+    async fn on_ignite(&self, rocket: Rocket<Build>) -> rocket::fairing::Result {
+        // Grab the shared broadcaster from state
+        let chat = rocket
+            .state::<Arc<ChatBroadcaster>>()
+            .expect("ChatBroadcaster not managed");
+        // Clone it so we can move into async block
+        let chat_clone = Arc::clone(chat);
+
+        // Spawn the background worker **inside** on_ignite
+        tokio::spawn(async move {
+            start_llm_worker(chat_clone).await;
+        });
+
+        Ok(rocket)
+    }
+}
+
 /// ---------- launch ----------
 #[launch]
 fn rocket() -> Rocket<Build> {
@@ -82,5 +151,9 @@ fn rocket() -> Rocket<Build> {
     rocket::build()
         .manage(chat)
         .attach(cors.to_cors().unwrap())
-        .mount("/", routes![post_message, event_stream])
+        .attach(LlmWorkerFairing {})
+        .mount(
+            "/",
+            routes![post_message, event_stream, auth::signup, auth::login],
+        )
 }