metric and dashoard context fix

api/libs/handlers/src/chats/context_loaders/dashboard_context.rs

@@ -1,19 +1,19 @@
-use anyhow::{Result, anyhow};
+use agents::{Agent, AgentMessage};
+use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 use database::{
-    models::{User, Dataset, MetricFile},
+    models::{Dataset, MetricFile},
     pool::get_pg_pool,
-    schema::{dashboard_files, metric_files, datasets},
+    schema::{dashboard_files, datasets, metric_files},
 };
 use diesel::prelude::*;
 use diesel_async::RunQueryDsl;
-use agents::{AgentMessage, Agent};
 use litellm::MessageProgress;
 use middleware::AuthenticatedUser;
 use serde_json::Value;
+use std::collections::HashSet;
 use std::sync::Arc;
 use uuid::Uuid;
-use std::collections::HashSet;
 
 use super::ContextLoader;
 
@@ -29,9 +29,16 @@ impl DashboardContextLoader {
 
 #[async_trait]
 impl ContextLoader for DashboardContextLoader {
-    async fn load_context(&self, user: &AuthenticatedUser, agent: &Arc<Agent>) -> Result<Vec<AgentMessage>> {
+    async fn load_context(
+        &self,
+        user: &AuthenticatedUser,
+        agent: &Arc<Agent>,
+    ) -> Result<Vec<AgentMessage>> {
         let mut conn = get_pg_pool().get().await.map_err(|e| {
-            anyhow!("Failed to get database connection for dashboard context loading: {}", e)
+            anyhow!(
+                "Failed to get database connection for dashboard context loading: {}",
+                e
+            )
         })?;
 
         // First verify the dashboard exists and user has access
@@ -46,9 +53,7 @@ impl ContextLoader for DashboardContextLoader {
             })?;
 
         // Parse dashboard content to DashboardYml
-        let dashboard_yml: database::types::DashboardYml = 
+        let dashboard_yml = dashboard.content;
-            serde_json::from_value(dashboard.content.clone())
-            .map_err(|e| anyhow!("Failed to parse dashboard content as YAML for dashboard {}: {}", dashboard.name, e))?;
 
         // Collect all metric IDs from the dashboard
         let mut metric_ids = HashSet::new();
@@ -71,16 +76,9 @@ impl ContextLoader for DashboardContextLoader {
             {
                 Ok(metric) => {
                     // Parse metric content
-                    match serde_json::from_value::<database::types::MetricYml>(metric.content.clone()) {
+                    all_dataset_ids.extend(metric.content.dataset_ids.clone());
-                        Ok(metric_yml) => {
-                            all_dataset_ids.extend(metric_yml.dataset_ids);
                     metrics_vec.push(metric);
                 }
-                        Err(e) => {
-                            failed_metric_loads.push((metric_id, format!("Failed to parse metric content: {}", e)));
-                        }
-                    }
-                }
                 Err(e) => {
                     failed_metric_loads.push((metric_id, format!("Failed to load metric: {}", e)));
                 }
@@ -119,25 +117,34 @@ impl ContextLoader for DashboardContextLoader {
         }
 
         // Set agent state based on loaded assets
-        agent.set_state_value(String::from("dashboards_available"), Value::Bool(true))
+        agent
+            .set_state_value(String::from("dashboards_available"), Value::Bool(true))
             .await;
 
-        agent.set_state_value(String::from("files_available"), Value::Bool(true))
+        agent
+            .set_state_value(String::from("files_available"), Value::Bool(true))
             .await;
 
         if !metrics_vec.is_empty() {
-            agent.set_state_value(String::from("metrics_available"), Value::Bool(true))
+            agent
+                .set_state_value(String::from("metrics_available"), Value::Bool(true))
                 .await;
         };
 
         if !datasets_vec.is_empty() {
-            agent.set_state_value(String::from("data_context"), Value::Bool(true))
+            agent
+                .set_state_value(String::from("data_context"), Value::Bool(true))
                 .await;
         };
 
         // Format the context message with dashboard, metrics, and dataset information
-        let dashboard_yaml = serde_yaml::to_string(&dashboard_yml)
+        let dashboard_yaml = serde_yaml::to_string(&dashboard_yml).map_err(|e| {
-            .map_err(|e| anyhow!("Failed to serialize dashboard {} to YAML: {}", dashboard.name, e))?;
+            anyhow!(
+                "Failed to serialize dashboard {} to YAML: {}",
+                dashboard.name,
+                e
+            )
+        })?;
 
         let mut context_message = format!(
             "This conversation is continuing with context from the dashboard. Here is the relevant information:\n\nDashboard Definition:\n{}\n\n",
@@ -147,15 +154,10 @@ impl ContextLoader for DashboardContextLoader {
         if !metrics_vec.is_empty() {
             context_message.push_str("Referenced Metrics:\n");
             for metric in metrics_vec {
-                match serde_json::from_value::<database::types::MetricYml>(metric.content) {
+                context_message.push_str(&format!(
-                    Ok(metric_yml) => {
+                    "\n{}\n",
-                        match serde_yaml::to_string(&metric_yml) {
+                    serde_yaml::to_string(&metric.content).unwrap()
-                            Ok(yaml) => context_message.push_str(&format!("\n{}\n", yaml)),
+                ));
-                            Err(e) => tracing::warn!("Failed to serialize metric {} to YAML: {}", metric.id, e),
-                        }
-                    }
-                    Err(e) => tracing::warn!("Failed to parse metric {} content: {}", metric.id, e),
-                }
             }
         }
 

api/libs/handlers/src/chats/context_loaders/metric_context.rs

@@ -47,15 +47,7 @@ impl ContextLoader for MetricContextLoader {
             })?;
 
         // Get the metric content as MetricYml
-        let metric_yml: database::types::MetricYml =
+        let metric_yml = metric.content;
-            serde_json::from_value(metric.content.clone()).map_err(|e| {
-                anyhow!(
-                    "Failed to parse metric content as YAML for metric {}: {}",
-                    metric.name,
-                    e
-                )
-            })?;
-
         // Load all referenced datasets
         let dataset_ids = &metric_yml.dataset_ids;
         let mut datasets_vec = Vec::new();