From 55e36c7e9da2e1c93ebeaabdc8355d0a22361f03 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sat, 31 Aug 2024 22:02:39 +0800 Subject: feat: webui support RAG (#815) --- src/serve.rs | 92 +++++++++++++++++++++++++++++++++++++++++++----------------- 1 file changed, 67 insertions(+), 25 deletions(-) (limited to 'src/serve.rs') diff --git a/src/serve.rs b/src/serve.rs index 65ac4cf..95b1402 100644 --- a/src/serve.rs +++ b/src/serve.rs @@ -1,4 +1,4 @@ -use crate::{client::*, config::*, utils::*}; +use crate::{client::*, config::*, function::*, rag::*, utils::*}; use anyhow::{anyhow, bail, Result}; use bytes::Bytes; @@ -65,20 +65,18 @@ pub async fn run(config: GlobalConfig, addr: Option) -> Result<()> { } struct Server { - clients: Vec, - model: Model, + config: Config, models: Vec, roles: Vec, + rags: Vec, } impl Server { fn new(config: &GlobalConfig) -> Self { - let config = config.read(); - let clients = config.clients.clone(); - let model = config.model.clone(); - let roles = Config::all_roles(); + let mut config = config.read().clone(); + config.functions = Functions::default(); let mut models = list_models(&config); - let mut default_model = model.clone(); + let mut default_model = config.model.clone(); default_model.data_mut().name = DEFAULT_MODEL_NAME.into(); models.insert(0, &default_model); let models: Vec = models @@ -101,12 +99,13 @@ impl Server { }) .collect(); Self { - clients, - model, - roles, + config, models, + roles: Config::all_roles(), + rags: Config::list_rags(), } } + async fn run(self: Arc, listener: TcpListener) -> Result> { let (tx, rx) = oneshot::channel(); tokio::spawn(async move { @@ -164,6 +163,10 @@ impl Server { self.list_models() } else if path == "/v1/roles" { self.list_roles() + } else if path == "/v1/rags" { + self.list_rags() + } else if path == "/v1/rags/search" { + self.search_rag(req).await } else if path == "/playground" || path == "/playground.html" { self.playground_page() } else if path == "/arena" || path == "/arena.html" { @@ -220,6 +223,43 @@ impl Server { Ok(res) } + fn list_rags(&self) -> Result { + let data = json!({ "data": self.rags }); + let res = Response::builder() + .header("Content-Type", "application/json; charset=utf-8") + .body(Full::new(Bytes::from(data.to_string())).boxed())?; + Ok(res) + } + + async fn search_rag(&self, req: hyper::Request) -> Result { + let req_body = req.collect().await?.to_bytes(); + let req_body: Value = serde_json::from_slice(&req_body) + .map_err(|err| anyhow!("Invalid request json, {err}"))?; + + debug!("search rag request: {req_body}"); + let SearchRagReqBody { name, input } = serde_json::from_value(req_body) + .map_err(|err| anyhow!("Invalid request body, {err}"))?; + + let config = Arc::new(RwLock::new(self.config.clone())); + + let abort_signal = create_abort_signal(); + + let rag = config + .read() + .rag_file(&name) + .ok() + .and_then(|rag_path| Rag::load(&config, &name, &rag_path).ok()) + .ok_or_else(|| anyhow!("Invalid rag"))?; + + let rag_result = Config::search_rag(&config, &rag, &input, abort_signal).await?; + + let data = json!({ "data": rag_result }); + let res = Response::builder() + .header("Content-Type", "application/json; charset=utf-8") + .body(Full::new(Bytes::from(data.to_string())).boxed())?; + Ok(res) + } + async fn chat_completions(&self, req: hyper::Request) -> Result { let req_body = req.collect().await?.to_bytes(); let req_body: Value = serde_json::from_slice(&req_body) @@ -238,16 +278,15 @@ impl Server { stream, } = req_body; - let config = Config { - clients: self.clients.to_vec(), - model: self.model.clone(), - ..Default::default() - }; + let config = self.config.clone(); + + let default_model = config.model.clone(); + let config = Arc::new(RwLock::new(config)); let (model_name, change) = if model == DEFAULT_MODEL_NAME { - (self.model.id(), true) - } else if self.model.id() == model { + (default_model.id(), true) + } else if default_model.id() == model { (model, false) } else { (model, true) @@ -261,7 +300,7 @@ impl Server { if max_tokens.is_some() { client.model_mut().set_max_tokens(max_tokens, true); } - let abort = create_abort_signal(); + let abort_signal = create_abort_signal(); let http_client = client.build_client()?; let completion_id = generate_completion_id(); @@ -280,7 +319,7 @@ impl Server { tokio::spawn(async move { let is_first = Arc::new(AtomicBool::new(true)); let (sse_tx, sse_rx) = unbounded_channel(); - let mut handler = SseHandler::new(sse_tx, abort); + let mut handler = SseHandler::new(sse_tx, abort_signal); async fn map_event( mut sse_rx: UnboundedReceiver, tx: &UnboundedSender, @@ -399,11 +438,8 @@ impl Server { model: embedding_model_id, } = req_body; - let config = Config { - clients: self.clients.to_vec(), - ..Default::default() - }; - let config = Arc::new(RwLock::new(config)); + let config = Arc::new(RwLock::new(self.config.clone())); + let embedding_model = Model::retrieve_embedding(&config.read(), &embedding_model_id)?; let texts = match input { @@ -444,6 +480,12 @@ impl Server { } } +#[derive(Debug, Deserialize)] +struct SearchRagReqBody { + name: String, + input: String, +} + #[derive(Debug, Deserialize)] struct ChatCompletionsReqBody { model: String, -- cgit v1.2.3