proxysql/lib/AI_Features_Manager.cpp

#include "AI_Features_Manager.h"
#include "GenAI_Thread.h"
#include "LLM_Bridge.h"
#include "Anomaly_Detector.h"
#include "sqlite3db.h"
#include "proxysql_utils.h"
#include <cstring>
#include <cstdlib>
#include <sys/stat.h>
#include <libgen.h>  // for dirname

// Global instance is defined in src/main.cpp
extern AI_Features_Manager *GloAI;

// GenAI module - configuration is now managed here
extern GenAI_Threads_Handler *GloGATH;

// Forward declaration to avoid header ordering issues
class ProxySQL_Admin;
extern ProxySQL_Admin *GloAdmin;

AI_Features_Manager::AI_Features_Manager()
	: shutdown_(0), llm_bridge(NULL), anomaly_detector(NULL), vector_db(NULL)
{
	pthread_rwlock_init(&rwlock, NULL);

	// Initialize status counters
	memset(&status_variables, 0, sizeof(status_variables));

	// Note: Configuration is now managed by GenAI module (GloGATH)
	// All genai-* variables are accessible via GloGATH->get_variable()
}

AI_Features_Manager::~AI_Features_Manager() {
	shutdown();

	// Note: Configuration strings are owned by GenAI module, not freed here
	pthread_rwlock_destroy(&rwlock);
}

int AI_Features_Manager::init_vector_db() {
	proxy_info("AI: Initializing vector storage at %s\n", GloGATH->variables.genai_vector_db_path);

	// Ensure directory exists
	char* path_copy = strdup(GloGATH->variables.genai_vector_db_path);
	if (!path_copy) {
		proxy_error("AI: Failed to allocate memory for path copy in init_vector_db\n");
		return -1;
	}
	char* dir = dirname(path_copy);
	struct stat st;
	if (stat(dir, &st) != 0) {
		// Create directory if it doesn't exist
		char cmd[512];
		snprintf(cmd, sizeof(cmd), "mkdir -p %s", dir);
		system(cmd);
	}
	free(path_copy);

	vector_db = new SQLite3DB();
	char path_buf[512];
	strncpy(path_buf, GloGATH->variables.genai_vector_db_path, sizeof(path_buf) - 1);
	path_buf[sizeof(path_buf) - 1] = '\0';
	int rc = vector_db->open(path_buf, SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE);
	if (rc != SQLITE_OK) {
		proxy_error("AI: Failed to open vector database: %s\n", GloGATH->variables.genai_vector_db_path);
		delete vector_db;
		vector_db = NULL;
		return -1;
	}

	// Create tables for LLM cache
	const char* create_llm_cache =
		"CREATE TABLE IF NOT EXISTS llm_cache ("
		"id INTEGER PRIMARY KEY AUTOINCREMENT , "
		"prompt TEXT NOT NULL , "
		"response TEXT NOT NULL , "
		"system_message TEXT , "
		"embedding BLOB , "
		"hit_count INTEGER DEFAULT 0 , "
		"last_hit INTEGER , "
		"created_at INTEGER DEFAULT (strftime('%s' ,  'now'))"
		");";

	if (vector_db->execute(create_llm_cache) != 0) {
		proxy_error("AI: Failed to create llm_cache table\n");
		return -1;
	}

	// Create table for anomaly patterns
	const char* create_anomaly_patterns =
		"CREATE TABLE IF NOT EXISTS anomaly_patterns ("
		"id INTEGER PRIMARY KEY AUTOINCREMENT , "
		"pattern_name TEXT , "
		"pattern_type TEXT , "  // 'sql_injection', 'dos', 'privilege_escalation'
		"query_example TEXT , "
		"embedding BLOB , "
		"severity INTEGER , "  // 1-10
		"created_at INTEGER DEFAULT (strftime('%s' ,  'now'))"
		");";

	if (vector_db->execute(create_anomaly_patterns) != 0) {
		proxy_error("AI: Failed to create anomaly_patterns table\n");
		return -1;
	}

	// Create table for query history
	const char* create_query_history =
		"CREATE TABLE IF NOT EXISTS query_history ("
		"id INTEGER PRIMARY KEY AUTOINCREMENT , "
		"prompt TEXT NOT NULL , "
		"response TEXT , "
		"embedding BLOB , "
		"execution_time_ms INTEGER , "
		"success BOOLEAN , "
		"timestamp INTEGER DEFAULT (strftime('%s' ,  'now'))"
		");";

	if (vector_db->execute(create_query_history) != 0) {
		proxy_error("AI: Failed to create query_history table\n");
		return -1;
	}

	// Create virtual vector tables for similarity search using sqlite-vec
	// Note: sqlite-vec extension is auto-loaded in Admin_Bootstrap.cpp:612

	// 1. LLM cache virtual table
	const char* create_llm_vec =
		"CREATE VIRTUAL TABLE IF NOT EXISTS llm_cache_vec USING vec0("
		"embedding float(1536)"
		");";

	if (vector_db->execute(create_llm_vec) != 0) {
		proxy_error("AI: Failed to create llm_cache_vec virtual table\n");
		// Virtual table creation failure is not critical - log and continue
		proxy_debug(PROXY_DEBUG_GENAI, 3, "Continuing without llm_cache_vec");
	}

	// 2. Anomaly patterns virtual table
	const char* create_anomaly_vec =
		"CREATE VIRTUAL TABLE IF NOT EXISTS anomaly_patterns_vec USING vec0("
		"embedding float(1536)"
		");";

	if (vector_db->execute(create_anomaly_vec) != 0) {
		proxy_error("AI: Failed to create anomaly_patterns_vec virtual table\n");
		proxy_debug(PROXY_DEBUG_GENAI, 3, "Continuing without anomaly_patterns_vec");
	}

	// 3. Query history virtual table
	const char* create_history_vec =
		"CREATE VIRTUAL TABLE IF NOT EXISTS query_history_vec USING vec0("
		"embedding float(1536)"
		");";

	if (vector_db->execute(create_history_vec) != 0) {
		proxy_error("AI: Failed to create query_history_vec virtual table\n");
		proxy_debug(PROXY_DEBUG_GENAI, 3, "Continuing without query_history_vec");
	}

	proxy_info("AI: Vector storage initialized successfully with virtual tables\n");
	return 0;
}

int AI_Features_Manager::init_llm_bridge() {
	if (!GloGATH->variables.genai_llm_enabled) {
		proxy_info("AI: LLM bridge disabled ,  skipping initialization\n");
		return 0;
	}

	proxy_info("AI: Initializing LLM Bridge\n");

	llm_bridge = new LLM_Bridge();

	// Set vector database
	llm_bridge->set_vector_db(vector_db);

	// Update config with current variables from GenAI module
	llm_bridge->update_config(
		GloGATH->variables.genai_llm_provider,
		GloGATH->variables.genai_llm_provider_url,
		GloGATH->variables.genai_llm_provider_model,
		GloGATH->variables.genai_llm_provider_key,
		GloGATH->variables.genai_llm_cache_similarity_threshold,
		GloGATH->variables.genai_llm_timeout_ms
	);

	if (llm_bridge->init() != 0) {
		proxy_error("AI: Failed to initialize LLM Bridge\n");
		delete llm_bridge;
		llm_bridge = NULL;
		return -1;
	}

	proxy_info("AI: LLM Bridge initialized\n");
	return 0;
}

int AI_Features_Manager::init_anomaly_detector() {
	if (!GloGATH->variables.genai_anomaly_enabled) {
		proxy_info("AI: Anomaly detection disabled ,  skipping initialization\n");
		return 0;
	}

	proxy_info("AI: Initializing Anomaly Detector\n");

	anomaly_detector = new Anomaly_Detector();
	if (anomaly_detector->init() != 0) {
		proxy_error("AI: Failed to initialize Anomaly Detector\n");
		delete anomaly_detector;
		anomaly_detector = NULL;
		return -1;
	}

	proxy_info("AI: Anomaly Detector initialized\n");
	return 0;
}

void AI_Features_Manager::close_vector_db() {
	if (vector_db) {
		delete vector_db;
		vector_db = NULL;
	}
}

void AI_Features_Manager::close_llm_bridge() {
	if (llm_bridge) {
		llm_bridge->close();
		delete llm_bridge;
		llm_bridge = NULL;
	}
}

void AI_Features_Manager::close_anomaly_detector() {
	if (anomaly_detector) {
		anomaly_detector->close();
		delete anomaly_detector;
		anomaly_detector = NULL;
	}
}

int AI_Features_Manager::init() {
	proxy_info("AI: Initializing AI Features Manager v%s\n", AI_FEATURES_MANAGER_VERSION);

	if (!GloGATH || !GloGATH->variables.genai_enabled) {
		proxy_info("AI: AI features disabled by configuration\n");
		return 0;
	}

	// Initialize vector storage first (needed by both LLM bridge and Anomaly Detector)
	if (init_vector_db() != 0) {
		proxy_error("AI: Failed to initialize vector storage\n");
		return -1;
	}

	// Initialize LLM bridge
	if (init_llm_bridge() != 0) {
		proxy_error("AI: Failed to initialize LLM bridge\n");
		return -1;
	}

	// Initialize Anomaly Detector
	if (init_anomaly_detector() != 0) {
		proxy_error("AI: Failed to initialize Anomaly Detector\n");
		return -1;
	}

	proxy_info("AI: AI Features Manager initialized successfully\n");
	return 0;
}

void AI_Features_Manager::shutdown() {
	if (shutdown_) return;
	shutdown_ = 1;

	proxy_info("AI: Shutting down AI Features Manager\n");

	close_llm_bridge();
	close_anomaly_detector();
	close_vector_db();

	proxy_info("AI: AI Features Manager shutdown complete\n");
}

void AI_Features_Manager::wrlock() {
	pthread_rwlock_wrlock(&rwlock);
}

void AI_Features_Manager::wrunlock() {
	pthread_rwlock_unlock(&rwlock);
}

// Note: Configuration get/set methods have been removed - they are now
// handled by the GenAI module (GloGATH). Use GloGATH->get_variable()
// and GloGATH->set_variable() for configuration access.

std::string AI_Features_Manager::get_status_json() {
	char buf[2048];
	snprintf(buf, sizeof(buf),
		"{"
		"\"version\": \"%s\" , "
		"\"llm\": {"
		"\"total_requests\": %llu , "
		"\"cache_hits\": %llu , "
		"\"local_calls\": %llu , "
		"\"cloud_calls\": %llu , "
		"\"total_response_time_ms\": %llu , "
		"\"cache_total_lookup_time_ms\": %llu , "
		"\"cache_total_store_time_ms\": %llu , "
		"\"cache_lookups\": %llu , "
		"\"cache_stores\": %llu , "
		"\"cache_misses\": %llu"
		"} , "
		"\"anomaly\": {"
		"\"total_checks\": %llu , "
		"\"blocked\": %llu , "
		"\"flagged\": %llu"
		"} , "
		"\"spend\": {"
		"\"daily_usd\": %.2f"
		"}"
		"}",
		AI_FEATURES_MANAGER_VERSION,
		status_variables.llm_total_requests,
		status_variables.llm_cache_hits,
		status_variables.llm_local_model_calls,
		status_variables.llm_cloud_model_calls,
		status_variables.llm_total_response_time_ms,
		status_variables.llm_cache_total_lookup_time_ms,
		status_variables.llm_cache_total_store_time_ms,
		status_variables.llm_cache_lookups,
		status_variables.llm_cache_stores,
		status_variables.llm_cache_misses,
		status_variables.anomaly_total_checks,
		status_variables.anomaly_blocked_queries,
		status_variables.anomaly_flagged_queries,
		status_variables.daily_cloud_spend_usd
	);

	return std::string(buf);
}