From b6e8384c32579f2f97e135026eb16f2dc94f44cb Mon Sep 17 00:00:00 2001
From: Din <dinmukhamed.mailibay@gmail.com>
Date: Fri, 22 Nov 2024 10:34:23 -0800
Subject: [PATCH] fix: response model would overwrite request model

---
 app-server/src/traces/attributes.rs | 3 +--
 app-server/src/traces/utils.rs      | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/app-server/src/traces/attributes.rs b/app-server/src/traces/attributes.rs
index 7922560e..0c8dc578 100644
--- a/app-server/src/traces/attributes.rs
+++ b/app-server/src/traces/attributes.rs
@@ -1,7 +1,6 @@
 use std::collections::HashMap;
 
 use chrono::{DateTime, Utc};
-use serde_json::Value;
 use uuid::Uuid;
 
 use crate::db::trace::TraceType;
@@ -17,7 +16,7 @@ pub struct TraceAttributes {
     pub total_token_count: Option<i64>,
     pub input_cost: Option<f64>,
     pub output_cost: Option<f64>,
-    /// Total costis not calculated on this struct and must be set manually
+    /// Total cost is not calculated on this struct and must be set manually
     pub cost: Option<f64>,
     pub session_id: Option<String>,
     pub trace_type: Option<TraceType>,
diff --git a/app-server/src/traces/utils.rs b/app-server/src/traces/utils.rs
index 49e3ee51..39abdb8c 100644
--- a/app-server/src/traces/utils.rs
+++ b/app-server/src/traces/utils.rs
@@ -78,7 +78,7 @@ pub async fn get_llm_usage_for_span(
         input_cost,
         output_cost,
         total_cost,
-        response_model: attributes.request_model().clone(),
+        response_model: attributes.response_model().clone(),
         request_model: attributes.request_model().clone(),
         provider_name,
     }