polymathy 0.2.0

//! API endpoints and server implementation.
//!
//! This module handles:
//! - HTTP server setup and configuration
//! - API endpoint definitions
//! - Request/response handling
//! - OpenAPI documentation generation

use actix_web::{web, App, HttpServer, HttpResponse};
use serde_json::Value;
use futures::future::join_all;
use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use usearch::Index;
use log::debug;
use anyhow::{Result, Context};
use std::env;
use dotenv::dotenv;

use crate::search::{SearchQuery, ProcessedContent};
use crate::index::create_index;

use actix_web::middleware::Logger;

use apistos::api_operation;
use apistos::app::{BuildConfig, OpenApiWrapper};
use apistos::info::Info;
use apistos::server::Server;
use apistos::spec::Spec;
use apistos::web::{get, resource, scope};
use apistos::{RapidocConfig, RedocConfig, ScalarConfig, SwaggerUIConfig};

/// Processes search results by fetching content from URLs, chunking it, and generating embeddings.
async fn process_search_results(search_results: Value, _index: Arc<Mutex<Index>>) -> Result<HashMap<u64, (String, String)>> {
    let chunk_map: Arc<Mutex<HashMap<u64, (String, String)>>> = Arc::new(Mutex::new(HashMap::new()));
    let chunk_counter: Arc<Mutex<u64>> = Arc::new(Mutex::new(0));

    let futures: Vec<_> = search_results["results"]
        .as_array()
        .context("Results is not an array")?
        .iter()
        .filter_map(|result| result["url"].as_str().map(String::from))
        .take(10)
        .map(|url| {
            let client: reqwest::Client = reqwest::Client::new();

            let chunk_map = Arc::clone(&chunk_map);
            let chunk_counter = Arc::clone(&chunk_counter);
            let _index = Arc::clone(&_index);

            async move {
                let params = serde_json::json!({
                    "config": {
                        "chunking_size": 100,
                        "chunking_type": "words",
                        "embedding_model": "AllMiniLML6V2"
                    },
                    "url": url
                });
            
                let processor_url = env::var("PROCESSOR_URL").expect("PROCESSOR_URL must be set");
                let response = match client.post(&processor_url)
                    .json(&params)
                    .send()
                    .await
                {
                    Ok(response) => response,
                    Err(e) => {
                        log::debug!("Error processing URL {}: {}", url, e);
                        return Ok::<(), anyhow::Error>(());
                    }
                };
                
                log::debug!("Processed URL {}", url);
                
                // Check the status code of the response
                if response.status().is_server_error() {
                    log::debug!("Received 500 Internal Server Error for URL {}", url);
                    return Ok::<(), anyhow::Error>(());
                }

                log::debug!("Not a 500 {}", url);

                let processed_content: ProcessedContent = match response.json().await {
                    Ok(content) => content,
                    Err(e) => {
                        log::debug!("Error parsing processed content for URL {}: {}", url, e);
                        return Ok::<(), anyhow::Error>(());
                    }
                };
                log::debug!("Can parse {}", url);

                // Check if processed_content is empty
                if processed_content.chunks.is_empty() || processed_content.embeddings.is_empty() {
                    log::debug!("Skipping URL {} due to empty processed content", url);
                    return Ok::<(), anyhow::Error>(());
                }

                log::debug!("Parsing starts {}", url);

                // Add each chunk to the chunk map
                for (chunk_id, chunk_text) in &processed_content.chunks {
                    // Skip if no embedding exists for this chunk
                    if !processed_content.embeddings.contains_key(chunk_id) {
                        log::debug!("Embedding not found for chunk {} in URL {}", chunk_id, url);
                        continue;
                    }

                    let mut chunk_counter = chunk_counter.lock().unwrap();
                    let key: u64 = *chunk_counter;
                    chunk_map.lock().unwrap().insert(key, (processed_content.url.clone(), chunk_text.clone()));
                    *chunk_counter += 1;
                    log::debug!("Added chunk to map: {}", key);
                }
                Ok::<(), anyhow::Error>(())
            }
        })
        .collect();

    join_all(futures).await;
    let chunk_map_clone = chunk_map.lock().unwrap().clone();
    log::debug!("Chunk map: {:?}", chunk_map_clone);
    Ok(chunk_map_clone)
}

/// The main search endpoint.
/// It takes a search query, retrieves search results, processes them, and returns a map of chunks.
#[api_operation(summary = "Process a query and return processed content")]
async fn search_and_index(
    query: web::Query<SearchQuery>,
) -> actix_web::Result<HttpResponse> {
    let searxng_url = env::var("SEARXNG_URL").expect("SEARXNG_URL must be set");
    let client = reqwest::Client::new();

    let index: Arc<Mutex<Index>> = Arc::new(Mutex::new(create_index()));

    debug!("Search results is being called");

    let search_results: Value = client
        .get(searxng_url)
        .query(&[("q", &query.q), ("format", &"json".to_string())])
        .send()
        .await
        .map_err(actix_web::error::ErrorInternalServerError)?
        .json()
        .await
        .map_err(actix_web::error::ErrorInternalServerError)?;
    
    debug!("Search results are in");

    let chunk_map = process_search_results(search_results, Arc::clone(&index))
        .await
        .map_err(actix_web::error::ErrorInternalServerError)?;
    
    debug!("Processed search results");

    Ok(HttpResponse::Ok().json(chunk_map))
}

/// Runs the Polymath web service.
/// This function initializes the logger, loads environment variables, and starts the Actix web server.
///
/// # Returns
///
/// A `std::io::Result<()>` indicating success or failure of the server startup.
pub async fn run() -> std::io::Result<()> {
    dotenv().ok(); // Load .env file
    env_logger::init();

    let host = env::var("SERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
    let port = env::var("SERVER_PORT").unwrap_or_else(|_| "8080".to_string());
    let bind_address = format!("{}:{}", host, port);

    HttpServer::new(move || {
        let spec = Spec {
            info: Info {
              title: "Polymath API".to_string(),
              description: Some(
                "This is the polymath API".to_string(),
              ),
              ..Default::default()
            },
            servers: vec![Server {
              url: "/".to_string(),
              ..Default::default()
            }],
            ..Default::default()
          };
        
      App::new()
          .document(spec)
          .wrap(Logger::default())
          .service(scope("/v1")
              .service(resource("/search").route(get().to(search_and_index)))
      )
      .build_with(
          "/openapi.json",
          BuildConfig::default()
            .with(RapidocConfig::new(&"/rapidoc"))
            .with(RedocConfig::new(&"/redoc"))
            .with(ScalarConfig::new(&"/scalar"))
            .with(SwaggerUIConfig::new(&"/swagger"))
        )
    })
    .bind(bind_address)?
    .run()
    .await
}