read llm costs from cache, TTL to 24 hours (#981)

* read llm costs from cache, TTL to 7 days

* read from DB even if error reading from cache, add error logs

* return optional from db on price entry

Author: dinmukhamedm

app-server/src/db/prices.rs

@@ -11,7 +11,11 @@ pub struct DBPriceEntry {
     pub additional_prices: Value,
 }
 
-pub async fn get_price(pool: &PgPool, provider: &str, model: &str) -> anyhow::Result<DBPriceEntry> {
+pub async fn get_price(
+    pool: &PgPool,
+    provider: &str,
+    model: &str,
+) -> anyhow::Result<Option<DBPriceEntry>> {
     let price = sqlx::query_as::<_, DBPriceEntry>(
         "SELECT
             provider,
@@ -32,7 +36,7 @@ pub async fn get_price(pool: &PgPool, provider: &str, model: &str) -> anyhow::Re
     )
     .bind(provider)
     .bind(model)
-    .fetch_one(pool)
+    .fetch_optional(pool)
     .await?;
 
     Ok(price)

app-server/src/language_model/costs/mod.rs

@@ -13,6 +13,8 @@ use utils::calculate_cost;
 
 mod utils;
 
+const LLM_PRICES_CACHE_TTL_SECONDS: u64 = 60 * 60 * 24; // 24 hours
+
 #[derive(Clone, Deserialize, Serialize)]
 pub struct LLMPriceEntry {
     _provider: String,
@@ -44,15 +46,30 @@ pub async fn estimate_output_cost(
     num_tokens: i64,
 ) -> Option<f64> {
     let cache_key = format!("{LLM_PRICES_CACHE_KEY}:{provider}:{model}");
-    let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await.ok()?;
+    let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await;
 
     let price_per_million_tokens = match cache_res {
-        Some(price) => price.output_price_per_million,
-        None => {
-            let price = get_price(&db.pool, provider, model).await.ok()?;
+        Ok(Some(price)) => price.output_price_per_million,
+        Ok(None) | Err(_) => {
+            let price = get_price(&db.pool, provider, model)
+                .await
+                .map_err(|e| {
+                    log::error!(
+                        "Error getting price from DB for provider: {}, model: {}: {:?}",
+                        provider,
+                        model,
+                        e
+                    );
+                    e
+                })
+                .unwrap_or_default()?;
             let price = LLMPriceEntry::from(price);
             let _ = cache
-                .insert::<LLMPriceEntry>(&cache_key, price.clone())
+                .insert_with_ttl::<LLMPriceEntry>(
+                    &cache_key,
+                    price.clone(),
+                    LLM_PRICES_CACHE_TTL_SECONDS,
+                )
                 .await;
             price.output_price_per_million
         }
@@ -68,16 +85,30 @@ pub async fn estimate_input_cost(
     input_tokens: InputTokens,
 ) -> Option<f64> {
     let cache_key = format!("{LLM_PRICES_CACHE_KEY}:{provider}:{model}");
-    // let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await.ok()?;
-    let cache_res = None;
+    let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await;
 
     let price = match cache_res {
-        Some(price) => price,
-        None => {
-            let price = get_price(&db.pool, provider, model).await.ok()?;
+        Ok(Some(price)) => price,
+        Ok(None) | Err(_) => {
+            let price = get_price(&db.pool, provider, model)
+                .await
+                .map_err(|e| {
+                    log::error!(
+                        "Error getting price from DB for provider: {}, model: {}: {:?}",
+                        provider,
+                        model,
+                        e
+                    );
+                    e
+                })
+                .unwrap_or_default()?;
             let price = LLMPriceEntry::from(price);
             let _ = cache
-                .insert::<LLMPriceEntry>(&cache_key, price.clone())
+                .insert_with_ttl::<LLMPriceEntry>(
+                    &cache_key,
+                    price.clone(),
+                    LLM_PRICES_CACHE_TTL_SECONDS,
+                )
                 .await;
             price
         }