From 324465ed9539c7472736e6211c91b8b7a5496904 Mon Sep 17 00:00:00 2001
From: qooba <dev@qooba.net>
Date: Mon, 15 Jan 2024 01:39:12 +0100
Subject: [PATCH] add rest api embeddings

---
 yummy-rs/yummy/Cargo.toml  |  4 ++-
 yummy-rs/yummy/Makefile    |  5 ++++
 yummy-rs/yummy/src/main.rs | 56 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/yummy-rs/yummy/Cargo.toml b/yummy-rs/yummy/Cargo.toml
index d3385a7..48235e4 100644
--- a/yummy-rs/yummy/Cargo.toml
+++ b/yummy-rs/yummy/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "yummy"
-version = "0.0.8"
+version.workspace = true
 edition = "2021"
 repository = "https://github.com/yummyml/yummy"
 description = "The yummy library"
@@ -17,8 +17,10 @@ rustflags = ["-C", "target-cpu=native"]
 clap = "4.2.1"
 yummy-delta = { version = "0.0.9",  path = "../yummy-delta", optional = true }
 yummy-ml = { version = "0.0.9",  path = "../yummy-ml", optional = true }
+yummy-llm = { version = "0.0.9",  path = "../yummy-llm", optional = true }
 tokio = { version = "1.20.0", features = ["full"] }
 
 [features]
 yummy-delta = ["dep:yummy-delta"]
 yummy-ml = ["dep:yummy-ml"]
+yummy-llm = ["dep:yummy-llm"]
diff --git a/yummy-rs/yummy/Makefile b/yummy-rs/yummy/Makefile
index bfe9f24..8f5d785 100644
--- a/yummy-rs/yummy/Makefile
+++ b/yummy-rs/yummy/Makefile
@@ -30,7 +30,12 @@ test-delta:
 run-ml-serve:
 	cargo run --features yummy-ml ml serve --model /home/jovyan/yummy/yummy-rs/tests/mlflow/lightgbm_model/lightgbm_my_model/ --port 8080 --host 0.0.0.0 --loglevel debug
 
+run-llm-serve-embeddings-e5:
+	cargo run --features yummy-llm llm serve --config /home/jovyan/yummy/yummy-rs/tests/llm/config.yaml --port 8080 --host 0.0.0.0 --loglevel debug
+
+
 run-delta-serve:
 	cargo run --features yummy-delta delta serve --config /home/jovyan/yummy/yummy-rs/tests/delta/apply.yaml --port 8080 --host 0.0.0.0 --loglevel debug
 
 
+#curl -d '{"input":["Elixir of Eternal Twilight: Grants visions of realms beyond the veil."]}' -H "Content-Type: application/json" -X POST http://localhost:8080/embeddings | python3 -mjson.tool
diff --git a/yummy-rs/yummy/src/main.rs b/yummy-rs/yummy/src/main.rs
index 2a8f3c4..7902e47 100644
--- a/yummy-rs/yummy/src/main.rs
+++ b/yummy-rs/yummy/src/main.rs
@@ -44,6 +44,22 @@ fn cli() -> Command {
                         .arg_required_else_help(true),
                 ),
         )
+        .subcommand(
+            Command::new("llm")
+                .about("yummy llm")
+                .subcommand_required(true)
+                .subcommand(
+                    Command::new("serve")
+                        .about("yummy llm serve")
+                        .args(vec![
+                            arg!(--config <FILE> "config path"),
+                            arg!(--host <HOST> "host"),
+                            arg!(--port <PORT> "port"),
+                            arg!(--loglevel <LOGLEVEL> "log level"),
+                        ])
+                        .arg_required_else_help(true),
+                ),
+        )
 }
 
 #[tokio::main]
@@ -95,6 +111,26 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             }
             _ => unreachable!(),
         },
+        Some(("llm", sub_matches)) => match sub_matches.subcommand() {
+            Some(("serve", sub_sub_matches)) => {
+                let config = sub_sub_matches
+                    .get_one::<String>("config")
+                    .expect("required");
+                let host = sub_sub_matches.get_one::<String>("host").expect("required");
+                let port = sub_sub_matches
+                    .get_one::<String>("port")
+                    .expect("required")
+                    .parse::<u16>()?;
+
+                let log_level = sub_sub_matches
+                    .get_one::<String>("loglevel")
+                    .expect("required");
+
+                llm_serve(config.clone(), host.clone(), port, log_level.clone()).await?
+            }
+            _ => unreachable!(),
+        },
+
         _ => unreachable!(),
     }
 
@@ -150,3 +186,23 @@ async fn ml_serve(
 ) -> std::io::Result<()> {
     unreachable!()
 }
+
+#[cfg(feature = "yummy-llm")]
+async fn llm_serve(
+    config: String,
+    host: String,
+    port: u16,
+    log_level: String,
+) -> std::io::Result<()> {
+    yummy_llm::serve_llm(config, host, port, log_level).await
+}
+
+#[cfg(not(feature = "yummy-llm"))]
+async fn llm_serve(
+    _config: String,
+    _host: String,
+    _port: u16,
+    _log_level: String,
+) -> std::io::Result<()> {
+    unreachable!()
+}