use axum::{ body::Body, extract::{Path, State}, http::{header, HeaderMap, HeaderValue, StatusCode}, response::IntoResponse, }; use futures::try_join; use serde::Serialize; use crate::{ html_state::HtmlState, middlewares::{ auth_middleware::RequireUser, response_middleware::{HtmlError, TemplateResponse}, }, utils::text_content_preview::truncate_text_contents, AuthSessionType, }; use common::storage::store; use common::storage::types::user::DashboardStats; use common::{ error::AppError, storage::types::{ conversation::Conversation, file_info::FileInfo, ingestion_task::IngestionTask, knowledge_entity::KnowledgeEntity, knowledge_relationship::KnowledgeRelationship, text_chunk::TextChunk, text_content::TextContent, user::User, }, }; #[derive(Serialize)] pub struct IndexPageData { user: Option, text_contents: Vec, stats: DashboardStats, active_jobs: Vec, conversation_archive: Vec, } pub async fn index_handler( State(state): State, auth: AuthSessionType, ) -> Result { let Some(user) = auth.current_user else { return Ok(TemplateResponse::redirect("/signin")); }; let (text_contents, conversation_archive, stats, active_jobs) = try_join!( User::get_latest_text_contents(&user.id, &state.db), User::get_user_conversations(&user.id, &state.db), User::get_dashboard_stats(&user.id, &state.db), User::get_unfinished_ingestion_tasks(&user.id, &state.db) )?; let text_contents = truncate_text_contents(text_contents); Ok(TemplateResponse::new_template( "dashboard/base.html", IndexPageData { user: Some(user), text_contents, stats, conversation_archive, active_jobs, }, )) } #[derive(Serialize)] pub struct LatestTextContentData { text_contents: Vec, user: User, } pub async fn delete_text_content( State(state): State, RequireUser(user): RequireUser, Path(id): Path, ) -> Result { // Get and validate TextContent let text_content = get_and_validate_text_content(&state, &id, &user).await?; // Remove stored assets before deleting the text content record if let Some(file_info) = text_content.file_info.as_ref() { let file_in_use = TextContent::has_other_with_file(&file_info.id, &text_content.id, &state.db).await?; if !file_in_use { FileInfo::delete_by_id(&file_info.id, &state.db, &state.config).await?; } } // Delete the text content and any related data TextChunk::delete_by_source_id(&text_content.id, &state.db).await?; KnowledgeEntity::delete_by_source_id(&text_content.id, &state.db).await?; KnowledgeRelationship::delete_relationships_by_source_id(&text_content.id, &state.db).await?; state .db .delete_item::(&text_content.id) .await?; // Render updated content let text_contents = truncate_text_contents(User::get_latest_text_contents(&user.id, &state.db).await?); Ok(TemplateResponse::new_partial( "dashboard/recent_content.html", "latest_content_section", LatestTextContentData { user: user.to_owned(), text_contents, }, )) } // Helper function to get and validate text content async fn get_and_validate_text_content( state: &HtmlState, id: &str, user: &User, ) -> Result { let text_content = state .db .get_item::(id) .await? .ok_or_else(|| AppError::NotFound("Item was not found".to_string()))?; if text_content.user_id != user.id { return Err(AppError::Auth( "You are not the owner of that content".to_string(), )); } Ok(text_content) } #[derive(Serialize)] pub struct ActiveJobsData { pub active_jobs: Vec, pub user: User, } pub async fn delete_job( State(state): State, RequireUser(user): RequireUser, Path(id): Path, ) -> Result { User::validate_and_delete_job(&id, &user.id, &state.db).await?; let active_jobs = User::get_unfinished_ingestion_tasks(&user.id, &state.db).await?; Ok(TemplateResponse::new_partial( "dashboard/active_jobs.html", "active_jobs_section", ActiveJobsData { user: user.clone(), active_jobs, }, )) } pub async fn show_active_jobs( State(state): State, RequireUser(user): RequireUser, ) -> Result { let active_jobs = User::get_unfinished_ingestion_tasks(&user.id, &state.db).await?; Ok(TemplateResponse::new_template( "dashboard/active_jobs.html", ActiveJobsData { user: user.clone(), active_jobs, }, )) } pub async fn serve_file( State(state): State, RequireUser(user): RequireUser, Path(file_id): Path, ) -> Result { let file_info = match FileInfo::get_by_id(&file_id, &state.db).await { Ok(info) => info, _ => return Ok(TemplateResponse::not_found().into_response()), }; if file_info.user_id != user.id { return Ok(TemplateResponse::unauthorized().into_response()); } let stream = match store::get_stream_at(&file_info.path, &state.config).await { Ok(s) => s, Err(_) => return Ok(TemplateResponse::server_error().into_response()), }; let body = Body::from_stream(stream); let mut headers = HeaderMap::new(); headers.insert( header::CONTENT_TYPE, HeaderValue::from_str(&file_info.mime_type) .unwrap_or_else(|_| HeaderValue::from_static("application/octet-stream")), ); let Ok(disposition_value) = HeaderValue::from_str(&format!("attachment; filename=\"{}\"", file_info.file_name)) else { headers.insert( header::CONTENT_DISPOSITION, HeaderValue::from_static("attachment"), ); return Ok((StatusCode::OK, headers, body).into_response()); }; headers.insert(header::CONTENT_DISPOSITION, disposition_value); headers.insert( header::CACHE_CONTROL, HeaderValue::from_static("private, max-age=31536000, immutable"), ); Ok((StatusCode::OK, headers, body).into_response()) }