add embedding result

mcharytoniuk · mcharytoniuk · commit ee0786257a8d · 2025-08-07T00:08:45.000+02:00
diff --git a/src/agent/jsonrpc/response.rs b/src/agent/jsonrpc/response.rs
@@ -2,12 +2,14 @@ use serde::Deserialize;
 use serde::Serialize;
 
 use crate::chat_template::ChatTemplate;
+use crate::embedding_result::EmbeddingResult;
 use crate::generated_token_envelope::GeneratedTokenEnvelope;
 use crate::model_metadata::ModelMetadata;
 
 #[derive(Deserialize, Serialize)]
 pub enum Response {
     ChatTemplateOverride(Option<ChatTemplate>),
+    Embedding(EmbeddingResult),
     GeneratedToken(GeneratedTokenEnvelope),
     ModelMetadata(Option<ModelMetadata>),
 }
diff --git a/src/balancer/agent_controller.rs b/src/balancer/agent_controller.rs
@@ -22,14 +22,17 @@ use crate::atomic_value::AtomicValue;
 use crate::balancer::agent_controller_snapshot::AgentControllerSnapshot;
 use crate::balancer::agent_controller_update_result::AgentControllerUpdateResult;
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
+use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
 use crate::balancer::manages_senders_controller::ManagesSendersController;
 use crate::balancer::model_metadata_sender_collection::ModelMetadataSenderCollection;
 use crate::balancer::receive_tokens_controller::ReceiveTokensController;
+use crate::embedding_result::EmbeddingResult;
 use crate::jsonrpc::RequestEnvelope;
 use crate::produces_snapshot::ProducesSnapshot;
 use crate::request_params::ContinueFromConversationHistoryParams;
 use crate::request_params::ContinueFromRawPromptParams;
+use crate::request_params::GenerateEmbeddingBatchParams;
 use crate::sends_rpc_message::SendsRpcMessage;
 use crate::sets_desired_state::SetsDesiredState;
 use crate::slot_aggregated_status_snapshot::SlotAggregatedStatusSnapshot;
@@ -42,6 +45,7 @@ pub struct AgentController {
     pub download_current: AtomicValue<AtomicUsize>,
     pub download_filename: RwLock<Option<String>>,
     pub download_total: AtomicValue<AtomicUsize>,
+    pub embedding_sender_collection: Arc<EmbeddingSenderCollection>,
     pub generate_tokens_sender_collection: Arc<GenerateTokensSenderCollection>,
     pub id: String,
     pub issues: RwLock<BTreeSet<AgentIssue>>,
@@ -59,15 +63,13 @@ impl AgentController {
     pub async fn continue_from_conversation_history(
         &self,
         request_id: String,
-        continue_from_conversation_history_params: ContinueFromConversationHistoryParams,
+        params: ContinueFromConversationHistoryParams,
     ) -> Result<ReceiveTokensController> {
         self.receiver_from_message(
             request_id.clone(),
             AgentJsonRpcMessage::Request(RequestEnvelope {
                 id: request_id,
-                request: AgentJsonRpcRequest::ContinueFromConversationHistory(
-                    continue_from_conversation_history_params.clone(),
-                ),
+                request: AgentJsonRpcRequest::ContinueFromConversationHistory(params.clone()),
             }),
         )
         .await
@@ -76,20 +78,32 @@ impl AgentController {
     pub async fn continue_from_raw_prompt(
         &self,
         request_id: String,
-        continue_from_raw_prompt_params: ContinueFromRawPromptParams,
+        params: ContinueFromRawPromptParams,
     ) -> Result<ReceiveTokensController> {
         self.receiver_from_message(
             request_id.clone(),
             AgentJsonRpcMessage::Request(RequestEnvelope {
                 id: request_id,
-                request: AgentJsonRpcRequest::ContinueFromRawPrompt(
-                    continue_from_raw_prompt_params.clone(),
-                ),
+                request: AgentJsonRpcRequest::ContinueFromRawPrompt(params.clone()),
             }),
         )
         .await
     }
 
+    pub async fn generate_embedding_batch(
+        &self,
+        embedding_rx: mpsc::UnboundedReceiver<EmbeddingResult>,
+        embedding_tx: mpsc::UnboundedSender<EmbeddingResult>,
+        request_id: String,
+        params: GenerateEmbeddingBatchParams,
+    ) -> Result<ManagesSendersController<EmbeddingSenderCollection>> {
+        Ok(ManagesSendersController {
+            request_id,
+            response_rx: embedding_rx,
+            response_sender_collection: self.embedding_sender_collection.clone(),
+        })
+    }
+
     pub async fn get_chat_template_override(
         &self,
     ) -> Result<ManagesSendersController<ChatTemplateOverrideSenderCollection>> {
diff --git a/src/balancer/embedding_sender_collection.rs b/src/balancer/embedding_sender_collection.rs
@@ -0,0 +1,27 @@
+use async_trait::async_trait;
+use dashmap::DashMap;
+use tokio::sync::mpsc;
+
+use crate::balancer::manages_senders::ManagesSenders;
+use crate::embedding_result::EmbeddingResult;
+
+pub struct EmbeddingSenderCollection {
+    senders: DashMap<String, mpsc::UnboundedSender<EmbeddingResult>>,
+}
+
+impl EmbeddingSenderCollection {
+    pub fn new() -> Self {
+        Self {
+            senders: DashMap::new(),
+        }
+    }
+}
+
+#[async_trait]
+impl ManagesSenders for EmbeddingSenderCollection {
+    type Value = EmbeddingResult;
+
+    fn get_sender_collection(&self) -> &DashMap<String, mpsc::UnboundedSender<Self::Value>> {
+        &self.senders
+    }
+}
diff --git a/src/balancer/inference_service/http_route/api/post_generate_embedding_batch.rs b/src/balancer/inference_service/http_route/api/post_generate_embedding_batch.rs
@@ -7,21 +7,17 @@ use actix_web::web;
 use actix_web::Error;
 use actix_web::HttpResponse;
 use actix_web::Responder;
-use async_trait::async_trait;
 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::inference_service::app_data::AppData;
-use crate::balancer::inference_service::chunk_forwarding_session_controller::ChunkForwardingSessionController;
-use crate::balancer::inference_service::controls_inference_endpoint::ControlsInferenceEndpoint;
 use crate::request_params::GenerateEmbeddingBatchParams;
 
-const CHARACTERS_PER_TOKEN_MORE_OR_LESS: usize = 4;
+const CHARACTERS_PER_TOKEN_APPROXIMATELY: usize = 3;
 
 pub fn register(cfg: &mut web::ServiceConfig) {
     cfg.service(respond);
@@ -48,12 +44,12 @@ async fn respond(
         ));
     }
 
-    let request_id: String = nanoid!();
     let (connection_close_tx, mut connection_close_rx) = broadcast::channel::<()>(1);
     let (chunk_tx, chunk_rx) = mpsc::unbounded_channel();
 
     let request_batches = params.chunk_by_input_size(
-        agent_desired_state.inference_parameters.batch_n_tokens * CHARACTERS_PER_TOKEN_MORE_OR_LESS,
+        agent_desired_state.inference_parameters.batch_n_tokens
+            * CHARACTERS_PER_TOKEN_APPROXIMATELY,
     );
 
     rt::spawn(async move {});
diff --git a/src/balancer/inference_service/mod.rs b/src/balancer/inference_service/mod.rs
@@ -80,6 +80,7 @@ impl Service for InferenceService {
                 .configure(common_http_route::get_health::register)
                 .configure(http_route::api::post_continue_from_conversation_history::register)
                 .configure(http_route::api::post_continue_from_raw_prompt::register)
+                .configure(http_route::api::post_generate_embedding_batch::register)
                 .configure(http_route::api::ws_inference_socket::register)
         })
         .shutdown_signal(async move {
diff --git a/src/balancer/management_service/app_data.rs b/src/balancer/management_service/app_data.rs
@@ -3,6 +3,7 @@ use std::sync::Arc;
 use crate::balancer::agent_controller_pool::AgentControllerPool;
 use crate::balancer::buffered_request_manager::BufferedRequestManager;
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
+use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
 use crate::balancer::model_metadata_sender_collection::ModelMetadataSenderCollection;
 use crate::balancer::state_database::StateDatabase;
@@ -13,6 +14,7 @@ pub struct AppData {
     pub balancer_applicable_state_holder: Arc<BalancerApplicableStateHolder>,
     pub buffered_request_manager: Arc<BufferedRequestManager>,
     pub chat_template_override_sender_collection: Arc<ChatTemplateOverrideSenderCollection>,
+    pub embedding_sender_collection: Arc<EmbeddingSenderCollection>,
     pub generate_tokens_sender_collection: Arc<GenerateTokensSenderCollection>,
     pub model_metadata_sender_collection: Arc<ModelMetadataSenderCollection>,
     pub state_database: Arc<dyn StateDatabase>,
diff --git a/src/balancer/management_service/http_route/api/ws_agent_socket/agent_socket_controller_context.rs b/src/balancer/management_service/http_route/api/ws_agent_socket/agent_socket_controller_context.rs
@@ -5,15 +5,17 @@ use log::info;
 
 use crate::balancer::agent_controller_pool::AgentControllerPool;
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
+use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
 use crate::balancer::model_metadata_sender_collection::ModelMetadataSenderCollection;
 use crate::balancer_applicable_state_holder::BalancerApplicableStateHolder;
 
 pub struct AgentSocketControllerContext {
-    pub balancer_applicable_state_holder: Arc<BalancerApplicableStateHolder>,
     pub agent_controller_pool: Arc<AgentControllerPool>,
     pub agent_id: String,
+    pub balancer_applicable_state_holder: Arc<BalancerApplicableStateHolder>,
     pub chat_template_override_sender_collection: Arc<ChatTemplateOverrideSenderCollection>,
+    pub embedding_sender_collection: Arc<EmbeddingSenderCollection>,
     pub generate_tokens_sender_collection: Arc<GenerateTokensSenderCollection>,
     pub model_metadata_sender_collection: Arc<ModelMetadataSenderCollection>,
 }
diff --git a/src/balancer/management_service/http_route/api/ws_agent_socket/mod.rs b/src/balancer/management_service/http_route/api/ws_agent_socket/mod.rs
@@ -41,6 +41,7 @@ use crate::balancer::agent_controller::AgentController;
 use crate::balancer::agent_controller_pool::AgentControllerPool;
 use crate::balancer::agent_controller_update_result::AgentControllerUpdateResult;
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
+use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
 use crate::balancer::management_service::app_data::AppData;
 use crate::balancer::manages_senders::ManagesSenders as _;
@@ -63,6 +64,7 @@ struct AgentSocketController {
     agent_id: String,
     balancer_applicable_state_holder: Arc<BalancerApplicableStateHolder>,
     chat_template_override_sender_collection: Arc<ChatTemplateOverrideSenderCollection>,
+    embedding_sender_collection: Arc<EmbeddingSenderCollection>,
     generate_tokens_sender_collection: Arc<GenerateTokensSenderCollection>,
     model_metadata_sender_collection: Arc<ModelMetadataSenderCollection>,
 }
@@ -81,6 +83,7 @@ impl ControlsWebSocketEndpoint for AgentSocketController {
             chat_template_override_sender_collection: self
                 .chat_template_override_sender_collection
                 .clone(),
+            embedding_sender_collection: self.embedding_sender_collection.clone(),
             generate_tokens_sender_collection: self.generate_tokens_sender_collection.clone(),
             model_metadata_sender_collection: self.model_metadata_sender_collection.clone(),
         }
@@ -136,6 +139,7 @@ impl ControlsWebSocketEndpoint for AgentSocketController {
                     download_current: AtomicValue::<AtomicUsize>::new(download_current),
                     download_filename: RwLock::new(download_filename),
                     download_total: AtomicValue::<AtomicUsize>::new(download_total),
+                    embedding_sender_collection: context.embedding_sender_collection.clone(),
                     generate_tokens_sender_collection: context
                         .generate_tokens_sender_collection
                         .clone(),
@@ -241,6 +245,17 @@ impl ControlsWebSocketEndpoint for AgentSocketController {
 
                 Ok(ContinuationDecision::Continue)
             }
+            ManagementJsonRpcMessage::Response(ResponseEnvelope {
+                request_id,
+                response: AgentJsonRpcResponse::Embedding(embedding_result),
+            }) => {
+                context
+                    .embedding_sender_collection
+                    .forward_response_safe(request_id, embedding_result)
+                    .await;
+
+                Ok(ContinuationDecision::Continue)
+            }
             ManagementJsonRpcMessage::Response(ResponseEnvelope {
                 request_id,
                 response: AgentJsonRpcResponse::GeneratedToken(generated_token_envelope),
@@ -310,6 +325,7 @@ async fn respond(
         chat_template_override_sender_collection: app_data
             .chat_template_override_sender_collection
             .clone(),
+        embedding_sender_collection: app_data.embedding_sender_collection.clone(),
         generate_tokens_sender_collection: app_data.generate_tokens_sender_collection.clone(),
         model_metadata_sender_collection: app_data.model_metadata_sender_collection.clone(),
     };
diff --git a/src/balancer/management_service/mod.rs b/src/balancer/management_service/mod.rs
@@ -15,6 +15,7 @@ use tokio::sync::broadcast;
 use crate::balancer::agent_controller_pool::AgentControllerPool;
 use crate::balancer::buffered_request_manager::BufferedRequestManager;
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
+use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
 use crate::balancer::http_route as common_http_route;
 use crate::balancer::management_service::app_data::AppData;
@@ -33,6 +34,7 @@ pub struct ManagementService {
     pub buffered_request_manager: Arc<BufferedRequestManager>,
     pub chat_template_override_sender_collection: Arc<ChatTemplateOverrideSenderCollection>,
     pub configuration: ManagementServiceConfiguration,
+    pub embedding_sender_collection: Arc<EmbeddingSenderCollection>,
     pub generate_tokens_sender_collection: Arc<GenerateTokensSenderCollection>,
     pub model_metadata_sender_collection: Arc<ModelMetadataSenderCollection>,
     pub state_database: Arc<dyn StateDatabase>,
@@ -64,6 +66,7 @@ impl Service for ManagementService {
             chat_template_override_sender_collection: self
                 .chat_template_override_sender_collection
                 .clone(),
+            embedding_sender_collection: self.embedding_sender_collection.clone(),
             generate_tokens_sender_collection: self.generate_tokens_sender_collection.clone(),
             model_metadata_sender_collection: self.model_metadata_sender_collection.clone(),
             state_database: self.state_database.clone(),
diff --git a/src/balancer/mod.rs b/src/balancer/mod.rs
@@ -11,6 +11,7 @@ pub mod buffered_request_manager;
 mod buffered_request_manager_snapshot;
 pub mod chat_template_override_sender_collection;
 mod controls_manages_senders_endpoint;
+pub mod embedding_sender_collection;
 pub mod generate_tokens_sender_collection;
 mod http_route;
 pub mod inference_service;
diff --git a/src/cmd/balancer.rs b/src/cmd/balancer.rs
@@ -14,6 +14,7 @@ use super::parse_socket_addr;
 use crate::balancer::agent_controller_pool::AgentControllerPool;
 use crate::balancer::buffered_request_manager::BufferedRequestManager;
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
+use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
 use crate::balancer::inference_service::configuration::Configuration as InferenceServiceConfiguration;
 use crate::balancer::inference_service::InferenceService;
@@ -138,6 +139,7 @@ impl Handler for Balancer {
         ));
         let chat_template_override_sender_collection =
             Arc::new(ChatTemplateOverrideSenderCollection::new());
+        let embedding_sender_collection = Arc::new(EmbeddingSenderCollection::new());
         let generate_tokens_sender_collection = Arc::new(GenerateTokensSenderCollection::new());
         let model_metadata_sender_collection = Arc::new(ModelMetadataSenderCollection::new());
         let mut service_manager = ServiceManager::new();
@@ -167,7 +169,8 @@ impl Handler for Balancer {
             buffered_request_manager: buffered_request_manager.clone(),
             chat_template_override_sender_collection,
             configuration: self.get_management_service_configuration(),
-            generate_tokens_sender_collection: generate_tokens_sender_collection.clone(),
+            embedding_sender_collection,
+            generate_tokens_sender_collection,
             model_metadata_sender_collection,
             state_database: state_database.clone(),
             #[cfg(feature = "web_admin_panel")]
@@ -184,7 +187,7 @@ impl Handler for Balancer {
 
         if let Some(statsd_addr) = self.statsd_addr {
             service_manager.add_service(StatsdService {
-                agent_controller_pool: agent_controller_pool.clone(),
+                agent_controller_pool,
                 buffered_request_manager,
                 configuration: StatsdServiceConfiguration {
                     statsd_addr,
diff --git a/src/embedding.rs b/src/embedding.rs
@@ -1,36 +1,13 @@
 use anyhow::anyhow;
 use anyhow::Result;
+use serde::Deserialize;
+use serde::Serialize;
 
 use crate::embedding_normalization_method::EmbeddingNormalizationMethod;
+use crate::normalization::l2;
+use crate::normalization::rms_norm;
 
-fn l2(embedding: &[f32]) -> Vec<f32> {
-    let magnitude = embedding
-        .iter()
-        .fold(0.0, |acc, &val| val.mul_add(val, acc))
-        .sqrt();
-
-    if magnitude == 0.0 {
-        return vec![0.0; embedding.len()];
-    }
-
-    embedding.iter().map(|&val| val / magnitude).collect()
-}
-
-fn rms_norm(embedding: &[f32], eps: f32) -> Vec<f32> {
-    let mean_square = embedding
-        .iter()
-        .fold(0.0, |acc, &val| val.mul_add(val, acc))
-        / embedding.len() as f32;
-
-    let rms = (mean_square + eps).sqrt();
-
-    if rms == 0.0 {
-        return vec![0.0; embedding.len()];
-    }
-
-    embedding.iter().map(|&val| val / rms).collect()
-}
-
+#[derive(Debug, Deserialize, Serialize)]
 pub struct Embedding {
     pub embedding: Vec<f32>,
     pub normalization_method: EmbeddingNormalizationMethod,
@@ -69,21 +46,3 @@ impl Embedding {
         })
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_normalize_l2() {
-        let embedding = vec![3.0, 4.0];
-        let normalized = l2(&embedding);
-
-        assert_eq!(normalized, vec![0.6, 0.8]);
-
-        let zero_embedding = vec![0.0, 0.0];
-        let normalized_zero = l2(&zero_embedding);
-
-        assert_eq!(normalized_zero, vec![0.0, 0.0]);
-    }
-}
diff --git a/src/embedding_result.rs b/src/embedding_result.rs
diff --git a/src/main.rs b/src/main.rs
diff --git a/src/normalization/l2.rs b/src/normalization/l2.rs
diff --git a/src/normalization/mod.rs b/src/normalization/mod.rs
diff --git a/src/normalization/rms_norm.rs b/src/normalization/rms_norm.rs
