generalize agent response handlers

mcharytoniuk · mcharytoniuk · commit da3fa1d68e85 · 2025-08-07T14:20:24.000+02:00
diff --git a/src/agent/continue_from_conversation_history_request.rs b/src/agent/continue_from_conversation_history_request.rs
@@ -16,11 +16,12 @@ pub struct ContinueFromConversationHistoryRequest {
 
 impl FromRequestParams for ContinueFromConversationHistoryRequest {
     type RequestParams = ContinueFromConversationHistoryParams;
+    type Response = GeneratedTokenEnvelope;
 
     fn from_request_params(
         params: Self::RequestParams,
+        generated_tokens_tx: mpsc::UnboundedSender<Self::Response>,
         generate_tokens_stop_rx: mpsc::UnboundedReceiver<()>,
-        generated_tokens_tx: mpsc::UnboundedSender<GeneratedTokenEnvelope>,
     ) -> Self {
         ContinueFromConversationHistoryRequest {
             generate_tokens_stop_rx,
diff --git a/src/agent/continue_from_raw_prompt_request.rs b/src/agent/continue_from_raw_prompt_request.rs
@@ -16,11 +16,12 @@ pub struct ContinueFromRawPromptRequest {
 
 impl FromRequestParams for ContinueFromRawPromptRequest {
     type RequestParams = ContinueFromRawPromptParams;
+    type Response = GeneratedTokenEnvelope;
 
     fn from_request_params(
         params: Self::RequestParams,
+        generated_tokens_tx: mpsc::UnboundedSender<Self::Response>,
         generate_tokens_stop_rx: mpsc::UnboundedReceiver<()>,
-        generated_tokens_tx: mpsc::UnboundedSender<GeneratedTokenEnvelope>,
     ) -> Self {
         ContinueFromRawPromptRequest {
             generate_tokens_stop_rx,
diff --git a/src/agent/from_request_params.rs b/src/agent/from_request_params.rs
@@ -1,13 +1,14 @@
 use tokio::sync::mpsc;
 
-use crate::generated_token_envelope::GeneratedTokenEnvelope;
+use crate::agent::jsonrpc::response::Response;
 
 pub trait FromRequestParams: Send + Sync {
     type RequestParams;
+    type Response: Into<Response>;
 
     fn from_request_params(
-        request_params: Self::RequestParams,
-        generate_tokens_stop_rx: mpsc::UnboundedReceiver<()>,
-        generated_tokens_tx: mpsc::UnboundedSender<GeneratedTokenEnvelope>,
+        params: Self::RequestParams,
+        response_tx: mpsc::UnboundedSender<Self::Response>,
+        stop_rx: mpsc::UnboundedReceiver<()>,
     ) -> Self;
 }
diff --git a/src/agent/generate_embedding_batch_request.rs b/src/agent/generate_embedding_batch_request.rs
@@ -2,6 +2,7 @@ use actix::Message;
 use anyhow::Result;
 use tokio::sync::mpsc;
 
+use crate::agent::from_request_params::FromRequestParams;
 use crate::embedding_result::EmbeddingResult;
 use crate::request_params::GenerateEmbeddingBatchParams;
 
@@ -12,3 +13,20 @@ pub struct GenerateEmbeddingBatchRequest {
     pub generated_embedding_tx: mpsc::UnboundedSender<EmbeddingResult>,
     pub params: GenerateEmbeddingBatchParams,
 }
+
+impl FromRequestParams for GenerateEmbeddingBatchRequest {
+    type RequestParams = GenerateEmbeddingBatchParams;
+    type Response = EmbeddingResult;
+
+    fn from_request_params(
+        params: Self::RequestParams,
+        generated_embedding_tx: mpsc::UnboundedSender<Self::Response>,
+        generate_embedding_stop_rx: mpsc::UnboundedReceiver<()>,
+    ) -> Self {
+        GenerateEmbeddingBatchRequest {
+            generate_embedding_stop_rx,
+            generated_embedding_tx,
+            params,
+        }
+    }
+}
diff --git a/src/agent/jsonrpc/mod.rs b/src/agent/jsonrpc/mod.rs
@@ -2,7 +2,7 @@ mod message;
 mod notification;
 pub mod notification_params;
 mod request;
-mod response;
+pub mod response;
 
 pub use self::message::Message;
 pub use self::notification::Notification;
diff --git a/src/agent/jsonrpc/response.rs b/src/agent/jsonrpc/response.rs
@@ -13,3 +13,27 @@ pub enum Response {
     GeneratedToken(GeneratedTokenEnvelope),
     ModelMetadata(Option<ModelMetadata>),
 }
+
+impl From<Option<ChatTemplate>> for Response {
+    fn from(chat_template: Option<ChatTemplate>) -> Self {
+        Response::ChatTemplateOverride(chat_template)
+    }
+}
+
+impl From<EmbeddingResult> for Response {
+    fn from(embedding_result: EmbeddingResult) -> Self {
+        Response::Embedding(embedding_result)
+    }
+}
+
+impl From<GeneratedTokenEnvelope> for Response {
+    fn from(generated_token_envelope: GeneratedTokenEnvelope) -> Self {
+        Response::GeneratedToken(generated_token_envelope)
+    }
+}
+
+impl From<Option<ModelMetadata>> for Response {
+    fn from(model_metadata: Option<ModelMetadata>) -> Self {
+        Response::ModelMetadata(model_metadata)
+    }
+}
diff --git a/src/agent/management_socket_client_service.rs b/src/agent/management_socket_client_service.rs
@@ -20,7 +20,7 @@ use tokio_tungstenite::connect_async;
 use tokio_tungstenite::tungstenite::protocol::Message;
 
 use crate::agent_desired_state::AgentDesiredState;
-use crate::agent::receive_tokens_stopper_collection::ReceiveTokensStopperCollection;
+use crate::agent::receive_stream_stopper_collection::ReceiveStreamStopperCollection;
 use crate::agent::jsonrpc::Message as JsonRpcMessage;
 use crate::agent::jsonrpc::Notification as JsonRpcNotification;
 use crate::agent::jsonrpc::Request as JsonRpcRequest;
@@ -40,7 +40,6 @@ use crate::balancer::management_service::http_route::api::ws_agent_socket::jsonr
 use crate::jsonrpc::Error as JsonRpcError;
 use crate::jsonrpc::ResponseEnvelope;
 use crate::jsonrpc::RequestEnvelope;
-use crate::generated_token_envelope::GeneratedTokenEnvelope;
 use crate::produces_snapshot::ProducesSnapshot;
 use crate::jsonrpc::ErrorEnvelope;
 use crate::service::Service;
@@ -53,8 +52,9 @@ struct IncomingMessageContext {
     continue_from_conversation_history_request_tx:
         mpsc::UnboundedSender<ContinueFromConversationHistoryRequest>,
     continue_from_raw_prompt_request_tx: mpsc::UnboundedSender<ContinueFromRawPromptRequest>,
+    generate_embedding_batch_request_tx: mpsc::UnboundedSender<GenerateEmbeddingBatchRequest>,
     model_metadata_holder: Arc<ModelMetadataHolder>,
-    receive_tokens_stopper_collection: Arc<ReceiveTokensStopperCollection>,
+    receive_stream_stopper_collection: Arc<ReceiveStreamStopperCollection>,
     message_tx: mpsc::UnboundedSender<ManagementJsonRpcMessage>,
 }
 
@@ -67,47 +67,46 @@ pub struct ManagementSocketClientService {
     pub generate_embedding_batch_request_tx: mpsc::UnboundedSender<GenerateEmbeddingBatchRequest>,
     pub model_metadata_holder: Arc<ModelMetadataHolder>,
     pub name: Option<String>,
-    pub receive_tokens_stopper_collection: Arc<ReceiveTokensStopperCollection>,
+    pub receive_stream_stopper_collection: Arc<ReceiveStreamStopperCollection>,
     pub slot_aggregated_status: Arc<SlotAggregatedStatus>,
     pub socket_url: String,
 }
 
 impl ManagementSocketClientService {
-    async fn generate_tokens<TRequest: FromRequestParams + 'static>(
+    async fn generate_responses<TRequest: FromRequestParams + 'static>(
         connection_close_tx: broadcast::Sender<()>,
         id: String,
         message_tx: mpsc::UnboundedSender<ManagementJsonRpcMessage>,
         request_params: TRequest::RequestParams,
-        receive_tokens_stopper_collection: Arc<ReceiveTokensStopperCollection>,
+        receive_stream_stopper_collection: Arc<ReceiveStreamStopperCollection>,
         request_tx: mpsc::UnboundedSender<TRequest>,
     ) -> Result<()> {
-        let (generated_tokens_tx, mut generated_tokens_rx) =
-            mpsc::unbounded_channel::<GeneratedTokenEnvelope>();
-        let (generate_tokens_stop_tx, generate_tokens_stop_rx) = mpsc::unbounded_channel::<()>();
+        let (response_tx, mut response_rx) = mpsc::unbounded_channel::<TRequest::Response>();
+        let (stop_tx, stop_rx) = mpsc::unbounded_channel::<()>();
 
-        let _guard = receive_tokens_stopper_collection
-            .register_stopper_with_guard(id.clone(), generate_tokens_stop_tx)
-            .context(format!("Failed to register stopper for request ID: {id}"))?;
+        let _guard = receive_stream_stopper_collection
+            .register_stopper_with_guard(id.clone(), stop_tx)
+            .context(format!("Failed to register stopper for request: {id}"))?;
 
         request_tx.send(TRequest::from_request_params(
             request_params,
-            generate_tokens_stop_rx,
-            generated_tokens_tx,
+            response_tx,
+            stop_rx,
         ))?;
 
         let mut connection_close_rx = connection_close_tx.subscribe();
 
         loop {
             tokio::select! {
                 _ = connection_close_rx.recv() => break,
-                generated_token_envelope = generated_tokens_rx.recv() => {
-                    match generated_token_envelope {
-                        Some(generated_token_envelope) => {
+                response = response_rx.recv() => {
+                    match response {
+                        Some(response) => {
                             message_tx.send(
                                 ManagementJsonRpcMessage::Response(
                                     ResponseEnvelope {
                                         request_id: id.clone(),
-                                        response: JsonRpcResponse::GeneratedToken(generated_token_envelope),
+                                        response: response.into(),
                                     }
                                 ),
                             )?;
@@ -128,9 +127,10 @@ impl ManagementSocketClientService {
             connection_close_tx,
             continue_from_conversation_history_request_tx,
             continue_from_raw_prompt_request_tx,
+            generate_embedding_batch_request_tx,
             message_tx,
             model_metadata_holder,
-            receive_tokens_stopper_collection,
+            receive_stream_stopper_collection,
         }: IncomingMessageContext,
         deserialized_message: JsonRpcMessage,
     ) -> Result<()> {
@@ -154,7 +154,7 @@ impl ManagementSocketClientService {
             }
             JsonRpcMessage::Notification(JsonRpcNotification::StopGeneratingTokens(request_id)) => {
                 debug!("Received StopGeneratingTokens notification for request ID: {request_id:?}");
-                receive_tokens_stopper_collection
+                receive_stream_stopper_collection
                     .stop(request_id.clone())
                     .context(format!(
                         "Failed to stop generating tokens for request ID: {request_id}"
@@ -181,12 +181,12 @@ impl ManagementSocketClientService {
                         continue_from_conversation_history_params,
                     ),
             }) => {
-                Self::generate_tokens(
+                Self::generate_responses(
                     connection_close_tx,
                     id,
                     message_tx,
                     continue_from_conversation_history_params,
-                    receive_tokens_stopper_collection,
+                    receive_stream_stopper_collection,
                     continue_from_conversation_history_request_tx,
                 )
                 .await
@@ -195,20 +195,30 @@ impl ManagementSocketClientService {
                 id,
                 request: JsonRpcRequest::ContinueFromRawPrompt(generate_tokens_params),
             }) => {
-                Self::generate_tokens(
+                Self::generate_responses(
                     connection_close_tx,
                     id,
                     message_tx,
                     generate_tokens_params,
-                    receive_tokens_stopper_collection,
+                    receive_stream_stopper_collection,
                     continue_from_raw_prompt_request_tx,
                 )
                 .await
             }
             JsonRpcMessage::Request(RequestEnvelope {
                 id,
                 request: JsonRpcRequest::GenerateEmbeddingBatch(generate_embedding_batch_params),
-            }) => Ok(()),
+            }) => {
+                Self::generate_responses(
+                    connection_close_tx,
+                    id,
+                    message_tx,
+                    generate_embedding_batch_params,
+                    receive_stream_stopper_collection,
+                    generate_embedding_batch_request_tx,
+                )
+                .await
+            }
             JsonRpcMessage::Request(RequestEnvelope {
                 id,
                 request: JsonRpcRequest::GetChatTemplateOverride,
@@ -429,8 +439,9 @@ impl ManagementSocketClientService {
                                         connection_close_tx: connection_close_tx.clone(),
                                         continue_from_conversation_history_request_tx: self.continue_from_conversation_history_request_tx.clone(),
                                         continue_from_raw_prompt_request_tx: self.continue_from_raw_prompt_request_tx.clone(),
+                                        generate_embedding_batch_request_tx: self.generate_embedding_batch_request_tx.clone(),
                                         model_metadata_holder: self.model_metadata_holder.clone(),
-                                        receive_tokens_stopper_collection: self.receive_tokens_stopper_collection.clone(),
+                                        receive_stream_stopper_collection: self.receive_stream_stopper_collection.clone(),
                                         message_tx: message_tx.clone(),
                                     },
                                     msg,
diff --git a/src/agent/mod.rs b/src/agent/mod.rs
@@ -11,6 +11,6 @@ mod llamacpp_slot;
 mod llamacpp_slot_context;
 pub mod management_socket_client_service;
 pub mod model_metadata_holder;
-pub mod receive_tokens_stopper_collection;
-mod receive_tokens_stopper_drop_guard;
+pub mod receive_stream_stopper_collection;
+mod receive_stream_stopper_drop_guard;
 pub mod reconciliation_service;
diff --git a/src/agent/receive_stream_stopper_collection.rs b/src/agent/receive_stream_stopper_collection.rs
@@ -5,21 +5,21 @@ use anyhow::Result;
 use dashmap::DashMap;
 use tokio::sync::mpsc;
 
-use crate::agent::receive_tokens_stopper_drop_guard::ReceiveTokensStopperDropGuard;
+use crate::agent::receive_stream_stopper_drop_guard::ReceiveStreamStopperDropGuard;
 
-pub struct ReceiveTokensStopperCollection {
-    receive_tokens_stoppers: DashMap<String, mpsc::UnboundedSender<()>>,
+pub struct ReceiveStreamStopperCollection {
+    receive_stoppers: DashMap<String, mpsc::UnboundedSender<()>>,
 }
 
-impl ReceiveTokensStopperCollection {
+impl ReceiveStreamStopperCollection {
     pub fn new() -> Self {
         Self {
-            receive_tokens_stoppers: DashMap::new(),
+            receive_stoppers: DashMap::new(),
         }
     }
 
     pub fn deregister_stopper(&self, request_id: String) -> Result<()> {
-        if let Some(stopper) = self.receive_tokens_stoppers.remove(&request_id) {
+        if let Some(stopper) = self.receive_stoppers.remove(&request_id) {
             drop(stopper);
 
             Ok(())
@@ -33,13 +33,13 @@ impl ReceiveTokensStopperCollection {
         request_id: String,
         stopper: mpsc::UnboundedSender<()>,
     ) -> Result<()> {
-        if self.receive_tokens_stoppers.contains_key(&request_id) {
+        if self.receive_stoppers.contains_key(&request_id) {
             return Err(anyhow!(
                 "Stopper for request_id {request_id} already exists"
             ));
         }
 
-        self.receive_tokens_stoppers.insert(request_id, stopper);
+        self.receive_stoppers.insert(request_id, stopper);
 
         Ok(())
     }
@@ -48,17 +48,17 @@ impl ReceiveTokensStopperCollection {
         self: &Arc<Self>,
         request_id: String,
         stopper: mpsc::UnboundedSender<()>,
-    ) -> Result<ReceiveTokensStopperDropGuard> {
+    ) -> Result<ReceiveStreamStopperDropGuard> {
         self.register_stopper(request_id.clone(), stopper)?;
 
-        Ok(ReceiveTokensStopperDropGuard {
-            receive_tokens_stopper_collection: self.clone(),
+        Ok(ReceiveStreamStopperDropGuard {
+            receive_stream_stopper_collection: self.clone(),
             request_id,
         })
     }
 
     pub fn stop(&self, request_id: String) -> Result<()> {
-        if let Some(stopper) = self.receive_tokens_stoppers.get(&request_id) {
+        if let Some(stopper) = self.receive_stoppers.get(&request_id) {
             stopper.send(())?;
 
             Ok(())
diff --git a/src/agent/receive_stream_stopper_drop_guard.rs b/src/agent/receive_stream_stopper_drop_guard.rs
@@ -2,17 +2,17 @@ use std::sync::Arc;
 
 use log::error;
 
-use crate::agent::receive_tokens_stopper_collection::ReceiveTokensStopperCollection;
+use crate::agent::receive_stream_stopper_collection::ReceiveStreamStopperCollection;
 
-pub struct ReceiveTokensStopperDropGuard {
-    pub receive_tokens_stopper_collection: Arc<ReceiveTokensStopperCollection>,
+pub struct ReceiveStreamStopperDropGuard {
+    pub receive_stream_stopper_collection: Arc<ReceiveStreamStopperCollection>,
     pub request_id: String,
 }
 
-impl Drop for ReceiveTokensStopperDropGuard {
+impl Drop for ReceiveStreamStopperDropGuard {
     fn drop(&mut self) {
         if let Err(err) = self
-            .receive_tokens_stopper_collection
+            .receive_stream_stopper_collection
             .deregister_stopper(self.request_id.clone())
         {
             error!(
diff --git a/src/balancer/agent_controller.rs b/src/balancer/agent_controller.rs
diff --git a/src/cmd/agent.rs b/src/cmd/agent.rs
