From 7e1e3ac38e754482c7193c26ab93617a2184235e Mon Sep 17 00:00:00 2001
From: prplV <maseeeeeeeed@gmail.com>
Date: Tue, 1 Jul 2025 04:03:01 -0400
Subject: [PATCH] ws lug fixed + docker support

---
 .env.example       |  4 ++--
 Dockerfile         |  1 +
 docker-compose.yml |  5 +++--
 src/endpoints.rs   | 13 +++++++++----
 src/models.rs      |  2 +-
 5 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/.env.example b/.env.example
index a031e10..303099d 100644
--- a/.env.example
+++ b/.env.example
@@ -13,9 +13,9 @@ ML_MODEL_NAME="kis-test"
 # `ML_REQUEST_TIMEOUT` for setting ML response wait time
 # ------------------------------------------------------
 # - Measures in `secs`
-# - Default value - 10
+# - Default value - 15
 # ------------------------------------------------------
-ML_REQUEST_TIMEOUT="10" 
+ML_REQUEST_TIMEOUT="15" 
 
 
 # `ML_API_LOG_LEVEL` log level selection 
diff --git a/Dockerfile b/Dockerfile
index a3555bf..04e49de 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -17,5 +17,6 @@ COPY --from=builder /app/target/x86_64-unknown-linux-musl/release/ml-api /app/ml
 RUN apk add --no-cache ca-certificates
 
 EXPOSE 5134/tcp
+EXPOSE 1560/tcp
 
 ENTRYPOINT ["/app/ml-api"]
diff --git a/docker-compose.yml b/docker-compose.yml
index dfe390d..d4fe86d 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -8,13 +8,14 @@ services:
       - NOXIS_SOCKET_PATH=./noxis.sock
       - NOXIS_PROXY_PORT=7654
       - NOXIS_LOG_LEVEL=TRACE
-      - ML_TARGET_URL=http://192.168.231.89:11434/api/generate
+      - ML_TARGET_URL=http://192.168.249.5:11434/api/generate
       - ML_MODEL_NAME=kis-test
-      - ML_REQUEST_TIMEOUT=10
+      - ML_REQUEST_TIMEOUT=20
       - ML_LOG_LEVEL=TRACE
       - ML_API_PORT=5134
     ports:
       - 5134:5134
+      - 5135:1560
     restart: always
 
 networks:
diff --git a/src/endpoints.rs b/src/endpoints.rs
index 2ece019..5d60805 100644
--- a/src/endpoints.rs
+++ b/src/endpoints.rs
@@ -38,7 +38,7 @@ pub mod rest {
 
 pub mod ws {
     use axum::extract::{ws::{Message, WebSocket}, WebSocketUpgrade};
-    use tracing::trace;
+    use tracing::{debug, trace};
     use super::{IntoResponse, InputMetric, Config, State};
 
     pub async fn model_ws_handler(
@@ -63,17 +63,22 @@ pub mod ws {
                     Ok(msg) => {
                         match serde_json::from_str::<Vec<InputMetric>>(msg) {
                             Err(er) => {
-                                let _ = ws.send(
+                                match ws.send(
                                     Message::Text(
                                         format!("Cannot convert input message: {er}")
                                         .into()
                                     )
-                                ).await;
+                                ).await {
+                                    Ok(_) => trace!("cannot convert given msg"),
+                                    Err(er) => {let _ = ws.send(Message::Text(format!("Request was aborted: {}", er).into())).await;},
+                                }
                             },
                             Ok(req) => {
                                 match super::send_message_to_model(config.clone(), req).await {
                                     Ok(resp) => {
-                                        let _ = ws.send(Message::Text(resp.into()));
+                                        trace!("Successfull request to model");
+                                        debug!("Response from the Model: {}", &resp);
+                                        let _ = ws.send(Message::Text(resp.into())).await;
                                     },
                                     Err(er) => {
                                         let _ = ws.send(Message::Text(format!("Cannot get model's response: {er}").into())).await;
diff --git a/src/models.rs b/src/models.rs
index 16a6dbd..14ee7b7 100644
--- a/src/models.rs
+++ b/src/models.rs
@@ -48,7 +48,7 @@ impl ApiSessionConfig {
             model_name : std::env::var("ML_MODEL_NAME")?,
             request_timeout : std::env::var("ML_REQUEST_TIMEOUT")?.parse().unwrap_or_else(|_| {
                 warn!("invalid port was given, setting up default one ...");
-                10
+                15
             }),
             client : Client::new(),
         };