add message transformer

Author: mcharytoniuk

src/balancer/chunk_forwarding_session_controller/identity_transformer.rs

@@ -0,0 +1,23 @@
+use anyhow::Result;
+use async_trait::async_trait;
+
+use super::transforms_outgoing_message::TransformsOutgoingMessage;
+use crate::balancer::inference_client::Message as OutgoingMessage;
+
+#[derive(Clone)]
+pub struct IdentityTransformer;
+
+impl IdentityTransformer {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+#[async_trait]
+impl TransformsOutgoingMessage for IdentityTransformer {
+    type TransformedMessage = OutgoingMessage;
+
+    async fn transform(&self, message: OutgoingMessage) -> Result<Self::TransformedMessage> {
+        return Ok(message);
+    }
+}

src/balancer/chunk_forwarding_session_controller/mod.rs

@@ -0,0 +1,49 @@
+pub mod identity_transformer;
+pub mod transforms_outgoing_message;
+
+use async_trait::async_trait;
+use tokio::sync::mpsc;
+
+use self::transforms_outgoing_message::TransformsOutgoingMessage;
+use crate::balancer::inference_client::Message as OutgoingMessage;
+use crate::controls_session::ControlsSession;
+
+#[derive(Clone)]
+pub struct ChunkForwardingSessionController<TTransformsOutgoingMessage>
+where
+    TTransformsOutgoingMessage: Clone + TransformsOutgoingMessage + Send + Sync,
+{
+    chunk_tx: mpsc::UnboundedSender<String>,
+    transformer: TTransformsOutgoingMessage,
+}
+
+impl<TTransformsOutgoingMessage> ChunkForwardingSessionController<TTransformsOutgoingMessage>
+where
+    TTransformsOutgoingMessage: Clone + TransformsOutgoingMessage + Send + Sync,
+{
+    pub fn new(
+        chunk_tx: mpsc::UnboundedSender<String>,
+        transformer: TTransformsOutgoingMessage,
+    ) -> Self {
+        Self {
+            chunk_tx,
+            transformer,
+        }
+    }
+}
+
+#[async_trait]
+impl<TTransformsOutgoingMessage> ControlsSession<OutgoingMessage>
+    for ChunkForwardingSessionController<TTransformsOutgoingMessage>
+where
+    TTransformsOutgoingMessage: Clone + TransformsOutgoingMessage + Send + Sync,
+{
+    async fn send_response(&mut self, message: OutgoingMessage) -> anyhow::Result<()> {
+        let transformed_message = self.transformer.transform(message).await?;
+        let stringified_message = serde_json::to_string(&transformed_message)?;
+
+        self.chunk_tx.send(stringified_message)?;
+
+        Ok(())
+    }
+}

src/balancer/chunk_forwarding_session_controller/transforms_outgoing_message.rs

@@ -0,0 +1,12 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use serde::Serialize;
+
+use crate::balancer::inference_client::Message as OutgoingMessage;
+
+#[async_trait]
+pub trait TransformsOutgoingMessage {
+    type TransformedMessage: Serialize;
+
+    async fn transform(&self, message: OutgoingMessage) -> Result<Self::TransformedMessage>;
+}

src/balancer/inference_service/http_route/api/ws_inference_socket/client/message.rs renamed to src/balancer/inference_client/message.rs

@@ -1,29 +1,15 @@
-use actix_web::http::header;
-use log::error;
 use actix_web::post;
 use actix_web::error::ErrorBadRequest;
-use actix_web::rt;
 use actix_web::web;
 use actix_web::Error;
-use actix_web::HttpResponse;
 use actix_web::Responder;
-use bytes::Bytes;
-use futures::stream::StreamExt;
-use nanoid::nanoid;
-use tokio::sync::broadcast;
-use tokio::sync::mpsc;
-use tokio_stream::wrappers::UnboundedReceiverStream;
 
 use crate::validates::Validates as _;
 use crate::balancer::inference_service::app_data::AppData;
-use crate::balancer::inference_service::chunk_forwarding_session_controller::ChunkForwardingSessionController;
-use crate::balancer::request_from_agent::request_from_agent;
 use crate::request_params::ContinueFromConversationHistoryParams;
 use crate::request_params::continue_from_conversation_history_params::tool::tool_params::function_call::parameters_schema::raw_parameters_schema::RawParametersSchema;
-use crate::jsonrpc::Error as JsonRpcError;
-use crate::controls_session::ControlsSession as _;
-use crate::jsonrpc::ErrorEnvelope;
-use crate::balancer::inference_service::http_route::api::ws_inference_socket::client::Message as OutgoingMessage;
+use crate::balancer::stream_from_agent::stream_from_agent;
+use crate::balancer::chunk_forwarding_session_controller::identity_transformer::IdentityTransformer;
 
 pub fn register(cfg: &mut web::ServiceConfig) {
     cfg.service(respond);
@@ -34,50 +20,18 @@ async fn respond(
     app_data: web::Data<AppData>,
     params: web::Json<ContinueFromConversationHistoryParams<RawParametersSchema>>,
 ) -> Result<impl Responder, Error> {
-    let request_id: String = nanoid!();
-    let (connection_close_tx, _connection_close_rx) = broadcast::channel(1);
-    let (chunk_tx, chunk_rx) = mpsc::unbounded_channel();
-
-    let validated_params = match params.into_inner().validate() {
-        Ok(validated_params) => validated_params,
-        Err(validation_error) => {
-            return Err(ErrorBadRequest(format!(
-                "Invalid request parameters: {validation_error}"
-            )));
-        }
-    };
-
-    rt::spawn(async move {
-        let mut session_controller = ChunkForwardingSessionController { chunk_tx };
-
-        if let Err(err) = request_from_agent(
-            app_data.buffered_request_manager.clone(),
-            connection_close_tx,
-            app_data.inference_service_configuration.clone(),
-            validated_params,
-            request_id.clone(),
-            session_controller.clone(),
-        )
-        .await
-        {
-            error!("Failed to handle request: {err}");
-            session_controller
-                .send_response_safe(OutgoingMessage::Error(ErrorEnvelope {
-                    request_id: request_id.clone(),
-                    error: JsonRpcError {
-                        code: 500,
-                        description: format!("Request {request_id} failed: {err}"),
-                    },
-                }))
-                .await;
-        }
-    });
-
-    let stream = UnboundedReceiverStream::new(chunk_rx)
-        .map(|chunk: String| Ok::<_, Error>(Bytes::from(format!("{chunk}\n"))));
-
-    Ok(HttpResponse::Ok()
-        .insert_header(header::ContentType::json())
-        .insert_header((header::CACHE_CONTROL, "no-cache"))
-        .streaming(stream))
+    stream_from_agent(
+        app_data.buffered_request_manager.clone(),
+        app_data.inference_service_configuration.clone(),
+        match params.into_inner().validate() {
+            Ok(validated_params) => validated_params,
+            Err(validation_error) => {
+                return Err(ErrorBadRequest(format!(
+                    "Invalid request parameters: {validation_error}"
+                )));
+            }
+        },
+        IdentityTransformer::new(),
+    )
+    .await
 }

src/balancer/inference_service/http_route/api/ws_inference_socket/client/mod.rs renamed to src/balancer/inference_client/mod.rs

@@ -1,25 +1,11 @@
 use actix_web::Error;
-use actix_web::HttpResponse;
 use actix_web::Responder;
-use actix_web::http::header;
 use actix_web::post;
-use actix_web::rt;
 use actix_web::web;
-use bytes::Bytes;
-use futures::stream::StreamExt;
-use log::error;
-use nanoid::nanoid;
-use tokio::sync::broadcast;
-use tokio::sync::mpsc;
-use tokio_stream::wrappers::UnboundedReceiverStream;
 
+use crate::balancer::chunk_forwarding_session_controller::identity_transformer::IdentityTransformer;
 use crate::balancer::inference_service::app_data::AppData;
-use crate::balancer::inference_service::chunk_forwarding_session_controller::ChunkForwardingSessionController;
-use crate::balancer::inference_service::http_route::api::ws_inference_socket::client::Message as OutgoingMessage;
-use crate::balancer::request_from_agent::request_from_agent;
-use crate::controls_session::ControlsSession as _;
-use crate::jsonrpc::Error as JsonRpcError;
-use crate::jsonrpc::ErrorEnvelope;
+use crate::balancer::stream_from_agent::stream_from_agent;
 use crate::request_params::ContinueFromRawPromptParams;
 
 pub fn register(cfg: &mut web::ServiceConfig) {
@@ -31,41 +17,11 @@ async fn respond(
     app_data: web::Data<AppData>,
     params: web::Json<ContinueFromRawPromptParams>,
 ) -> Result<impl Responder, Error> {
-    let request_id: String = nanoid!();
-    let (connection_close_tx, _connection_close_rx) = broadcast::channel(1);
-    let (chunk_tx, chunk_rx) = mpsc::unbounded_channel();
-
-    rt::spawn(async move {
-        let mut session_controller = ChunkForwardingSessionController { chunk_tx };
-
-        if let Err(err) = request_from_agent(
-            app_data.buffered_request_manager.clone(),
-            connection_close_tx,
-            app_data.inference_service_configuration.clone(),
-            params.into_inner(),
-            request_id.clone(),
-            session_controller.clone(),
-        )
-        .await
-        {
-            error!("Failed to handle request: {err}");
-            session_controller
-                .send_response_safe(OutgoingMessage::Error(ErrorEnvelope {
-                    request_id: request_id.clone(),
-                    error: JsonRpcError {
-                        code: 500,
-                        description: format!("Request {request_id} failed: {err}"),
-                    },
-                }))
-                .await;
-        }
-    });
-
-    let stream = UnboundedReceiverStream::new(chunk_rx)
-        .map(|chunk: String| Ok::<_, Error>(Bytes::from(format!("{chunk}\n"))));
-
-    Ok(HttpResponse::Ok()
-        .insert_header(header::ContentType::json())
-        .insert_header((header::CACHE_CONTROL, "no-cache"))
-        .streaming(stream))
+    stream_from_agent(
+        app_data.buffered_request_manager.clone(),
+        app_data.inference_service_configuration.clone(),
+        params.into_inner(),
+        IdentityTransformer::new(),
+    )
+    .await
 }

src/balancer/inference_service/http_route/api/ws_inference_socket/client/response.rs renamed to src/balancer/inference_client/response.rs

@@ -15,9 +15,10 @@ use tokio::sync::broadcast;
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
+use crate::balancer::chunk_forwarding_session_controller::ChunkForwardingSessionController;
+use crate::balancer::chunk_forwarding_session_controller::identity_transformer::IdentityTransformer;
+use crate::balancer::inference_client::Message as OutgoingMessage;
 use crate::balancer::inference_service::app_data::AppData;
-use crate::balancer::inference_service::chunk_forwarding_session_controller::ChunkForwardingSessionController;
-use crate::balancer::inference_service::http_route::api::ws_inference_socket::client::Message as OutgoingMessage;
 use crate::balancer::request_from_agent::request_from_agent;
 use crate::controls_session::ControlsSession as _;
 use crate::jsonrpc::Error as JsonRpcError;
@@ -67,9 +68,8 @@ async fn respond(
 
         rt::spawn(async move {
             let request_id: String = nanoid!();
-            let mut session_controller = ChunkForwardingSessionController {
-                chunk_tx: chunk_tx_clone,
-            };
+            let mut session_controller =
+                ChunkForwardingSessionController::new(chunk_tx_clone, IdentityTransformer::new());
 
             if let Err(err) = request_from_agent(
                 buffered_request_manager_clone,