use anyhow::{Context, Result}; use std::fs; use std::path::PathBuf; use walkdir::WalkDir; use crate::api::{ApiClient, RecallRequest, ReflectRequest, MemoryItem, RetainRequest}; use crate::config; use crate::output::{self, OutputFormat}; use crate::ui; // Import types from generated client use hindsight_client::types::{Budget, ChunkIncludeOptions, IncludeOptions, TagsMatch}; use serde::Deserialize; use serde_json; // Local types for serde_json::Value deserialization #[derive(Debug, Deserialize)] struct MemoryUnitDetail { id: String, text: String, #[serde(rename = "type")] type_: Option, document_id: Option, context: Option, occurred_start: Option, occurred_end: Option, entities: Option>, tags: Option>, } #[derive(Debug, Deserialize)] struct EntityRef { id: String, name: String, } // Helper function to parse budget string to Budget enum fn parse_budget(budget: &str) -> Budget { match budget.to_lowercase().as_str() { "low" => Budget::Low, "high" => Budget::High, _ => Budget::Mid, // Default to mid } } /// List memory units with pagination and optional filters pub fn list( client: &ApiClient, bank_id: &str, type_filter: Option, query: Option, limit: i64, offset: i64, verbose: bool, output_format: OutputFormat, ) -> Result<()> { let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Fetching memories...")) } else { None }; let response = client.list_memories( bank_id, type_filter.as_deref(), query.as_deref(), Some(limit), Some(offset), verbose, ); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if output_format == OutputFormat::Pretty { ui::print_section_header(&format!("Memories: {} (showing {}-{})", bank_id, offset + 1, offset + result.items.len() as i64)); if result.items.is_empty() { println!(" {}", ui::dim("No memories found.")); } else { for item in &result.items { let fact_type = item.get("type") .and_then(|v| v.as_str()) .unwrap_or("unknown"); let type_t = match fact_type { "world" => 0.0, "experience" => 0.5, "opinion" => 1.0, _ => 0.5, }; let id = item.get("id") .and_then(|v| v.as_str()) .unwrap_or("unknown"); println!( " {} {}", ui::gradient(&format!("[{}]", fact_type.to_uppercase()), type_t), ui::dim(id) ); // Truncate text if too long if let Some(text) = item.get("text").and_then(|v| v.as_str()) { let text_preview: String = text.chars().take(100).collect(); let ellipsis = if text.len() > 100 { "..." } else { "" }; println!(" {}{}", text_preview, ellipsis); } if let Some(doc_id) = item.get("document_id").and_then(|v| v.as_str()) { println!(" {} {}", ui::dim("doc:"), ui::dim(doc_id)); } println!(); } println!(" {} {} total", ui::dim("Total:"), result.total); } } else { output::print_output(&result, output_format)?; } Ok(()) } Err(e) => Err(e), } } /// Get a specific memory unit by ID pub fn get( client: &ApiClient, bank_id: &str, memory_id: &str, verbose: bool, output_format: OutputFormat, ) -> Result<()> { let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Fetching memory...")) } else { None }; let response = client.get_memory(bank_id, memory_id, verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(value) => { if output_format == OutputFormat::Pretty { let result: MemoryUnitDetail = serde_json::from_value(value) .with_context(|| "Failed to parse memory response")?; let fact_type = result.type_.as_deref().unwrap_or("unknown"); let type_t = match fact_type { "world" => 0.0, "experience" => 0.5, "opinion" => 1.0, _ => 0.5, }; ui::print_section_header(&format!("Memory: {}", memory_id)); println!(" {} {}", ui::dim("Type:"), ui::gradient(&fact_type.to_uppercase(), type_t)); println!(" {} {}", ui::dim("ID:"), result.id); if let Some(doc_id) = &result.document_id { println!(" {} {}", ui::dim("Document:"), doc_id); } if let Some(context) = &result.context { println!(" {} {}", ui::dim("Context:"), context); } println!(); println!("{}", ui::gradient_text("─── Content ───")); println!(); println!("{}", result.text); // Show temporal info if available if result.occurred_start.is_some() || result.occurred_end.is_some() { println!(); println!("{}", ui::gradient_text("─── Temporal ───")); if let Some(start) = &result.occurred_start { println!(" {} {}", ui::dim("Start:"), start); } if let Some(end) = &result.occurred_end { println!(" {} {}", ui::dim("End:"), end); } } // Show entities if available if let Some(entities) = &result.entities { if !entities.is_empty() { println!(); println!("{}", ui::gradient_text("─── Entities ───")); for entity in entities { println!(" • {} ({})", entity.name, entity.id); } } } // Show tags if available if let Some(tags) = &result.tags { if !tags.is_empty() { println!(); println!("{}", ui::gradient_text("─── Tags ───")); println!(" {}", tags.join(", ")); } } println!(); } else { output::print_output(&value, output_format)?; } Ok(()) } Err(e) => Err(e), } } // Helper function to check if a file has a text-based extension fn is_text_file(path: &std::path::Path) -> bool { const TEXT_EXTENSIONS: &[&str] = &[ "txt", "md", "json", "yaml", "yml", "toml", "xml", "csv", "log", "rst", "adoc", ]; path.extension() .and_then(|ext| ext.to_str()) .map(|ext| TEXT_EXTENSIONS.contains(&ext.to_lowercase().as_str())) .unwrap_or(false) } pub fn recall( client: &ApiClient, agent_id: &str, query: String, fact_type: Vec, budget: String, max_tokens: i64, trace: bool, include_chunks: bool, chunk_max_tokens: i64, verbose: bool, output_format: OutputFormat, ) -> Result<()> { let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Recalling memories...")) } else { None }; // Build include options if chunks are requested let include = if include_chunks { Some(IncludeOptions { chunks: Some(ChunkIncludeOptions { max_tokens: chunk_max_tokens, }), entities: None, }) } else { None }; let request = RecallRequest { query, types: if fact_type.is_empty() { None } else { Some(fact_type) }, budget: Some(parse_budget(&budget)), max_tokens, trace, query_timestamp: None, include, tags: None, tags_match: TagsMatch::Any, }; let response = client.recall(agent_id, &request, verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if output_format == OutputFormat::Pretty { ui::print_search_results(&result, trace, include_chunks); } else { output::print_output(&result, output_format)?; } Ok(()) } Err(e) => Err(e) } } pub fn reflect( client: &ApiClient, agent_id: &str, query: String, budget: String, context: Option, max_tokens: Option, schema_path: Option, verbose: bool, output_format: OutputFormat, ) -> Result<()> { let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Reflecting...")) } else { None }; // Load and parse schema if provided let response_schema = if let Some(path) = schema_path { let schema_content = fs::read_to_string(&path) .with_context(|| format!("Failed to read schema file: {}", path.display()))?; let schema: serde_json::Map = serde_json::from_str(&schema_content) .with_context(|| format!("Failed to parse JSON schema from: {}", path.display()))?; Some(schema) } else { None }; let request = ReflectRequest { query, budget: Some(parse_budget(&budget)), context, max_tokens: max_tokens.unwrap_or(4096), include: None, response_schema, tags: None, tags_match: TagsMatch::Any, }; let response = client.reflect(agent_id, &request, verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if output_format == OutputFormat::Pretty { ui::print_think_response(&result); } else { output::print_output(&result, output_format)?; } Ok(()) } Err(e) => Err(e) } } pub fn retain( client: &ApiClient, agent_id: &str, content: String, doc_id: Option, context: Option, r#async: bool, verbose: bool, output_format: OutputFormat, ) -> Result<()> { let doc_id = doc_id.unwrap_or_else(config::generate_doc_id); let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Retaining memory...")) } else { None }; let item = MemoryItem { content: content.clone(), context, metadata: None, timestamp: None, document_id: Some(doc_id.clone()), entities: None, tags: None, }; let request = RetainRequest { items: vec![item], async_: r#async, document_tags: None, }; let response = client.retain(agent_id, &request, r#async, verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if output_format == OutputFormat::Pretty { ui::print_success(&format!( "Memory retained successfully (document: {})", doc_id )); if result.is_async { println!(" Status: queued for background processing"); println!(" Items: {}", result.items_count); } else { println!(" Stored count: {}", result.items_count); } } else { output::print_output(&result, output_format)?; } Ok(()) } Err(e) => Err(e) } } pub fn retain_files( client: &ApiClient, agent_id: &str, path: PathBuf, recursive: bool, context: Option, r#async: bool, verbose: bool, output_format: OutputFormat, ) -> Result<()> { if !path.exists() { anyhow::bail!("Path does not exist: {}", path.display()); } let mut files = Vec::new(); if path.is_file() { files.push(path); } else if path.is_dir() { if recursive { for entry in WalkDir::new(&path) .into_iter() .filter_map(|e| e.ok()) .filter(|e| e.file_type().is_file()) { let path = entry.path(); if is_text_file(&path) { files.push(path.to_path_buf()); } } } else { for entry in fs::read_dir(&path)? { let entry = entry?; let path = entry.path(); if path.is_file() && is_text_file(&path) { files.push(path); } } } } if files.is_empty() { ui::print_warning("No text files found (supported: txt, md, json, yaml, yml, toml, xml, csv, log, rst, adoc)"); return Ok(()); } ui::print_info(&format!("Found {} files to import", files.len())); let pb = ui::create_progress_bar(files.len() as u64, "Processing files"); let mut items = Vec::new(); for file_path in &files { let content = fs::read_to_string(file_path) .with_context(|| format!("Failed to read file: {}", file_path.display()))?; let doc_id = file_path .file_stem() .and_then(|s| s.to_str()) .map(|s| s.to_string()) .unwrap_or_else(config::generate_doc_id); items.push(MemoryItem { content, context: context.clone(), metadata: None, timestamp: None, document_id: Some(doc_id), entities: None, tags: None, }); pb.inc(1); } pb.finish_with_message("Files processed"); // Always use async mode for the API call let request = RetainRequest { items, async_: true, document_tags: None, }; let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Submitting retain request...")) } else { None }; let response = client.retain(agent_id, &request, true, verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if r#async { // User requested async mode - return immediately if output_format == OutputFormat::Pretty { ui::print_success("Files queued for processing"); println!(" Items: {}", result.items_count); if let Some(op_id) = &result.operation_id { println!(" Operation ID: {}", op_id); } } else { output::print_output(&result, output_format)?; } } else { // Poll until completion if let Some(operation_id) = &result.operation_id { let poll_spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Processing memories...")) } else { None }; let (success, error_msg) = client.poll_operation(agent_id, operation_id, verbose)?; if let Some(mut sp) = poll_spinner { sp.finish(); } if success { if output_format == OutputFormat::Pretty { ui::print_success("Files retained successfully"); println!(" Items processed: {}", result.items_count); } else { output::print_output(&result, output_format)?; } } else { let msg = error_msg.unwrap_or_else(|| "Unknown error".to_string()); if output_format == OutputFormat::Pretty { ui::print_error(&format!("Retain operation failed: {}", msg)); } anyhow::bail!("Retain operation failed: {}", msg); } } else { // No operation ID returned, shouldn't happen with async=true if output_format == OutputFormat::Pretty { ui::print_success("Files retained successfully"); println!(" Items processed: {}", result.items_count); } else { output::print_output(&result, output_format)?; } } } Ok(()) } Err(e) => Err(e) } } pub fn delete( client: &ApiClient, agent_id: &str, unit_id: &str, verbose: bool, output_format: OutputFormat, ) -> Result<()> { let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner("Deleting memory unit...")) } else { None }; let response = client.delete_memory(agent_id, unit_id, verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if output_format == OutputFormat::Pretty { if result.success { ui::print_success("Memory unit deleted successfully"); } else { ui::print_error("Failed to delete memory unit"); } } else { output::print_output(&result, output_format)?; } Ok(()) } Err(e) => Err(e) } } pub fn clear( client: &ApiClient, agent_id: &str, fact_type: Option, yes: bool, verbose: bool, output_format: OutputFormat, ) -> Result<()> { // Confirmation prompt unless -y flag is used if !yes && output_format == OutputFormat::Pretty { let message = if let Some(ft) = &fact_type { format!( "Are you sure you want to clear all '{}' memories for bank '{}'? This cannot be undone.", ft, agent_id ) } else { format!( "Are you sure you want to clear ALL memories for bank '{}'? This cannot be undone.", agent_id ) }; let confirmed = ui::prompt_confirmation(&message)?; if !confirmed { ui::print_info("Operation cancelled"); return Ok(()); } } let spinner_msg = if let Some(ft) = &fact_type { format!("Clearing {} memories...", ft) } else { "Clearing all memories...".to_string() }; let spinner = if output_format == OutputFormat::Pretty { Some(ui::create_spinner(&spinner_msg)) } else { None }; let response = client.clear_memories(agent_id, fact_type.as_deref(), verbose); if let Some(mut sp) = spinner { sp.finish(); } match response { Ok(result) => { if output_format == OutputFormat::Pretty { if result.success { let msg = if fact_type.is_some() { "Memories cleared successfully" } else { "All memories cleared successfully" }; ui::print_success(msg); } else { ui::print_error("Failed to clear memories"); } } else { output::print_output(&result, output_format)?; } Ok(()) } Err(e) => Err(e) } } #[cfg(test)] mod tests { use super::*; use std::path::Path; #[test] fn test_is_text_file_supported_extensions() { let supported = [ "file.txt", "file.md", "file.json", "file.yaml", "file.yml", "file.toml", "file.xml", "file.csv", "file.log", "file.rst", "file.adoc", ]; for filename in supported { assert!( is_text_file(Path::new(filename)), "{} should be recognized as a text file", filename ); } } #[test] fn test_is_text_file_case_insensitive() { assert!(is_text_file(Path::new("file.JSON"))); assert!(is_text_file(Path::new("file.TXT"))); assert!(is_text_file(Path::new("file.Md"))); assert!(is_text_file(Path::new("file.YAML"))); } #[test] fn test_is_text_file_unsupported_extensions() { let unsupported = [ "file.pdf", "file.doc", "file.docx", "file.png", "file.jpg", "file.exe", "file.bin", "file.zip", "file.tar", "file.gz", ]; for filename in unsupported { assert!( !is_text_file(Path::new(filename)), "{} should NOT be recognized as a text file", filename ); } } #[test] fn test_is_text_file_no_extension() { assert!(!is_text_file(Path::new("README"))); assert!(!is_text_file(Path::new("Makefile"))); assert!(!is_text_file(Path::new(".gitignore"))); } #[test] fn test_is_text_file_with_path() { assert!(is_text_file(Path::new("/some/path/to/file.json"))); assert!(is_text_file(Path::new("../relative/path/file.md"))); assert!(!is_text_file(Path::new("/path/to/image.png"))); } #[test] fn test_parse_budget_valid_values() { assert!(matches!(parse_budget("low"), Budget::Low)); assert!(matches!(parse_budget("mid"), Budget::Mid)); assert!(matches!(parse_budget("high"), Budget::High)); } #[test] fn test_parse_budget_case_insensitive() { assert!(matches!(parse_budget("LOW"), Budget::Low)); assert!(matches!(parse_budget("MID"), Budget::Mid)); assert!(matches!(parse_budget("HIGH"), Budget::High)); assert!(matches!(parse_budget("Low"), Budget::Low)); assert!(matches!(parse_budget("High"), Budget::High)); } #[test] fn test_parse_budget_defaults_to_mid() { assert!(matches!(parse_budget("invalid"), Budget::Mid)); assert!(matches!(parse_budget(""), Budget::Mid)); assert!(matches!(parse_budget("unknown"), Budget::Mid)); } }