add to collections

This commit is contained in:
dal 2025-03-20 09:57:01 -06:00
parent b8e4a95316
commit 5b484510c3
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
4 changed files with 70 additions and 45 deletions

View File

@ -558,6 +558,7 @@ pub async fn process_metric_file(
tool_call_id: String, tool_call_id: String,
file_name: String, file_name: String,
yml_content: String, yml_content: String,
user_id: &Uuid,
) -> Result< ) -> Result<
( (
MetricFile, MetricFile,
@ -599,7 +600,7 @@ pub async fn process_metric_file(
name: file_name.clone(), name: file_name.clone(),
file_name: file_name.clone(), file_name: file_name.clone(),
content: metric_yml.clone(), content: metric_yml.clone(),
created_by: Uuid::new_v4(), created_by: user_id.clone(),
verification: Verification::NotRequested, verification: Verification::NotRequested,
evaluation_obj: None, evaluation_obj: None,
evaluation_summary: None, evaluation_summary: None,

View File

@ -71,6 +71,7 @@ impl FileModificationTool for CreateDashboardFilesTool {}
async fn process_dashboard_file( async fn process_dashboard_file(
tool_call_id: String, tool_call_id: String,
file: DashboardFileParams, file: DashboardFileParams,
user_id: &Uuid,
) -> Result<(DashboardFile, DashboardYml), String> { ) -> Result<(DashboardFile, DashboardYml), String> {
debug!("Processing dashboard file creation: {}", file.name); debug!("Processing dashboard file creation: {}", file.name);
@ -106,7 +107,7 @@ async fn process_dashboard_file(
content: dashboard_yml.clone(), content: dashboard_yml.clone(),
filter: None, filter: None,
organization_id: Uuid::new_v4(), organization_id: Uuid::new_v4(),
created_by: Uuid::new_v4(), created_by: user_id.clone(),
created_at: Utc::now(), created_at: Utc::now(),
updated_at: Utc::now(), updated_at: Utc::now(),
deleted_at: None, deleted_at: None,
@ -152,7 +153,7 @@ impl ToolExecutor for CreateDashboardFilesTool {
// First pass - validate and prepare all records // First pass - validate and prepare all records
for file in files { for file in files {
match process_dashboard_file(tool_call_id.clone(), file.clone()).await { match process_dashboard_file(tool_call_id.clone(), file.clone(), &self.agent.get_user_id()).await {
Ok((dashboard_file, dashboard_yml)) => { Ok((dashboard_file, dashboard_yml)) => {
dashboard_records.push(dashboard_file); dashboard_records.push(dashboard_file);
dashboard_ymls.push(dashboard_yml); dashboard_ymls.push(dashboard_yml);

View File

@ -5,10 +5,10 @@ use async_trait::async_trait;
use braintrust::{get_prompt_system_message, BraintrustClient}; use braintrust::{get_prompt_system_message, BraintrustClient};
use chrono::Utc; use chrono::Utc;
use database::{ use database::{
pool::get_pg_pool, enums::{AssetPermissionRole, AssetType, IdentityType},
schema::{metric_files, asset_permissions},
models::AssetPermission, models::AssetPermission,
enums::{AssetType, IdentityType, AssetPermissionRole}, pool::get_pg_pool,
schema::{asset_permissions, metric_files},
}; };
use diesel::insert_into; use diesel::insert_into;
use diesel_async::RunQueryDsl; use diesel_async::RunQueryDsl;
@ -23,13 +23,16 @@ use crate::{
tools::{ tools::{
file_tools::{ file_tools::{
common::{process_metric_file, METRIC_YML_SCHEMA}, common::{process_metric_file, METRIC_YML_SCHEMA},
file_types::{file::FileWithId}, file_types::file::FileWithId,
}, },
ToolExecutor, ToolExecutor,
}, },
}; };
use super::{common::{validate_sql, generate_deterministic_uuid}, FileModificationTool}; use super::{
common::{generate_deterministic_uuid, validate_sql},
FileModificationTool,
};
#[derive(Debug, Serialize, Deserialize, Clone)] #[derive(Debug, Serialize, Deserialize, Clone)]
pub struct MetricFileParams { pub struct MetricFileParams {
@ -99,7 +102,14 @@ impl ToolExecutor for CreateMetricFilesTool {
let mut results_vec = vec![]; let mut results_vec = vec![];
// First pass - validate and prepare all records // First pass - validate and prepare all records
for file in files { for file in files {
match process_metric_file(tool_call_id.clone(), file.name.clone(), file.yml_content.clone()).await { match process_metric_file(
tool_call_id.clone(),
file.name.clone(),
file.yml_content.clone(),
&self.agent.get_user_id(),
)
.await
{
Ok((metric_file, metric_yml, message, results)) => { Ok((metric_file, metric_yml, message, results)) => {
metric_records.push(metric_file); metric_records.push(metric_file);
metric_ymls.push(metric_yml); metric_ymls.push(metric_yml);
@ -153,8 +163,11 @@ impl ToolExecutor for CreateMetricFilesTool {
.await .await
{ {
Ok(_) => { Ok(_) => {
tracing::debug!("Successfully inserted asset permissions for {} metric files", asset_permissions.len()); tracing::debug!(
}, "Successfully inserted asset permissions for {} metric files",
asset_permissions.len()
);
}
Err(e) => { Err(e) => {
tracing::error!("Error inserting asset permissions: {}", e); tracing::error!("Error inserting asset permissions: {}", e);
// Continue with the process even if permissions failed // Continue with the process even if permissions failed

View File

@ -1,12 +1,21 @@
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
use database::{collections::fetch_collection, enums::{AssetPermissionRole, AssetType, IdentityType}, pool::get_pg_pool, schema::{asset_permissions, collections, collections_to_assets, dashboard_files, metric_files, users}}; use database::{
use diesel::{ExpressionMethods, JoinOnDsl, QueryDsl, Queryable}; collections::fetch_collection,
enums::{AssetPermissionRole, AssetType, IdentityType},
pool::get_pg_pool,
schema::{
asset_permissions, collections, collections_to_assets, dashboard_files, metric_files, users,
},
};
use diesel::{ExpressionMethods, JoinOnDsl, NullableExpressionMethods, QueryDsl, Queryable};
use diesel_async::RunQueryDsl; use diesel_async::RunQueryDsl;
use tracing; use tracing;
use uuid::Uuid; use uuid::Uuid;
use crate::collections::types::{AssetUser, BusterShareIndividual, CollectionAsset, CollectionState, GetCollectionRequest}; use crate::collections::types::{
AssetUser, BusterShareIndividual, CollectionAsset, CollectionState, GetCollectionRequest,
};
#[derive(Queryable)] #[derive(Queryable)]
struct AssetPermissionInfo { struct AssetPermissionInfo {
@ -22,7 +31,7 @@ struct AssetQueryResult {
id: Uuid, id: Uuid,
name: String, name: String,
user_name: Option<String>, user_name: Option<String>,
email: String, email: Option<String>,
created_at: DateTime<Utc>, created_at: DateTime<Utc>,
updated_at: DateTime<Utc>, updated_at: DateTime<Utc>,
asset_type: AssetType, asset_type: AssetType,
@ -41,18 +50,16 @@ struct AssetQueryResult {
fn format_assets(assets: Vec<AssetQueryResult>) -> Vec<CollectionAsset> { fn format_assets(assets: Vec<AssetQueryResult>) -> Vec<CollectionAsset> {
assets assets
.into_iter() .into_iter()
.map(|asset| { .map(|asset| CollectionAsset {
CollectionAsset { id: asset.id,
id: asset.id, name: asset.name,
name: asset.name, created_by: AssetUser {
created_by: AssetUser { name: asset.user_name,
name: asset.user_name, email: asset.email.unwrap_or("chad@buster.so".to_string()),
email: asset.email, },
}, created_at: asset.created_at,
created_at: asset.created_at, updated_at: asset.updated_at,
updated_at: asset.updated_at, asset_type: asset.asset_type,
asset_type: asset.asset_type,
}
}) })
.collect() .collect()
} }
@ -90,7 +97,8 @@ pub async fn get_collection_handler(
// For collections, we'll default public fields to false/none // For collections, we'll default public fields to false/none
// since the schema doesn't have these fields yet // since the schema doesn't have these fields yet
let public_info: Result<(bool, Option<Uuid>, Option<DateTime<Utc>>), anyhow::Error> = Ok((false, None, None)); let public_info: Result<(bool, Option<Uuid>, Option<DateTime<Utc>>), anyhow::Error> =
Ok((false, None, None));
// Convert AssetPermissionInfo to BusterShareIndividual // Convert AssetPermissionInfo to BusterShareIndividual
let individual_permissions = match individual_permissions_query { let individual_permissions = match individual_permissions_query {
@ -136,7 +144,7 @@ pub async fn get_collection_handler(
// Query for metric assets in the collection // Query for metric assets in the collection
let metric_assets_result = collections_to_assets::table let metric_assets_result = collections_to_assets::table
.inner_join(metric_files::table.on(metric_files::id.eq(collections_to_assets::asset_id))) .inner_join(metric_files::table.on(metric_files::id.eq(collections_to_assets::asset_id)))
.inner_join(users::table.on(users::id.eq(metric_files::created_by))) .left_join(users::table.on(users::id.eq(metric_files::created_by)))
.filter(collections_to_assets::collection_id.eq(req.id)) .filter(collections_to_assets::collection_id.eq(req.id))
.filter(collections_to_assets::asset_type.eq(AssetType::MetricFile)) .filter(collections_to_assets::asset_type.eq(AssetType::MetricFile))
.filter(collections_to_assets::deleted_at.is_null()) .filter(collections_to_assets::deleted_at.is_null())
@ -144,8 +152,8 @@ pub async fn get_collection_handler(
.select(( .select((
metric_files::id, metric_files::id,
metric_files::name, metric_files::name,
users::name, users::name.nullable(),
users::email, users::email.nullable(),
metric_files::created_at, metric_files::created_at,
metric_files::updated_at, metric_files::updated_at,
collections_to_assets::asset_type, collections_to_assets::asset_type,
@ -155,8 +163,10 @@ pub async fn get_collection_handler(
// Query for dashboard assets in the collection // Query for dashboard assets in the collection
let dashboard_assets_result = collections_to_assets::table let dashboard_assets_result = collections_to_assets::table
.inner_join(dashboard_files::table.on(dashboard_files::id.eq(collections_to_assets::asset_id))) .inner_join(
.inner_join(users::table.on(users::id.eq(dashboard_files::created_by))) dashboard_files::table.on(dashboard_files::id.eq(collections_to_assets::asset_id)),
)
.left_join(users::table.on(users::id.eq(dashboard_files::created_by)))
.filter(collections_to_assets::collection_id.eq(req.id)) .filter(collections_to_assets::collection_id.eq(req.id))
.filter(collections_to_assets::asset_type.eq(AssetType::DashboardFile)) .filter(collections_to_assets::asset_type.eq(AssetType::DashboardFile))
.filter(collections_to_assets::deleted_at.is_null()) .filter(collections_to_assets::deleted_at.is_null())
@ -164,8 +174,8 @@ pub async fn get_collection_handler(
.select(( .select((
dashboard_files::id, dashboard_files::id,
dashboard_files::name, dashboard_files::name,
users::name, users::name.nullable(),
users::email, users::email.nullable(),
dashboard_files::created_at, dashboard_files::created_at,
dashboard_files::updated_at, dashboard_files::updated_at,
collections_to_assets::asset_type, collections_to_assets::asset_type,