send embeddings back through channel

Author: mcharytoniuk

src/agent/generate_embedding_batch_request.rs

@@ -1,10 +1,14 @@
 use actix::Message;
 use anyhow::Result;
+use tokio::sync::mpsc;
 
+use crate::embedding_result::EmbeddingResult;
 use crate::request_params::GenerateEmbeddingBatchParams;
 
 #[derive(Debug, Message)]
 #[rtype(result = "Result<()>")]
 pub struct GenerateEmbeddingBatchRequest {
+    pub generate_embedding_stop_rx: mpsc::UnboundedReceiver<()>,
+    pub generated_embedding_tx: mpsc::UnboundedSender<EmbeddingResult>,
     pub params: GenerateEmbeddingBatchParams,
 }

src/agent/jsonrpc/request.rs

@@ -3,11 +3,13 @@ use serde::Serialize;
 
 use crate::request_params::ContinueFromConversationHistoryParams;
 use crate::request_params::ContinueFromRawPromptParams;
+use crate::request_params::GenerateEmbeddingBatchParams;
 
 #[derive(Deserialize, Serialize)]
 pub enum Request {
     ContinueFromConversationHistory(ContinueFromConversationHistoryParams),
     ContinueFromRawPrompt(ContinueFromRawPromptParams),
+    GenerateEmbeddingBatch(GenerateEmbeddingBatchParams),
     GetChatTemplateOverride,
     GetModelMetadata,
 }

src/agent/llamacpp_arbiter_service.rs

@@ -19,6 +19,7 @@ use tokio::time::MissedTickBehavior;
 
 use crate::agent::continue_from_conversation_history_request::ContinueFromConversationHistoryRequest;
 use crate::agent::continue_from_raw_prompt_request::ContinueFromRawPromptRequest;
+use crate::agent::generate_embedding_batch_request::GenerateEmbeddingBatchRequest;
 use crate::agent::llamacpp_arbiter::LlamaCppArbiter;
 use crate::agent::llamacpp_arbiter_handle::LlamaCppArbiterHandle;
 use crate::agent::llamacpp_slot::LlamaCppSlot;
@@ -39,6 +40,7 @@ pub struct LlamaCppArbiterService {
         mpsc::UnboundedReceiver<ContinueFromConversationHistoryRequest>,
     pub continue_from_raw_prompt_request_rx: mpsc::UnboundedReceiver<ContinueFromRawPromptRequest>,
     pub desired_slots_total: i32,
+    pub generate_embedding_batch_request_rx: mpsc::UnboundedReceiver<GenerateEmbeddingBatchRequest>,
     pub llamacpp_arbiter_handle: Option<LlamaCppArbiterHandle>,
     pub model_metadata_holder: Arc<ModelMetadataHolder>,
     pub slot_aggregated_status_manager: Arc<SlotAggregatedStatusManager>,
@@ -143,7 +145,7 @@ impl LlamaCppArbiterService {
         Ok(())
     }
 
-    async fn generate_tokens<TRequest>(
+    async fn forward_request_to_arbiter<TRequest>(
         &mut self,
         request: TRequest,
         mut shutdown: broadcast::Receiver<()>,
@@ -158,11 +160,11 @@ impl LlamaCppArbiterService {
             rt::spawn(async move {
                 tokio::select! {
                     _ = shutdown.recv() => {
-                        error!("Shutdown received, stopping ContinueFromRawPromptRequest processing");
+                        error!("Shutdown received, stopping request processing");
                     }
                     result = llamacpp_slot_addr.send(request) => {
                         if let Err(err) = result {
-                            error!("Failed to send ContinueFromRawPromptRequest: {err}");
+                            error!("Failed to forward request to arbiter: {err}");
                         }
                     }
                 }
@@ -222,7 +224,7 @@ impl Service for LlamaCppArbiterService {
                 continue_from_conversation_history_request = self.continue_from_conversation_history_request_rx.recv() => {
                     match continue_from_conversation_history_request {
                         Some(continue_from_conversation_history_request) => {
-                            self.generate_tokens(
+                            self.forward_request_to_arbiter(
                                 continue_from_conversation_history_request,
                                 shutdown.resubscribe(),
                             ).await
@@ -235,7 +237,7 @@ impl Service for LlamaCppArbiterService {
                 continue_from_raw_prompt_request = self.continue_from_raw_prompt_request_rx.recv() => {
                     match continue_from_raw_prompt_request {
                         Some(continue_from_raw_prompt_request) => {
-                            self.generate_tokens(
+                            self.forward_request_to_arbiter(
                                 continue_from_raw_prompt_request,
                                 shutdown.resubscribe(),
                             ).await
@@ -245,6 +247,19 @@ impl Service for LlamaCppArbiterService {
                         }
                     }
                 }
+                generate_embedding_batch_request = self.generate_embedding_batch_request_rx.recv() => {
+                    match generate_embedding_batch_request {
+                        Some(generate_embedding_batch_request) => {
+                            self.forward_request_to_arbiter(
+                                generate_embedding_batch_request,
+                                shutdown.resubscribe(),
+                            ).await
+                        }
+                        None => {
+                            break Err(anyhow!("GenerateEmbeddingBatchRequest channel closed unexpectedly"));
+                        }
+                    }
+                }
             }
         }
     }

src/agent/llamacpp_slot.rs

@@ -31,6 +31,7 @@ use crate::agent::llamacpp_slot_context::LlamaCppSlotContext;
 use crate::embedding::Embedding;
 use crate::embedding_input_tokenized::EmbeddingInputTokenized;
 use crate::embedding_normalization_method::EmbeddingNormalizationMethod;
+use crate::embedding_result::EmbeddingResult;
 use crate::generated_token_envelope::GeneratedTokenEnvelope;
 use crate::generated_token_result::GeneratedTokenResult;
 use crate::request_params::ContinueFromConversationHistoryParams;
@@ -123,8 +124,8 @@ impl LlamaCppSlot {
         &mut self,
         batch: &mut LlamaBatch,
         current_batch_embeddings: &Vec<&EmbeddingInputTokenized>,
+        generated_embedding_tx: &mpsc::UnboundedSender<EmbeddingResult>,
         normalization_method: &EmbeddingNormalizationMethod,
-        output: &mut Vec<Embedding>,
     ) -> Result<()> {
         self.llama_context.clear_kv_cache();
         self.llama_context.decode(batch)?;
@@ -135,14 +136,14 @@ impl LlamaCppSlot {
                 .embeddings_seq_ith(index as i32)
                 .context("Failed to get embeddings")?;
 
-            output.push(
+            generated_embedding_tx.send(EmbeddingResult::Embedding(
                 Embedding {
                     embedding: embedding.to_vec(),
                     normalization_method: EmbeddingNormalizationMethod::None,
                     source_document_id: embedding_input_tokenized.id.clone(),
                 }
                 .normalize(normalization_method)?,
-            );
+            ))?;
         }
 
         batch.clear();
@@ -369,6 +370,8 @@ impl Handler<GenerateEmbeddingBatchRequest> for LlamaCppSlot {
     fn handle(
         &mut self,
         GenerateEmbeddingBatchRequest {
+            generate_embedding_stop_rx,
+            generated_embedding_tx,
             params:
                 GenerateEmbeddingBatchParams {
                     input_batch,
@@ -408,7 +411,6 @@ impl Handler<GenerateEmbeddingBatchRequest> for LlamaCppSlot {
 
         let mut batch = LlamaBatch::new(self.slot_context.inference_parameters.batch_n_tokens, 1);
         let mut current_batch_embeddings: Vec<&EmbeddingInputTokenized> = Vec::new();
-        let mut output = Vec::with_capacity(tokens_lines_list.len());
 
         for embedding_input_tokenized in &tokens_lines_list {
             // Flush the batch if the next prompt would exceed our batch size
@@ -418,8 +420,8 @@ impl Handler<GenerateEmbeddingBatchRequest> for LlamaCppSlot {
                 self.embedding_batch_decode(
                     &mut batch,
                     &current_batch_embeddings,
+                    &generated_embedding_tx,
                     &normalization_method,
-                    &mut output,
                 )?;
 
                 current_batch_embeddings.clear();
@@ -436,8 +438,8 @@ impl Handler<GenerateEmbeddingBatchRequest> for LlamaCppSlot {
         self.embedding_batch_decode(
             &mut batch,
             &current_batch_embeddings,
+            &generated_embedding_tx,
             &normalization_method,
-            &mut output,
         )?;
 
         Ok(())

src/agent/management_socket_client_service.rs

@@ -34,6 +34,7 @@ use crate::agent::from_request_params::FromRequestParams;
 use crate::balancer::management_service::http_route::api::ws_agent_socket::jsonrpc::Message as ManagementJsonRpcMessage;
 use crate::balancer::management_service::http_route::api::ws_agent_socket::jsonrpc::Notification as ManagementJsonRpcNotification;
 use crate::agent::continue_from_conversation_history_request::ContinueFromConversationHistoryRequest;
+use crate::agent::generate_embedding_batch_request::GenerateEmbeddingBatchRequest;
 use crate::balancer::management_service::http_route::api::ws_agent_socket::jsonrpc::notification_params::RegisterAgentParams;
 use crate::balancer::management_service::http_route::api::ws_agent_socket::jsonrpc::notification_params::UpdateAgentStatusParams;
 use crate::jsonrpc::Error as JsonRpcError;
@@ -63,6 +64,7 @@ pub struct ManagementSocketClientService {
     pub continue_from_conversation_history_request_tx:
         mpsc::UnboundedSender<ContinueFromConversationHistoryRequest>,
     pub continue_from_raw_prompt_request_tx: mpsc::UnboundedSender<ContinueFromRawPromptRequest>,
+    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>,
@@ -203,6 +205,10 @@ impl ManagementSocketClientService {
                 )
                 .await
             }
+            JsonRpcMessage::Request(RequestEnvelope {
+                id,
+                request: JsonRpcRequest::GenerateEmbeddingBatch(generate_embedding_batch_params),
+            }) => Ok(()),
             JsonRpcMessage::Request(RequestEnvelope {
                 id,
                 request: JsonRpcRequest::GetChatTemplateOverride,

src/cmd/agent.rs

@@ -12,6 +12,7 @@ use super::handler::Handler;
 use super::parse_socket_addr;
 use crate::agent::continue_from_conversation_history_request::ContinueFromConversationHistoryRequest;
 use crate::agent::continue_from_raw_prompt_request::ContinueFromRawPromptRequest;
+use crate::agent::generate_embedding_batch_request::GenerateEmbeddingBatchRequest;
 use crate::agent::llamacpp_arbiter_service::LlamaCppArbiterService;
 use crate::agent::management_socket_client_service::ManagementSocketClientService;
 use crate::agent::model_metadata_holder::ModelMetadataHolder;
@@ -47,6 +48,8 @@ impl Handler for Agent {
         ) = mpsc::unbounded_channel::<ContinueFromConversationHistoryRequest>();
         let (continue_from_raw_prompt_request_tx, continue_from_raw_prompt_request_rx) =
             mpsc::unbounded_channel::<ContinueFromRawPromptRequest>();
+        let (generate_embedding_batch_request_tx, generate_embedding_batch_request_rx) =
+            mpsc::unbounded_channel::<GenerateEmbeddingBatchRequest>();
 
         let agent_applicable_state_holder = Arc::new(AgentApplicableStateHolder::new());
         let model_metadata_holder = Arc::new(ModelMetadataHolder::new());
@@ -60,6 +63,7 @@ impl Handler for Agent {
             continue_from_conversation_history_request_rx,
             continue_from_raw_prompt_request_rx,
             desired_slots_total: self.slots,
+            generate_embedding_batch_request_rx,
             llamacpp_arbiter_handle: None,
             model_metadata_holder: model_metadata_holder.clone(),
             slot_aggregated_status_manager: slot_aggregated_status_manager.clone(),
@@ -70,6 +74,7 @@ impl Handler for Agent {
             agent_desired_state_tx,
             continue_from_conversation_history_request_tx,
             continue_from_raw_prompt_request_tx,
+            generate_embedding_batch_request_tx,
             model_metadata_holder,
             name: self.name.clone(),
             receive_tokens_stopper_collection: Arc::new(ReceiveTokensStopperCollection::new()),