memista 0.1.1

//! Memista: High-performance vector search service
//!
//! Memista is a high-performance vector search library that combines SQLite for metadata storage
//! with USearch for efficient vector similarity search. It provides both a library interface for
//! embedding in Rust applications and a standalone HTTP server.
//!
//! ## Features
//!
//! - Fast Vector Similarity Search: Utilizes USearch for high-performance similarity search
//! - Persistent Storage: Stores text chunks and metadata in SQLite for durability
//! - Multi-Database Support: Supports multiple isolated databases through `database_id` partitioning
//! - Comprehensive API Documentation: Auto-generated OpenAPI documentation with Swagger, Redoc, and RapiDoc interfaces
//! - Environment-Based Configuration: Easily configurable through environment variables
//! - Asynchronous I/O: Built with async I/O for high performance and concurrency
//! - Memory Efficient: Uses optimized data structures for efficient memory usage
//!
//! ## Library Usage
//!
//! ```rust,no_run
//! use memista::{AppState, Config, create_app};
//! use async_sqlite::{PoolBuilder, JournalMode};
//! use actix_web::HttpServer;
//! use std::sync::Arc;
//!
//! #[actix_web::main]
//! async fn main() -> std::io::Result<()> {
//!     // Load configuration
//!     let config = Config::from_env().expect("Failed to load configuration");
//!     
//!     // Create a database pool
//!     let db_pool = PoolBuilder::new()
//!         .path(&config.database_path)
//!         .journal_mode(JournalMode::Wal)
//!         .open()
//!         .await
//!         .expect("Failed to create database pool");
//!
//!     // Create application state
//!     let app_state = Arc::new(AppState { db_pool });
//!
//!     // Start the HTTP server
//!     let bind_address = format!("{}:{}", config.server_host, config.server_port);
//!     HttpServer::new(move || {
//!         create_app(app_state.clone())
//!     })
//!     .bind(bind_address)?
//!     .run()
//!     .await
//! }
//! ```
//!
//! ## HTTP API
//!
//! Memista provides a RESTful HTTP API for vector search operations:
//!
//! ### POST /v1/insert
//!
//! Insert text chunks with their embeddings into a specified database.
//!
//! ```bash
//! curl -X POST http://localhost:8083/v1/insert 
//!   -H "Content-Type: application/json" 
//!   -d '{
//!     "database_id": "my_db",
//!     "chunks": [{
//!       "embedding": [0.1, 0.2],
//!       "text": "Sample text",
//!       "metadata": "{"source": "document1"}"
//!     }]
//!   }'
//! ```
//!
//! ### POST /v1/search
//!
//! Search for similar chunks using vector embeddings.
//!
//! ```bash
//! curl -X POST http://localhost:8083/v1/search 
//!   -H "Content-Type: application/json" 
//!   -d '{
//!     "database_id": "my_db",
//!     "embeddings": [[0.1, 0.2]],
//!     "num_results": 5
//!   }'
//! ```
//!
//! ### DELETE /v1/drop
//!
//! Drop a specific database and its associated vector index.
//!
//! ```bash
//! curl -X DELETE http://localhost:8083/v1/drop 
//!   -H "Content-Type: application/json" 
//!   -d '{
//!     "database_id": "my_db"
//!   }'
//! ```

use std::sync::Arc;
use actix_web::{web, HttpResponse};
use serde::{Deserialize, Serialize};
use serde_json::json;
use anyhow::Result;
use usearch::{Index, IndexOptions, MetricKind, ScalarKind, new_index};
use async_sqlite::Pool;
use apistos::{api_operation, ApiComponent};
use apistos::app::{BuildConfig, OpenApiWrapper};
use apistos::info::Info;
use apistos::server::Server;
use apistos::spec::Spec;
use apistos::web::{post, delete, resource, scope};
use apistos::{RapidocConfig, RedocConfig, ScalarConfig, SwaggerUIConfig};
use schemars::JsonSchema;
use std::env;

use log::debug;

/// Represents a chunk of text with its vector embedding and metadata
///
/// This struct is used to store and retrieve text chunks along with their
/// vector embeddings and associated metadata.
///
/// # Examples
///
/// ```
/// use memista::ChunkData;
///
/// let chunk = ChunkData {
///     embedding: vec![0.1, 0.2, 0.3],
///     text: "Hello, world!".to_string(),
///     metadata: "{\"source\": \"example\"}".to_string(),
/// };
/// ```
#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, ApiComponent)]
pub struct ChunkData {
    /// Vector embedding for the text chunk
    ///
    /// This is a dense vector representation of the text content,
    /// typically generated by an embedding model.
    pub embedding: Vec<f32>,
    /// The actual text content
    ///
    /// The original text that was embedded.
    pub text: String,
    /// Additional metadata in JSON string format
    ///
    /// Arbitrary metadata associated with this chunk, stored as a JSON string.
    /// This could include information like source document, timestamp, etc.
    pub metadata: String,
}

/// Request structure for inserting chunks into the database
///
/// This struct defines the format for requests to the insert endpoint.
/// It contains a database identifier and a list of chunks to insert.
#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, ApiComponent)]
pub struct InsertChunkRequest {
    /// Identifier for the database to insert into
    ///
    /// This identifier is used to partition data into separate databases.
    /// Each database has its own SQLite table and USearch index.
    pub database_id: String,
    /// List of chunks to insert
    ///
    /// The chunks to be inserted into the database, each with its
    /// embedding, text, and metadata.
    pub chunks: Vec<ChunkData>,
}

/// Request structure for searching similar chunks
///
/// This struct defines the format for requests to the search endpoint.
/// It contains a database identifier, query embeddings, and the number
/// of results to return.
#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, ApiComponent)]
pub struct SearchRequest {
    /// Identifier for the database to search in
    ///
    /// This identifier specifies which database to search in.
    /// Each database has its own SQLite table and USearch index.
    pub database_id: String,
    /// List of query embeddings to search for
    ///
    /// These are the vector embeddings to search for similar chunks.
    /// Multiple query embeddings can be provided in a single request.
    pub embeddings: Vec<Vec<f32>>,
    /// Maximum number of results to return per query
    ///
    /// This limits the number of results returned for each query embedding.
    pub num_results: usize,
}

/// Structure representing a search result
///
/// This struct represents a single result from a vector search,
/// including the text content, metadata, and similarity score.
#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, ApiComponent)]
pub struct SearchResult {
    /// The text content of the matching chunk
    ///
    /// The original text content of the chunk that matched the query.
    pub text: String,
    /// Optional metadata associated with the chunk
    ///
    /// Any metadata that was stored with the chunk, if available.
    pub metadata: Option<String>,
    /// Similarity score (higher means more similar)
    ///
    /// The similarity score between the query embedding and this chunk's
    /// embedding. Higher scores indicate more similar chunks.
    pub score: f32,
}

/// Request structure for dropping a database table and its index
///
/// This struct defines the format for requests to the drop endpoint.
/// It contains only the database identifier to drop.
#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, ApiComponent)]
pub struct DropTableRequest {
    /// Identifier for the database to drop
    ///
    /// This identifier specifies which database to drop.
    /// Both the SQLite table and USearch index will be removed.
    pub database_id: String,
}

/// Application state shared across HTTP handlers
///
/// This struct holds the shared state that is accessible to all
/// HTTP handlers, primarily the database connection pool.
pub struct AppState {
    /// Database connection pool for SQLite operations
    ///
    /// This pool provides connections to the SQLite database for
    /// storing and retrieving chunk metadata.
    pub db_pool: Pool,
}

/// Ensures that a database table exists for the given database ID
///
/// Creates the table if it doesn't exist with the proper schema.
/// Each database ID gets its own table with the name "chunks_{database_id}".
///
/// # Arguments
///
/// * `db_pool` - The database connection pool
/// * `database_id` - The identifier for the database
///
/// # Returns
///
/// A Result indicating success or an error if the table couldn't be created
pub async fn ensure_table_exists(db_pool: &Pool, database_id: &str) -> Result<(), actix_web::Error> {
    // Table name is prefixed with "chunks_" to avoid conflicts
    let table_name = format!("chunks_{}", database_id);
    db_pool.conn(move |conn| {
        conn.execute(
            &format!("CREATE TABLE IF NOT EXISTS {} (
                chunk_id INTEGER PRIMARY KEY AUTOINCREMENT,
                text TEXT,
                metadata TEXT
            )", table_name),
            [],
        )
    }).await.map_err(actix_web::error::ErrorInternalServerError)?;
    Ok(())
}

/// Loads an existing vector search index or creates a new one if it doesn't exist
///
/// The index is persisted to disk with a filename based on the database ID.
/// If an index file already exists for this database ID, it will be loaded.
/// Otherwise, a new empty index will be created.
///
/// # Arguments
///
/// * `database_id` - The identifier for the database
///
/// # Returns
///
/// A Result containing the loaded or created Index, or an error if the operation failed
pub fn load_or_create_index(database_id: &str) -> Result<Index, actix_web::Error> {
    // Index file name based on database ID
    let index_file = format!("{}.usearch", database_id);
    
    // Configure the index options for USearch
    let options = IndexOptions {
        // Currently hardcoded to 2 dimensions (should be configurable)
        dimensions: 2,
        // Using Inner Product metric for similarity
        metric: MetricKind::IP,
        // Using 32-bit floats for quantization
        quantization: ScalarKind::F32,
        // Default connectivity parameters
        connectivity: 0,
        expansion_add: 0,
        expansion_search: 0,
        // Enable multi-vector support
        multi: true,
    };
    
    // Create a new index with the specified options
    let index = new_index(&options).map_err(actix_web::error::ErrorInternalServerError)?;
    
    // Load existing index from disk if it exists
    if std::path::Path::new(&index_file).exists() {
        index.load(&index_file).map_err(actix_web::error::ErrorInternalServerError)?;
    }
    
    Ok(index)
}

/// API endpoint for inserting chunks into the database
/// 
/// This function handles the insertion of text chunks along with their vector embeddings
/// into both the SQLite database for metadata storage and the USearch index for fast
/// vector similarity search.
#[api_operation(summary = "Insert chunks into the database")]
#[allow(unused_mut)] // Mutable index is used for add operations later
pub async fn insert_chunk(
    app_state: web::Data<Arc<AppState>>,
    request: web::Json<InsertChunkRequest>,
) -> actix_web::Result<HttpResponse> {

    debug!("Loading index for database: {}", &request.database_id);

    // Load or create the vector search index
    let mut index = load_or_create_index(&request.database_id)?;

    // Get the current size of the index
    let index_size = index.size();
    
    // Reserve space in the index for better performance
    index.reserve(request.chunks.len() + index_size).map_err(actix_web::error::ErrorInternalServerError)?;

    debug!("Loaded index for database: {}", &request.database_id);

    // Ensure the database table exists
    ensure_table_exists(&app_state.db_pool, &request.database_id).await?;

    debug!("Ensured table exists for database: {}", &request.database_id);
    
    // Get the table name for this database
    let table_name = format!("chunks_{}", request.database_id);

    // Track the IDs of inserted chunks
    let mut inserted_ids = Vec::new();

    // Process each chunk in the request
    for chunk in &request.chunks {
        let chunk = chunk.clone();
        let table_name = table_name.clone();

        debug!("Inserting chunk into database");
        // Insert the chunk into SQLite database and get the assigned ID
        let chunk_id: i64 = app_state.db_pool.conn(move |conn| {
            conn.query_row(
                &format!("INSERT INTO {} (text, metadata) VALUES (?, ?) RETURNING chunk_id", table_name),
                [&chunk.text, &chunk.metadata],
                |row| row.get(0),
            )
        }).await.map_err(actix_web::error::ErrorInternalServerError)?;
        
        debug!("Inserting chunk into vector index");

        // Add the chunk's embedding to the vector search index
        index.add(chunk_id as u64, &chunk.embedding).map_err(actix_web::error::ErrorInternalServerError)?;

        // Track the inserted chunk ID
        inserted_ids.push(chunk_id);
    }

    // Save the updated index to disk
    let index_file = format!("{}.usearch", request.database_id);
    index.save(&index_file).map_err(actix_web::error::ErrorInternalServerError)?;

    // Return the IDs of inserted chunks
    Ok(HttpResponse::Ok().json(json!({ "inserted_ids": inserted_ids })))
}

/// API endpoint for searching similar chunks using vector embeddings
/// 
/// This function performs vector similarity search using the USearch index
/// and retrieves the corresponding text chunks and metadata from SQLite.
#[api_operation(summary = "Search for chunks")]
pub async fn search(
    app_state: web::Data<Arc<AppState>>,
    request: web::Json<SearchRequest>,
) -> actix_web::Result<HttpResponse> {
    // Load or create the vector search index
    let index = load_or_create_index(&request.database_id)?;

    // Ensure the database table exists
    ensure_table_exists(&app_state.db_pool, &request.database_id).await?;
    let table_name = format!("chunks_{}", request.database_id);

    // Store results for all query embeddings
    let mut all_results = Vec::new();

    // Process each query embedding
    for query_embedding in &request.embeddings {
        // Perform vector search using USearch
        let results = index.search(query_embedding, request.num_results).map_err(actix_web::error::ErrorInternalServerError)?;
        
        // Collect the search results with text and metadata
        let mut ranked_chunks = Vec::new();
        for (chunk_id, score) in results.keys.iter().zip(results.distances.iter()) {
            let chunk_id = *chunk_id;
            let score = *score;
            let table_name = table_name.clone();
            
            // Retrieve the text and metadata from SQLite
            let chunk = app_state.db_pool.conn(move |conn| {
                conn.query_row(
                    &format!("SELECT text, metadata FROM {} WHERE chunk_id = ?", table_name),
                    [chunk_id.to_string()],
                    |row| Ok((row.get::<_, String>(0)?, row.get::<_, Option<String>>(1)?)),
                )
            }).await.map_err(actix_web::error::ErrorInternalServerError)?;

            // Create a search result object
            ranked_chunks.push(SearchResult {
                text: chunk.0,
                metadata: chunk.1,
                score,
            });
        }

        // Add results for this query embedding
        all_results.push(ranked_chunks);
    }

    // Return all search results
    Ok(HttpResponse::Ok().json(all_results))
}

/// API endpoint for dropping a database table and its associated vector index
/// 
/// This function removes both the SQLite table and the USearch index file
/// for the specified database ID.
#[api_operation(summary = "Drop a table for a specific database")]
pub async fn drop_table(
    app_state: web::Data<Arc<AppState>>,
    request: web::Json<DropTableRequest>,
) -> actix_web::Result<HttpResponse> {
    // Get the table name for this database
    let table_name = format!("chunks_{}", request.database_id);
    
    // Drop the SQLite table if it exists
    app_state.db_pool.conn(move |conn| {
        conn.execute(
            &format!("DROP TABLE IF EXISTS {}", table_name),
            [],
        )
    }).await.map_err(actix_web::error::ErrorInternalServerError)?;
    
    // Remove the USearch index file if it exists
    let index_file = format!("{}.usearch", request.database_id);
    if std::path::Path::new(&index_file).exists() {
        std::fs::remove_file(index_file).map_err(actix_web::error::ErrorInternalServerError)?;
    }

    // Return success message
    Ok(HttpResponse::Ok().json(json!({"status": "success", "message": "Table and index dropped successfully"})))
}

/// Configuration structure for the application
///
/// This struct holds configuration values for the application,
/// typically loaded from environment variables.
#[derive(Debug, Clone)]
pub struct Config {
    /// Path to the SQLite database file
    pub database_path: String,
    /// Host address to bind the server to
    pub server_host: String,
    /// Port to listen on
    pub server_port: u16,
    /// Logging level (debug, info, warn, error)
    pub log_level: String,
}

impl Config {
    /// Loads configuration from environment variables
    ///
    /// Provides default values if environment variables are not set.
    /// Will panic if required environment variables are set but invalid.
    ///
    /// # Returns
    ///
    /// A Config struct with values loaded from environment variables
    pub fn from_env() -> Result<Self, env::VarError> {
        Ok(Config {
            database_path: env::var("DATABASE_PATH").unwrap_or_else(|_| "memista.db".to_string()),
            server_host: env::var("SERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()),
            server_port: env::var("SERVER_PORT")
                .unwrap_or_else(|_| "8083".to_string())
                .parse()
                .expect("SERVER_PORT must be a number"),
            log_level: env::var("LOG_LEVEL").unwrap_or_else(|_| "info".to_string()),
        })
    }
}

/// Creates and configures the HTTP server
///
/// This function creates an Actix-web application with all the
/// necessary routes and middleware for the Memista service.
///
/// # Arguments
///
/// * `app_state` - The shared application state
///
/// # Returns
///
/// A configured Actix-web application
pub fn create_app(app_state: Arc<AppState>) -> actix_web::App<impl actix_web::dev::ServiceFactory<actix_web::dev::ServiceRequest, Config = (), Response = actix_web::dev::ServiceResponse, Error = actix_web::Error, InitError = ()>> {
    // Configure OpenAPI specification
    let spec = Spec {
        info: Info {
            title: "Vector Search API".to_string(),
            description: Some("Vector Search API for chunk storage and retrieval".to_string()),
            ..Default::default()
        },
        servers: vec![Server {
            url: "/".to_string(),
            ..Default::default()
        }],
        ..Default::default()
    };

    // Build the Actix-web application with API routes
    actix_web::App::new()
        .app_data(web::Data::new(app_state.clone()))
        .document(spec)  // Enable OpenAPI documentation
        .service(scope("/v1")
            .service(resource("/insert").route(post().to(insert_chunk)))
            .service(resource("/search").route(post().to(search)))
            .service(resource("/drop").route(delete().to(drop_table)))
        )
        .build_with(
            "/openapi.json",
            BuildConfig::default()
                .with(RapidocConfig::new(&"/rapidoc"))
                .with(RedocConfig::new(&"/redoc"))
                .with(ScalarConfig::new(&"/scalar"))
                .with(SwaggerUIConfig::new(&"/swagger")),
        )
}