file_rag/

file_rag.rs

use std::fs::{self};
use std::path::{Path, PathBuf};
use std::time::Instant;

use anyhow::{Context, Result};
use llm_brain::LLMBrain;
use pdf_extract::extract_text;
use serde_json::json;

const DEFAULT_CHUNK_SIZE: usize = 500; // Approximate words per chunk

// --- Config Helper (Copied and adapted) ---
fn ensure_config_exists() -> Result<()> {
    let config_dir = PathBuf::from("./config");
    let default_config_path = config_dir.join("default.toml");
    if default_config_path.exists() {
        return Ok(());
    }
    if !config_dir.exists() {
        fs::create_dir_all(&config_dir)?;
    }
    let default_toml_content = r#"
[database]
path = "./llm_brain_file_rag_db" # DB path for this example
namespace = "file_rag_ns"
database = "file_rag_db"

[llm]
# Provide necessary LLM config
"#;
    fs::write(&default_config_path, default_toml_content)?;
    println!("Created default config: {}", default_config_path.display());
    Ok(())
}
// --- End Config Helper ---

// Simple text chunking by approximate word count
fn chunk_text_by_words(text: &str, chunk_size: usize) -> Vec<String> {
    let words: Vec<&str> = text.split_whitespace().collect();
    let mut chunks = Vec::new();
    for chunk_words in words.chunks(chunk_size) {
        chunks.push(chunk_words.join(" "));
    }
    chunks
}

async fn ingest_pdf(llm_brain: &LLMBrain, pdf_path: &Path) -> Result<Vec<String>> {
    let start_time = Instant::now();
    println!("Processing PDF: {}", pdf_path.display());

    // 1. Extract text
    let text = extract_text(pdf_path).context("Failed to extract text from PDF")?;
    println!("Extracted {} characters.", text.len());

    // 2. Chunk text
    let chunks = chunk_text_by_words(&text, DEFAULT_CHUNK_SIZE);
    println!("Split into {} chunks.", chunks.len());

    // 3. Add chunks as MemoryFragments
    let filename = pdf_path
        .file_name()
        .unwrap_or_default()
        .to_string_lossy()
        .to_string();
    let file_stem = pdf_path
        .file_stem()
        .unwrap_or_default()
        .to_string_lossy()
        .to_string();

    let mut chunk_ids = Vec::new();

    for (i, chunk_content) in chunks.into_iter().enumerate() {
        let chunk_num = i + 1;
        let metadata = json!({
            "memory_type": "DocumentChunk",
            "source_file": filename,
            "chunk_number": chunk_num,
            // Optionally add entity_name for direct reference if needed
            "entity_name": format!("{}_chunk_{}", file_stem, chunk_num)
        });

        // Add memory (handle potential errors)
        match llm_brain.add_memory(chunk_content, metadata).await {
            Ok(id) => {
                chunk_ids.push(id.to_string());
                if chunk_num % 10 == 0 {
                    println!("Added chunk {chunk_num}...");
                }
            }
            Err(e) => {
                eprintln!("Failed to add chunk {chunk_num}: {e}");
                // Decide whether to continue or stop on error
            }
        }
    }

    // 4. (Optional) Add a MemoryFragment for the document itself, linking to chunks
    let doc_content = format!(
        "Metadata for document: {}. Contains {} chunks.",
        filename,
        chunk_ids.len()
    );
    let doc_metadata = json!({
        "memory_type": "DocumentMeta",
        "entity_name": file_stem,
        "properties": {
            "file_path": pdf_path.to_string_lossy(),
            "file_name": filename,
            "chunk_count": chunk_ids.len()
        },
        "relationships": {
            "contains_chunks": chunk_ids // Store the actual DB IDs
        }
    });
    llm_brain.add_memory(doc_content, doc_metadata).await?;
    println!("Added document metadata entry.");

    println!("Ingestion completed in {:?}.", start_time.elapsed());
    Ok(chunk_ids) // Return chunk IDs if needed elsewhere
}

async fn run_queries(llm_brain: &LLMBrain, queries: &[&str]) -> Result<()> {
    println!("\n--- Running Queries ---");
    for query in queries {
        println!("\nQuery: {query}");
        let start_time = Instant::now();
        match llm_brain.recall(query, 3).await {
            // Retrieve top 3 relevant chunks
            Ok(results) => {
                println!(
                    "Found {} relevant chunks in {:?}:",
                    results.len(),
                    start_time.elapsed()
                );
                if results.is_empty() {
                    println!("  (No relevant chunks found)");
                } else {
                    for (fragment, score) in results {
                        println!(
                            "  - Score: {:.4}, Source: {:?}, Chunk: {:?}\n    Content: {:.150}...",
                            score,
                            fragment.metadata.get("source_file"),
                            fragment.metadata.get("chunk_number"),
                            fragment.content.replace('\n', " ") // Basic formatting
                        );
                    }
                }
            }
            Err(e) => {
                eprintln!("  Error during recall: {e}");
            }
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    println!("--- Starting File RAG Example ---");
    ensure_config_exists()?;
    let llm_brain = LLMBrain::launch().await?;

    // --- Ingestion Step ---
    // IMPORTANT: Adjust this path to your actual test PDF file
    // Using the path relative to workspace root as seen in python example context
    let pdf_path = PathBuf::from("tests/document.pdf");

    if !pdf_path.exists() {
        eprintln!("Error: PDF file not found at {}", pdf_path.display());
        eprintln!("Please place a PDF file at that location or update the path in the code.");
        return Ok(()); // Exit gracefully
    }

    // Ingest the PDF content
    if let Err(e) = ingest_pdf(&llm_brain, &pdf_path).await {
        eprintln!("Error during PDF ingestion: {e}");
        // Decide if program should terminate or continue
    }

    // --- Query Step ---
    let test_queries = [
        "What is the main topic of the document?",
        "Summarize the key points from the document.",
        "What are the main conclusions drawn in the document?",
        "what is silha center",
        "who is Charlotte Higgins",
        "Explain the lawsuits",
        "Explain OpenAI's Involvement",
        "who is Mike Masnick",
        "content moderation liability shield", // Added query
    ];

    if let Err(e) = run_queries(&llm_brain, &test_queries).await {
        eprintln!("Error during querying: {e}");
    }

    println!("\n--- File RAG Example Finished ---");
    println!("Note: Database is at ./llm_brain_file_rag_db");

    Ok(())
}