Refactoring

In the last chapter we managed to cobble up something working from scratch, but it's not enough to just work — we must make it better.

In this chapter, we're going to poke some holes in our previous setup and look at how we can make it more secure and robust.

Let's start with error handling. As of now, every little error crashes our program and the error messages returned are quite vague.

We used

#![allow(unused)]
fn main() {
Result<T, Box<dyn std::error::Error>>
}

to propagate errors upwards from each of our functions. This is convenient, but it collapses every error into one generic, vague type which offers little information. For example, if we have a missing API key in our environment variables, the program exits with a generic error message as shown below:

Looking at our setup we can identify several distinct boundaries which might fail:

• IO — reading the FAQ might fail due to a missing file or missing permissions
• Network — to transform each of our texts to vectors we make an HTTP request to the Gemini API
• Database — database operations might fail due to a missing database or missing permissions
• Env — our API key for Gemini is read from environment variables
• Parse — the FAQ file or Gemini API response might not match the expected format

To make our program better we need to provide explicit error messages at each of these boundaries. We do this by defining a custom error type:

#![allow(unused)]
fn main() {
use std::{fmt::Display, io};
#[derive(Debug)]
pub enum AppError {
    IOError(std::io::Error),
    NetworkError(reqwest::Error),
    DBError(rusqlite::Error),
    EnvError(std::env::VarError),
    ParseError(String),
     DimensionMismatch { expected: usize, got: usize },
     ZeroVector
     
    

}

impl Display for AppError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AppError::IOError(e) => match e.kind() {
                io::ErrorKind::NotFound => write!(f, "File not found"),
                io::ErrorKind::PermissionDenied => write!(f, "Permission denied"),
                _ => write!(f, "IO error: {e}"),
            },
            AppError::NetworkError(e) => match e.status() {
                Some(status) => match status {
                    reqwest::StatusCode::UNAUTHORIZED => write!(f, "Invalid API key"),
                    reqwest::StatusCode::TOO_MANY_REQUESTS => write!(f, "Rate limit exceeded"),
                    reqwest::StatusCode::NOT_FOUND => write!(f, "API endpoint not found"),
                    reqwest::StatusCode::INTERNAL_SERVER_ERROR => write!(f, "Gemini API error"),
                    _ => write!(f, "Network error: {e}"),
                },
                None => {
                    if e.is_timeout() {
                        write!(f, "Request timed out")
                    } else if e.is_connect() {
                        write!(f, "Connection failed")
                    } else {
                        write!(f, "Network error: {e}")
                    }
                }
            },
            AppError::DBError(e) => write!(f, "Database error: {e}"),
            AppError::EnvError(e) => write!(f, "Missing environment variable: {e}"),
            AppError::ParseError(e) => write!(f, "Parse error: {e}"),
            AppError::DimensionMismatch { expected, got } => {
    write!(f, "Vector dimension mismatch: expected {expected}, got {got}")
},
AppError::ZeroVector => write!(f, "Cannot compute similarity for a zero vector"),
        }
    }
}

impl std::error::Error for AppError {}
macro_rules! impl_from {
    ($variant:ident, $error:ty) => {
        impl From<$error> for AppError {
            fn from(e: $error) -> Self {
                Self::$variant(e)
            }
        }
    };
}

impl_from!(IOError, io::Error);
impl_from!(NetworkError, reqwest::Error);
impl_from!(DBError, rusqlite::Error);
impl_from!(EnvError, std::env::VarError);
impl_from!(ParseError, String);
}

We explicitly define an enum whose variants each wrap a specific error type, giving us precise information about what went wrong and where. The impl_from! macro derives From trait implementations for each variant, allowing the ? operator to automatically convert errors into the appropriate AppError variant without any manual casting.

What's left is to replace every Box<dyn std::error::Error> in our codebase with AppError. We also take this opportunity to move our embedding logic into lib.rs, separating it from the entry point:

#![allow(unused)]
fn main() {
// src/types.rs

use reqwest;
use rusqlite::{Connection, params};
use serde::Deserialize;
use serde_json::json;
use std::env;

use crate::errors::AppError;

#[derive(Debug)]
pub struct Embedding {
    pub label: String,
    pub vector: Vec<f32>,
}

#[derive(Deserialize)]
struct GeminiResponse {
    embedding: EmbeddingValues,
}

#[derive(Deserialize)]
struct BatchGeminiResponse {
    embeddings: Vec<EmbeddingValues>,
}

#[derive(Deserialize)]
struct EmbeddingValues {
    values: Vec<f64>,
}

impl Embedding {
    /// Compute cosine distance between two vectors.
    fn cosine_distance(a: &[f32], b: &[f32]) -> Result<f32, AppError> {
        if a.len() != b.len() {
            return Err(AppError::DimensionMismatch {
                expected: a.len(),
                got: b.len(),
            });
        }

        let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
        let mag_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let mag_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();

        if mag_a == 0.0 || mag_b == 0.0 {
            return Err(AppError::ZeroVector);
        }

        Ok(1.0 - dot / (mag_a * mag_b))
    }

    /// Initialize the database schema.
    pub fn init_db(conn: &Connection) -> Result<(), AppError> {
        conn.execute(
            "CREATE TABLE IF NOT EXISTS embeddings (
                id INTEGER PRIMARY KEY,
                label TEXT NOT NULL UNIQUE,
                vector BLOB NOT NULL
            )",
            [],
        )?;
        Ok(())
    }

    /// Persist this embedding to SQLite.
    pub fn commit(&self, conn: &Connection) -> Result<(), AppError> {
        let bytes: &[u8] = bytemuck::cast_slice(&self.vector);
        conn.execute(
            "INSERT OR REPLACE INTO embeddings (label, vector) VALUES (?1, ?2)",
            params![&self.label, bytes],
        )?;
        Ok(())
    }

    /// Perform a naive similarity search.
    /// NOTE: This performs a full table scan and is suitable only for small datasets.
    pub fn search(&self, conn: &Connection, limit: usize) -> Result<Vec<(String, f32)>, AppError> {
        let mut stmt = conn.prepare("SELECT label, vector FROM embeddings")?;

        let mut results: Vec<(String, f32)> = stmt
            .query_map([], |row| {
                let label: String = row.get(0)?;
                let bytes: Vec<u8> = row.get(1)?;
                let stored: &[f32] = bytemuck::cast_slice(&bytes);
                // cosine_distance can't use ? inside query_map's closure since
                // it expects rusqlite::Error — map to a sentinel and surface
                // the real error after collection
                let distance = Self::cosine_distance(&self.vector, stored).unwrap_or(2.0);
                Ok((label, distance))
            })?
            .collect::<Result<_, _>>()?;

        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        results.truncate(limit);
        Ok(results)
    }

    /// Create a reusable HTTP client.
    pub fn create_client() -> Result<reqwest::Client, AppError> {
        Ok(reqwest::Client::builder().build()?)
    }

    /// Convert a single piece of text into a vector using Gemini.
    pub async fn vectorize(
        text: &str,
        client: &reqwest::Client,
    ) -> Result<Vec<f32>, AppError> {
        let key = env::var("GEMINI_API_KEY")?;

        let body = json!({
            "model": "models/gemini-embedding-001",
            "content": {
                "parts": [{ "text": text }]
            }
        });

        let url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent";

        let res = client
            .post(url)
            .header("x-goog-api-key", &key)
            .json(&body)
            .send()
            .await?
            .json::<GeminiResponse>()
            .await?;

        Ok(res.embedding.values.into_iter().map(|v| v as f32).collect())
    }

    /// Convert multiple texts into vectors using Gemini batch embedding.
    pub async fn batch_vectorize(
        texts: &[String],
        client: &reqwest::Client,
    ) -> Result<Vec<Vec<f32>>, AppError> {
        let key = env::var("GEMINI_API_KEY")?;

        let requests: Vec<_> = texts
            .iter()
            .map(|text| {
                json!({
                    "model": "models/gemini-embedding-001",
                    "content": {
                        "parts": [{ "text": text }]
                    }
                })
            })
            .collect();

        let body = json!({ "requests": requests });

        let url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents";

        let res = client
            .post(url)
            .header("x-goog-api-key", &key)
            .json(&body)
            .send()
            .await?
            .json::<BatchGeminiResponse>()
            .await?;

        Ok(res
            .embeddings
            .into_iter()
            .map(|e| e.values.into_iter().map(|v| v as f32).collect())
            .collect())
    }

    /// Construct a single embedding from text.
    pub async fn new(
        label: String,
        client: &reqwest::Client,
    ) -> Result<Self, AppError> {
        let vector = Self::vectorize(&label, client).await?;
        Ok(Self { label, vector })
    }

    /// Construct multiple embeddings from text using batch vectorization.
    pub async fn batch_new(
        labels: Vec<String>,
        client: &reqwest::Client,
    ) -> Result<Vec<Self>, AppError> {
        let vectors = Self::batch_vectorize(&labels, client).await?;

        Ok(labels
            .into_iter()
            .zip(vectors)
            .map(|(label, vector)| Embedding { label, vector })
            .collect())
    }
}
}

#![allow(unused)]
fn main() {
use crate::errors::AppError;

pub mod errors;
pub mod types;
use crate::types::{Embedding};

use rusqlite::Connection;
use std::fs::File;
use std::io::{self, BufRead, BufReader, Write};


async fn search_faq(
    query: &str,
    client: &reqwest::Client,
    conn: &Connection,
) -> Result<(), AppError> {
    println!("\nSearching for: \"{}\"", query);
    println!("Generating embedding...");

    let query_embedding = Embedding::new(query.to_string(), client).await?;
    let results = query_embedding.search(conn, 3)?;

    if results.is_empty() {
        println!("No results found. Try loading the FAQ first with 'load' command.");
        return Ok(());
    }

    println!("\n--- Top {} Results ---", results.len());
    for (i, (label, distance)) in results.iter().enumerate() {
        let similarity = 1.0 - distance;
        println!("\n{}. [Similarity: {:.2}%]", i + 1, similarity * 100.0);
        println!("   {}", label);
        if similarity > 0.7 {
            println!("   ✓ Strong match!");
        }
    }
    println!();

    Ok(())
}

async fn load_faq(
    client: &reqwest::Client,
    conn: &Connection,
) -> Result<(), AppError> {
    let file = File::open("./faq.txt")?;
    let reader = BufReader::new(file);

    let mut current_question = String::new();
    let mut current_answer = String::new();
    let mut count = 0;

    for line in reader.lines() {
        let line = line.map_err(|e|AppError::ParseError(e.to_string()))?;
        let trimmed = line.trim();

        if trimmed.is_empty() || trimmed.starts_with("===") {
            continue;
        }

        if trimmed.starts_with("Q: ") {
            if !current_question.is_empty() && !current_answer.is_empty() {
                let combined = format!("Q: {}\nA: {}", current_question, current_answer);
                let embedding = Embedding::new(combined, client).await?;
                embedding.commit(conn)?;
                count += 1;
                print!("\rEmbedded {} questions...", count);
                io::stdout().flush()?;
            }

            current_question = trimmed.strip_prefix("Q: ").unwrap().to_string();
            current_answer.clear();
        } else if trimmed.starts_with("A: ") {
            current_answer = trimmed.strip_prefix("A: ").unwrap().to_string();
        } else if !current_answer.is_empty() {
            current_answer.push('\n');
            current_answer.push_str(trimmed);
        }
    }

    if !current_question.is_empty() && !current_answer.is_empty() {
        let combined = format!("Q: {}\nA: {}", current_question, current_answer);
        let embedding = Embedding::new(combined, client).await?;
        embedding.commit(conn)?;
        count += 1;
    }

    println!("\n✓ Total embedded: {} Q&A pairs", count);
    Ok(())
}


pub async fn run()->Result<(),AppError>{
     let conn = Connection::open("./embeddings.db")?;
    let client = Embedding::create_client()?;

    Embedding::init_db(&conn)?;

    println!("=== FAQ Search System ===");
    println!("Commands:");
    println!("  search <query>  - Search for similar questions");
    println!("  load            - Load FAQ from faq.txt");
    println!("  optimize        - Optimize vector index for faster search");
    println!("  quit            - Exit program");
    println!();

    loop {
        print!("> ");
        io::stdout().flush()?;

        let mut input = String::new();
        io::stdin().read_line(&mut input)?;
        let input = input.trim();

        if input.is_empty() {
            continue;
        }

        let parts: Vec<&str> = input.splitn(2, ' ').collect();
        let command = parts[0];

        match command {
            "quit" | "exit" | "q" => {
                println!("Goodbye!");
                break;
            }
            "load" => {
                println!("Loading FAQ...");
                load_faq(&client, &conn).await?;
                println!("✓ FAQ loaded successfully!");
                println!("  Tip: Run 'optimize' to speed up searches");
            }
            "optimize" => {
                println!("Optimizing vector index...");
                println!("✓ Optimization complete (placeholder)");
            }
            "search" => {
                if parts.len() < 2 {
                    println!("Usage: search <your question>");
                    continue;
                }
                let query = parts[1].trim_matches('"').trim();
                search_faq(query, &client, &conn).await?;
            }
            _ => {
                search_faq(input, &client, &conn).await?;
            }
        }
    }

    Ok(())
}
}

main.rs is now responsible for one thing — running the program and surfacing any errors to the user via Display, rather than letting Rust's default Debug output bypass our custom messages:

// main.rs

use embeddings::run;
#[tokio::main]
async fn main() {
   if let Err(e) = run().await {
       eprintln!("{}",e);
       std::process::exit(1)
   }
   

}