refactor: env to args

Author: rrrodzilla

src/main.rs

@@ -9,32 +9,23 @@ mod transport; // Declare transport module
 
 // Use necessary items from modules and crates
 use crate::{
-    doc_loader::load_documents,
-    embeddings::{generate_embeddings, OPENAI_CLIENT},
-    embeddings::SerializableEmbedding,
+    doc_loader::{Document}, // Import Document struct and module
+    embeddings::{generate_embeddings, SerializableEmbedding, OPENAI_CLIENT}, // Group imports
     error::ServerError,
     server::RustDocsServer,
-    transport::StdioTransport, // Import StdioTransport
+    transport::StdioTransport,
 };
 use async_openai::Client as OpenAIClient;
-use rmcp::{
-    serve_server,
-    transport::io::stdio, // Keep stdio function
-    // Remove TransportAdapterAsyncCombinedRW import
-};
-use std::env;
+use bincode::config; // Keep config
 use ndarray::Array1;
-use std::fs::{self, File};
-use std::io::BufReader; // Removed unused BufWriter
-use std::path::PathBuf; // Removed unused Path
-use xdg::BaseDirectories;
-use bincode::{
-    config,
-    // serde::OwnedSerdeDecoder, // No longer needed
-    // decode_from_reader, // Removed unused import
-    // encode_to_vec, // Removed unused import
-    // Encode, Decode, // No longer needed directly
+use rmcp::{serve_server, transport::io::stdio};
+use std::{
+    env,
+    fs::{self, File},
+    io::BufReader,
+    path::PathBuf,
 };
+use xdg::BaseDirectories;
 
 #[tokio::main]
 async fn main() -> Result<(), ServerError> {
@@ -52,163 +43,179 @@ async fn main() -> Result<(), ServerError> {
         ServerError::MissingArgument("CRATE_VERSION".to_string())
     })?;
 
-    let _openai_api_key = env::var("OPENAI_API_KEY")
-        .map_err(|_| ServerError::MissingEnvVar("OPENAI_API_KEY".to_string()))?; // Needed later
-
-    // Load documents by generating them dynamically
-    println!("Loading documents for crate: {}", crate_name);
-    let documents = load_documents(&crate_name, &crate_version)?; // Pass crate_name and crate_version
-    println!("Loaded {} documents.", documents.len());
-
-    // Initialize OpenAI client and set it in the OnceLock
-    let openai_client = OpenAIClient::new();
-    OPENAI_CLIENT
-        .set(openai_client.clone()) // Clone for generate_embeddings
-        .expect("Failed to set OpenAI client");
-
-    // --- Persistence Logic ---
-    // Use XDG Base Directory specification for data storage
+    // --- Determine Paths ---
     let xdg_dirs = BaseDirectories::with_prefix("rustdocs-mcp-server")
-        .map_err(|e| ServerError::Xdg(format!("Failed to get XDG directories: {}", e)))?; // Use the new Xdg variant
+        .map_err(|e| ServerError::Xdg(format!("Failed to get XDG directories: {}", e)))?;
 
-    // Construct the path within the XDG data directory, including the crate name
-    let relative_path = PathBuf::from(&crate_name).join("embeddings.bin");
+    // Construct the path for embeddings file
+    let embeddings_relative_path = PathBuf::from(&crate_name).join("embeddings.bin");
+    let embeddings_file_path = xdg_dirs
+        .place_data_file(embeddings_relative_path)
+        .map_err(ServerError::Io)?;
 
-    // Use place_data_file to get the full path and ensure parent directories exist
-    let embeddings_file_path = xdg_dirs.place_data_file(relative_path)
-        .map_err(ServerError::Io)?; // Map IO error if directory creation fails
+    // --- Try Loading Embeddings ---
+    let mut loaded_from_cache = false;
+    let mut loaded_embeddings: Option<Vec<(String, Array1<f32>)>> = None;
 
-    let embeddings = if embeddings_file_path.exists() {
-        println!("Loading embeddings from: {:?}", embeddings_file_path);
+    if embeddings_file_path.exists() {
+        println!("Attempting to load embeddings from: {:?}", embeddings_file_path);
         match File::open(&embeddings_file_path) {
             Ok(file) => {
                 let reader = BufReader::new(file);
-                // Use top-level decode_from_reader now that bincode serde feature is enabled
-                match bincode::decode_from_reader::<Vec<SerializableEmbedding>, _, _>(reader, config::standard()) {
+                match bincode::decode_from_reader::<Vec<SerializableEmbedding>, _, _>(
+                    reader,
+                    config::standard(),
+                ) {
                     Ok(loaded_serializable) => {
-                        println!("Successfully loaded embeddings. Converting format...");
-                        // Convert back to Vec<(String, Array1<f32>)>
-                        let converted_embeddings = loaded_serializable
+                        println!("Successfully loaded embeddings from cache. Converting format...");
+                        let converted = loaded_serializable
                             .into_iter()
-                            .map(|se| (se.path, Array1::from(se.vector))) // Convert Vec to Array1
+                            .map(|se| (se.path, Array1::from(se.vector)))
                             .collect::<Vec<_>>();
-                        Some(converted_embeddings) // Wrap in Option for the outer match
+                        loaded_embeddings = Some(converted);
+                        loaded_from_cache = true; // Set flag
                     }
                     Err(e) => {
-                        println!("Failed to decode embeddings: {}. Regenerating...", e);
-                        // Fall through to regeneration
-                        None
+                        println!(
+                            "Failed to decode embeddings file: {}. Will regenerate.",
+                            e
+                        );
+                        // Proceed to generation
                     }
                 }
             }
             Err(e) => {
-                println!("Failed to open embeddings file: {}. Regenerating...", e);
-                // Fall through to regeneration
-                None
+                println!(
+                    "Failed to open embeddings file: {}. Will regenerate.",
+                    e
+                );
+                // Proceed to generation
             }
         }
     } else {
-        println!("Embeddings file not found. Generating...");
-        None
-    };
+        println!("Embeddings file not found. Will generate.");
+        // Proceed to generation
+    }
 
-    // Use loaded embeddings or generate new ones if loading failed or file didn't exist
-    // Variables to store generation stats if needed
+    // --- Generate or Use Loaded Embeddings ---
     let mut generated_tokens: Option<usize> = None;
     let mut generation_cost: Option<f64> = None;
+    let mut documents_for_server: Vec<Document> = Vec::new(); // Empty by default
 
-    let embeddings = match embeddings {
-        Some(e) => e,
+    let final_embeddings = match loaded_embeddings {
+        Some(embeddings) => {
+            println!("Using embeddings loaded from cache.");
+            embeddings // Use the ones loaded from the file
+        }
         None => {
-            // Directory creation is handled by xdg_dirs.place_data_file
-
-            // Generate embeddings
+            // --- Generation Path ---
+            println!("Proceeding with documentation loading and embedding generation.");
+
+            // Ensure OpenAI API key is available ONLY if generating
+            let _openai_api_key = env::var("OPENAI_API_KEY")
+                .map_err(|_| ServerError::MissingEnvVar("OPENAI_API_KEY".to_string()))?;
+
+            // Initialize OpenAI client ONLY if generating
+            let openai_client = OpenAIClient::new();
+            OPENAI_CLIENT
+                .set(openai_client.clone())
+                .expect("Failed to set OpenAI client");
+
+            // 1. Load documents
+            println!("Loading documents for crate: {}", crate_name);
+            // Use the imported module function directly
+            let loaded_documents = doc_loader::load_documents(&crate_name, &crate_version)?;
+            println!("Loaded {} documents.", loaded_documents.len());
+            documents_for_server = loaded_documents.clone(); // Clone for server if needed (though user said no)
+
+            // 2. Generate embeddings
             println!("Generating embeddings...");
-            // Capture the returned tuple (embeddings, total_tokens)
-            let (generated_embeddings, total_tokens) =
-                generate_embeddings(&openai_client, &documents, "text-embedding-3-small").await?;
-
-            // Calculate and print cost
-            // Price: $0.02 / 1M tokens for text-embedding-3-small
+            let (generated_embeddings, total_tokens) = generate_embeddings(
+                &openai_client,
+                &loaded_documents, // Use the just-loaded documents
+                "text-embedding-3-small",
+            )
+            .await?;
+
+            // Calculate and store cost
             let cost_per_million = 0.02;
             let estimated_cost = (total_tokens as f64 / 1_000_000.0) * cost_per_million;
             println!(
-                "Embedding generation cost for {} tokens: ${:.6}", // Format for cents/fractions
+                "Embedding generation cost for {} tokens: ${:.6}",
                 total_tokens, estimated_cost
             );
-            // Store generation stats
             generated_tokens = Some(total_tokens);
             generation_cost = Some(estimated_cost);
 
-
-            println!("Embeddings generated. Saving to: {:?}", embeddings_file_path);
-
-            // Convert to serializable format
-            let serializable_embeddings: Vec<SerializableEmbedding> = generated_embeddings // Use the embeddings from the tuple
+            // 3. Save embeddings
+            println!("Saving generated embeddings to: {:?}", embeddings_file_path);
+            let serializable_embeddings: Vec<SerializableEmbedding> = generated_embeddings
                 .iter()
                 .map(|(path, array)| SerializableEmbedding {
                     path: path.clone(),
-                    vector: array.to_vec(), // Convert Array1 to Vec
+                    vector: array.to_vec(),
                 })
                 .collect();
 
-            // Encode directly to Vec<u8>
             match bincode::encode_to_vec(&serializable_embeddings, config::standard()) {
                 Ok(encoded_bytes) => {
-                    // Write the bytes to the file
                     if let Err(e) = fs::write(&embeddings_file_path, encoded_bytes) {
                         println!("Warning: Failed to write embeddings file: {}", e);
                     } else {
                         println!("Embeddings saved successfully.");
                     }
                 }
                 Err(e) => {
-                    // Log error but continue
                     println!("Warning: Failed to encode embeddings to vec: {}", e);
                 }
             }
-            generated_embeddings
+            generated_embeddings // Return the generated embeddings
         }
     };
-    // --- End Persistence Logic ---
-
 
+    // --- Initialize and Start Server ---
     println!("Initializing server for crate: {}", crate_name);
 
-    // Create the service instance, passing embeddings
     // Prepare the startup summary message
-    let startup_message = {
-        let doc_count = documents.len();
-        match (generated_tokens, generation_cost) {
-            (Some(tokens), Some(cost)) => {
-                // Embeddings were generated
-                format!(
-                    "Server for crate '{}' initialized. Loaded {} documents. Generated embeddings for {} tokens (Est. Cost: ${:.6}).",
-                    crate_name, doc_count, tokens, cost
-                )
-            }
-            _ => {
-                // Embeddings were loaded from cache
-                format!(
-                    "Server for crate '{}' initialized. Loaded {} documents from cache.",
-                    crate_name, doc_count
-                )
-            }
-        }
+    let startup_message = if loaded_from_cache {
+        format!(
+            "Server for crate '{}' initialized. Loaded {} embeddings from cache.",
+            crate_name,
+            final_embeddings.len() // Use count from loaded/generated embeddings
+        )
+    } else {
+        // Embeddings were generated
+        let tokens = generated_tokens.unwrap_or(0);
+        let cost = generation_cost.unwrap_or(0.0);
+        format!(
+            "Server for crate '{}' initialized. Generated {} embeddings for {} tokens (Est. Cost: ${:.6}).",
+            crate_name,
+            final_embeddings.len(), // Use count from loaded/generated embeddings
+            tokens,
+            cost
+        )
     };
 
-    // Note: We still pass 'documents' which were loaded regardless of embedding source
-    let service = RustDocsServer::new(crate_name, documents, embeddings, startup_message)?;
+    // Create the service instance
+    // Pass the final embeddings and an empty Vec for documents as it's not needed by the service
+    let service = RustDocsServer::new(
+        crate_name,
+        documents_for_server, // Pass the (potentially empty) documents vec
+        final_embeddings,
+        startup_message,
+    )?;
 
     // Create the stdio transport
     let (stdin, stdout) = stdio();
-    // Use the custom StdioTransport wrapper
-    let transport = StdioTransport { reader: stdin, writer: stdout };
+    let transport = StdioTransport {
+        reader: stdin,
+        writer: stdout,
+    };
 
     println!("Rust Docs MCP server starting...");
 
     // Serve the server
-    serve_server(service, transport).await?; // Use imported serve_server
+    serve_server(service, transport).await?;
 
     println!("Rust Docs MCP server stopped.");
     Ok(())