From fe023d806554a4163eba8f27e7ed960a61f4021f Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil@katanemo.com>
Date: Mon, 4 Nov 2024 16:01:43 -0800
Subject: [PATCH] improve

---
 crates/prompt_gateway/src/filter_context.rs | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/crates/prompt_gateway/src/filter_context.rs b/crates/prompt_gateway/src/filter_context.rs
index 066f0072..0b44fac9 100644
--- a/crates/prompt_gateway/src/filter_context.rs
+++ b/crates/prompt_gateway/src/filter_context.rs
@@ -130,9 +130,9 @@ impl FilterContext {
 
     fn embedding_response_handler(
         &mut self,
-        body_size: usize,
         embedding_type: EmbeddingType,
         prompt_target_name: String,
+        body: Vec<u8>,
     ) {
         let prompt_target = self
             .prompt_targets
@@ -144,9 +144,6 @@ impl FilterContext {
                 )
             });
 
-        let body = self
-            .get_http_call_response_body(0, body_size)
-            .expect("No body in response");
         if !body.is_empty() {
             let mut embedding_response: CreateEmbeddingResponse =
                 match serde_json::from_slice(&body) {
@@ -225,17 +222,16 @@ impl Context for FilterContext {
             .remove(&token_id)
             .expect("invalid token_id");
 
-        let body_bytes = self.get_http_call_response_body(0, body_size).unwrap();
-
         self.active_embedding_calls_count -= 1;
         self.metrics.active_http_calls.increment(-1);
+        let body_bytes = self.get_http_call_response_body(0, body_size).unwrap();
 
         if let Some(status_code) = self.get_http_call_response_header(":status") {
             if status_code == StatusCode::OK.as_str() {
                 self.embedding_response_handler(
-                    body_size,
                     callout_data.embedding_type,
                     callout_data.prompt_target_name,
+                    body_bytes,
                 );
             } else {
                 warn!(