feat: refactor database connection handling and add LLM component

- Replace direct database connection establishment with shared `establish_pg_connection` utility
- Add "llm" to required components list in bootstrap manager
- Lower default RUST_LOG level from debug to info in VSCode config
- Clean up imports and connection error messages
- Remove hardcoded database URL strings in favor of centralized connection handling

The changes improve code maintainability by centralizing database connection logic and adding support for the new LLM component in the bootstrap process.
This commit is contained in:
Rodrigo Rodriguez (Pragmatismo) 2025-11-01 07:20:04 -03:00
parent 6f59cdaab6
commit 8bf347a9a2
7 changed files with 139 additions and 38 deletions

2
.vscode/launch.json vendored
View file

@ -17,7 +17,7 @@
}, },
"args": [], "args": [],
"env": { "env": {
"RUST_LOG": "debug" "RUST_LOG": "info"
}, },
"cwd": "${workspaceFolder}" "cwd": "${workspaceFolder}"
}, },

View file

View file

@ -1,7 +1,8 @@
use crate::config::AppConfig; use crate::config::AppConfig;
use crate::package_manager::{InstallMode, PackageManager}; use crate::package_manager::{InstallMode, PackageManager};
use crate::shared::utils::establish_pg_connection;
use anyhow::Result; use anyhow::Result;
use diesel::{connection::SimpleConnection, RunQueryDsl, Connection, QueryableByName}; use diesel::{connection::SimpleConnection, RunQueryDsl, QueryableByName, Connection};
use dotenvy::dotenv; use dotenvy::dotenv;
use log::{debug, error, info, trace}; use log::{debug, error, info, trace};
use aws_sdk_s3::Client; use aws_sdk_s3::Client;
@ -140,10 +141,7 @@ impl BootstrapManager {
if pm.is_installed(component.name) { if pm.is_installed(component.name) {
pm.start(component.name)?; pm.start(component.name)?;
} else { } else {
let database_url = std::env::var("DATABASE_URL") let mut conn = establish_pg_connection()?;
.unwrap_or_else(|_| "postgres://gbuser:@localhost:5432/botserver".to_string());
let mut conn = diesel::pg::PgConnection::establish(&database_url)
.map_err(|e| anyhow::anyhow!("Failed to connect to database: {}", e))?;
let default_bot_id: uuid::Uuid = diesel::sql_query("SELECT id FROM bots LIMIT 1") let default_bot_id: uuid::Uuid = diesel::sql_query("SELECT id FROM bots LIMIT 1")
.load::<BotIdRow>(&mut conn) .load::<BotIdRow>(&mut conn)
.map(|rows| rows.first().map(|r| r.id).unwrap_or_else(|| uuid::Uuid::new_v4())) .map(|rows| rows.first().map(|r| r.id).unwrap_or_else(|| uuid::Uuid::new_v4()))
@ -189,7 +187,7 @@ impl BootstrapManager {
return Ok(config); return Ok(config);
} }
match diesel::PgConnection::establish(&database_url) { match establish_pg_connection() {
Ok(mut conn) => { Ok(mut conn) => {
if let Err(e) = self.apply_migrations(&mut conn) { if let Err(e) = self.apply_migrations(&mut conn) {
log::warn!("Failed to apply migrations: {}", e); log::warn!("Failed to apply migrations: {}", e);
@ -197,7 +195,7 @@ impl BootstrapManager {
return Ok(AppConfig::from_database(&mut conn)); return Ok(AppConfig::from_database(&mut conn));
} }
Err(e) => { Err(e) => {
log::warn!("Failed to connect to legacy database: {}", e); log::warn!("Failed to connect to database: {}", e);
return Ok(AppConfig::from_env()); return Ok(AppConfig::from_env());
} }
} }
@ -205,7 +203,7 @@ impl BootstrapManager {
} }
let pm = PackageManager::new(self.install_mode.clone(), self.tenant.clone())?; let pm = PackageManager::new(self.install_mode.clone(), self.tenant.clone())?;
let required_components = vec!["tables", "drive", "cache"]; let required_components = vec!["tables", "drive", "cache", "llm"];
let mut config = AppConfig::from_env(); let mut config = AppConfig::from_env();
for component in required_components { for component in required_components {
@ -260,8 +258,7 @@ impl BootstrapManager {
futures::executor::block_on(pm.install(component))?; futures::executor::block_on(pm.install(component))?;
if component == "tables" { if component == "tables" {
let database_url = std::env::var("DATABASE_URL").unwrap(); let mut conn = establish_pg_connection()
let mut conn = diesel::PgConnection::establish(&database_url)
.map_err(|e| anyhow::anyhow!("Failed to connect to database: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to connect to database: {}", e))?;
let migration_dir = include_dir::include_dir!("./migrations"); let migration_dir = include_dir::include_dir!("./migrations");
@ -363,9 +360,7 @@ impl BootstrapManager {
fn update_bot_config(&self, bot_id: &uuid::Uuid, component: &str) -> Result<()> { fn update_bot_config(&self, bot_id: &uuid::Uuid, component: &str) -> Result<()> {
use diesel::sql_types::{Text, Uuid as SqlUuid}; use diesel::sql_types::{Text, Uuid as SqlUuid};
let database_url = std::env::var("DATABASE_URL") let mut conn = establish_pg_connection()?;
.unwrap_or_else(|_| "postgres://gbuser:@localhost:5432/botserver".to_string());
let mut conn = diesel::pg::PgConnection::establish(&database_url)?;
// Ensure globally unique keys and update values atomically // Ensure globally unique keys and update values atomically
let config_key = format!("{}_{}", bot_id, component); let config_key = format!("{}_{}", bot_id, component);
@ -388,8 +383,7 @@ impl BootstrapManager {
} }
pub async fn upload_templates_to_drive(&self, config: &AppConfig) -> Result<()> { pub async fn upload_templates_to_drive(&self, config: &AppConfig) -> Result<()> {
let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| config.database_url()); let mut conn = establish_pg_connection()?;
let mut conn = diesel::PgConnection::establish(&database_url)?;
self.create_bots_from_templates(&mut conn)?; self.create_bots_from_templates(&mut conn)?;
let templates_dir = Path::new("templates"); let templates_dir = Path::new("templates");
if !templates_dir.exists() { if !templates_dir.exists() {
@ -539,10 +533,8 @@ impl BootstrapManager {
let bytes = response.body.collect().await?.into_bytes(); let bytes = response.body.collect().await?.into_bytes();
let csv_content = String::from_utf8(bytes.to_vec())?; let csv_content = String::from_utf8(bytes.to_vec())?;
let database_url = std::env::var("DATABASE_URL")
.unwrap_or_else(|_| "postgres://gbuser:@localhost:5432/botserver".to_string());
// Create new connection for config loading // Create new connection for config loading
let config_conn = diesel::PgConnection::establish(&database_url)?; let config_conn = establish_pg_connection()?;
let config_manager = ConfigManager::new(Arc::new(Mutex::new(config_conn))); let config_manager = ConfigManager::new(Arc::new(Mutex::new(config_conn)));
// Use default bot ID or create one if needed // Use default bot ID or create one if needed
@ -556,7 +548,7 @@ impl BootstrapManager {
.map_err(|e| anyhow::anyhow!("Failed to sync gbot config: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to sync gbot config: {}", e))?;
// Load config from database which now has the CSV values // Load config from database which now has the CSV values
let mut config_conn = diesel::PgConnection::establish(&database_url)?; let mut config_conn = establish_pg_connection()?;
let config = AppConfig::from_database(&mut config_conn); let config = AppConfig::from_database(&mut config_conn);
info!("Successfully loaded config from CSV"); info!("Successfully loaded config from CSV");
Ok(config) Ok(config)

View file

@ -439,6 +439,32 @@ impl ConfigManager {
Self { conn } Self { conn }
} }
pub fn get_config(
&self,
bot_id: &uuid::Uuid,
key: &str,
fallback: Option<&str>,
) -> Result<String, diesel::result::Error> {
let mut conn = self.conn.lock().unwrap();
let fallback_str = fallback.unwrap_or("");
#[derive(Debug, QueryableByName)]
struct ConfigValue {
#[diesel(sql_type = Text)]
value: String,
}
let result = diesel::sql_query(
"SELECT get_bot_config($1, $2, $3) as value"
)
.bind::<diesel::sql_types::Uuid, _>(bot_id)
.bind::<Text, _>(key)
.bind::<Text, _>(fallback_str)
.get_result::<ConfigValue>(&mut *conn)
.map(|row| row.value)?;
Ok(result)
}
pub fn sync_gbot_config( pub fn sync_gbot_config(
&self, &self,
bot_id: &uuid::Uuid, bot_id: &uuid::Uuid,

View file

@ -1,10 +1,12 @@
use actix_web::{post, web, HttpRequest, HttpResponse, Result}; use actix_web::{post, web, HttpRequest, HttpResponse, Result};
use crate::config::{AppConfig, ConfigManager};
use dotenvy::dotenv; use dotenvy::dotenv;
use log::{error, info}; use log::{error, info};
use reqwest::Client; use reqwest::Client;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::env; use std::env;
use tokio::time::{sleep, Duration}; use tokio::time::{sleep, Duration};
use uuid::Uuid;
// OpenAI-compatible request/response structures // OpenAI-compatible request/response structures
#[derive(Debug, Serialize, Deserialize)] #[derive(Debug, Serialize, Deserialize)]
@ -62,13 +64,45 @@ pub async fn ensure_llama_servers_running() -> Result<(), Box<dyn std::error::Er
return Ok(()); return Ok(());
} }
// Get configuration from environment variables // Get configuration with fallback to default bot config
let llm_url = env::var("LLM_URL").unwrap_or_else(|_| "http://localhost:8081".to_string()); let default_bot_id = Uuid::parse_str("00000000-0000-0000-0000-000000000000").unwrap();
let embedding_url = let config_manager = AppConfig::from_env().db_conn.map(ConfigManager::new);
env::var("EMBEDDING_URL").unwrap_or_else(|_| "http://localhost:8082".to_string());
let llama_cpp_path = env::var("LLM_CPP_PATH").unwrap_or_else(|_| "~/llama.cpp".to_string()); let llm_url = match &config_manager {
let llm_model_path = env::var("LLM_MODEL_PATH").unwrap_or_else(|_| "".to_string()); Some(cm) => env::var("LLM_URL").unwrap_or_else(|_|
let embedding_model_path = env::var("EMBEDDING_MODEL_PATH").unwrap_or_else(|_| "".to_string()); cm.get_config(&default_bot_id, "LLM_URL", None)
.unwrap_or("http://localhost:8081".to_string())
),
None => env::var("LLM_URL").unwrap_or("http://localhost:8081".to_string())
};
let embedding_url = match &config_manager {
Some(cm) => env::var("EMBEDDING_URL").unwrap_or_else(|_|
cm.get_config(&default_bot_id, "EMBEDDING_URL", None)
.unwrap_or("http://localhost:8082".to_string())
),
None => env::var("EMBEDDING_URL").unwrap_or("http://localhost:8082".to_string())
};
let llama_cpp_path = match &config_manager {
Some(cm) => env::var("LLM_CPP_PATH").unwrap_or_else(|_|
cm.get_config(&default_bot_id, "LLM_CPP_PATH", None)
.unwrap_or("~/llama.cpp".to_string())
),
None => env::var("LLM_CPP_PATH").unwrap_or("~/llama.cpp".to_string())
};
let llm_model_path = match &config_manager {
Some(cm) => env::var("LLM_MODEL_PATH").unwrap_or_else(|_|
cm.get_config(&default_bot_id, "LLM_MODEL_PATH", None)
.unwrap_or("".to_string())
),
None => env::var("LLM_MODEL_PATH").unwrap_or("".to_string())
};
let embedding_model_path = match &config_manager {
Some(cm) => env::var("EMBEDDING_MODEL_PATH").unwrap_or_else(|_|
cm.get_config(&default_bot_id, "EMBEDDING_MODEL_PATH", None)
.unwrap_or("".to_string())
),
None => env::var("EMBEDDING_MODEL_PATH").unwrap_or("".to_string())
};
info!("🚀 Starting local llama.cpp servers..."); info!("🚀 Starting local llama.cpp servers...");
info!("📋 Configuration:"); info!("📋 Configuration:");
@ -189,18 +223,49 @@ async fn start_llm_server(
std::env::set_var("OMP_PROC_BIND", "close"); std::env::set_var("OMP_PROC_BIND", "close");
// "cd {} && numactl --interleave=all ./llama-server -m {} --host 0.0.0.0 --port {} --threads 20 --threads-batch 40 --temp 0.7 --parallel 1 --repeat-penalty 1.1 --ctx-size 8192 --batch-size 8192 -n 4096 --mlock --no-mmap --flash-attn --no-kv-offload --no-mmap &", // "cd {} && numactl --interleave=all ./llama-server -m {} --host 0.0.0.0 --port {} --threads 20 --threads-batch 40 --temp 0.7 --parallel 1 --repeat-penalty 1.1 --ctx-size 8192 --batch-size 8192 -n 4096 --mlock --no-mmap --flash-attn --no-kv-offload --no-mmap &",
// Read config values with defaults
let n_moe = env::var("LLM_SERVER_N_MOE").unwrap_or("4".to_string());
let ctx_size = env::var("LLM_SERVER_CTX_SIZE").unwrap_or("4096".to_string());
let parallel = env::var("LLM_SERVER_PARALLEL").unwrap_or("1".to_string());
let cont_batching = env::var("LLM_SERVER_CONT_BATCHING").unwrap_or("true".to_string());
let mlock = env::var("LLM_SERVER_MLOCK").unwrap_or("true".to_string());
let no_mmap = env::var("LLM_SERVER_NO_MMAP").unwrap_or("true".to_string());
let gpu_layers = env::var("LLM_SERVER_GPU_LAYERS").unwrap_or("20".to_string());
// Build command arguments dynamically
let mut args = format!(
"-m {} --host 0.0.0.0 --port {} --top_p 0.95 --temp 0.6 --ctx-size {} --repeat-penalty 1.2 -ngl {}",
model_path, port, ctx_size, gpu_layers
);
if n_moe != "0" {
args.push_str(&format!(" --n-moe {}", n_moe));
}
if parallel != "1" {
args.push_str(&format!(" --parallel {}", parallel));
}
if cont_batching == "true" {
args.push_str(" --cont-batching");
}
if mlock == "true" {
args.push_str(" --mlock");
}
if no_mmap == "true" {
args.push_str(" --no-mmap");
}
if cfg!(windows) { if cfg!(windows) {
let mut cmd = tokio::process::Command::new("cmd"); let mut cmd = tokio::process::Command::new("cmd");
cmd.arg("/C").arg(format!( cmd.arg("/C").arg(format!(
"cd {} && .\\llama-server.exe -m {} --host 0.0.0.0 --port {} --top_p 0.95 --temp 0.6 --flash-attn on --ctx-size 4096 --repeat-penalty 1.2 -ngl 20 ", "cd {} && .\\llama-server.exe {}",
llama_cpp_path, model_path, port llama_cpp_path, args
)); ));
cmd.spawn()?; cmd.spawn()?;
} else { } else {
let mut cmd = tokio::process::Command::new("sh"); let mut cmd = tokio::process::Command::new("sh");
cmd.arg("-c").arg(format!( cmd.arg("-c").arg(format!(
"cd {} && ./llama-server -m {} --host 0.0.0.0 --port {} --top_p 0.95 --temp 0.6 --flash-attn on --ctx-size 4096 --repeat-penalty 1.2 -ngl 20 &", "cd {} && ./llama-server {} &",
llama_cpp_path, model_path, port llama_cpp_path, args
)); ));
cmd.spawn()?; cmd.spawn()?;
} }

View file

@ -1,4 +1,6 @@
use crate::config::AIConfig; use crate::config::AIConfig;
use anyhow::{Context, Result};
use diesel::{Connection, PgConnection};
use futures_util::StreamExt; use futures_util::StreamExt;
use indicatif::{ProgressBar, ProgressStyle}; use indicatif::{ProgressBar, ProgressStyle};
use log::trace; use log::trace;
@ -177,3 +179,12 @@ pub async fn call_llm(
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> { ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
Ok(format!("Generated response for: {}", prompt)) Ok(format!("Generated response for: {}", prompt))
} }
/// Establishes a PostgreSQL connection using DATABASE_URL environment variable
pub fn establish_pg_connection() -> Result<PgConnection> {
let database_url = std::env::var("DATABASE_URL")
.unwrap_or_else(|_| "postgres://gbuser:@localhost:5432/botserver".to_string());
PgConnection::establish(&database_url)
.with_context(|| format!("Failed to connect to database at {}", database_url))
}

View file

@ -4,18 +4,25 @@ server_host,0.0.0.0
server_port,8080 server_port,8080
sites_root,/tmp sites_root,/tmp
llm-key,gsk_ llm-key,none
llm-model,openai/gpt-oss-20b
llm-url,https://api.groq.com/openai/v1/chat/completions
llm-url,http://localhost:8080/v1 llm-url,http://localhost:8080/v1
llm-model,./botserver-stack/llm/data/DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf llm-model,botserver-stack/data/llm/DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf
embedding-url,http://localhost:8082 embedding-url,http://localhost:8082
embedding-model-path,./botserver-stack/llm/data/bge-small-en-v1.5-f32.gguf embedding-model-path,botserver-stack/data/llm/bge-small-en-v1.5-f32.gguf
llm-server,false llm-server,false
llm-server-path,~/llama.cpp llm-server-path,botserver-stack/bin/llm/
llm-server-model,botserver-stack/data/llm/DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf
llm-server-host,0.0.0.0
llm-server-port,8080
llm-server-gpu-layers,35
llm-server-n-moe,4
llm-server-ctx-size,2048
llm-server-parallel,4
llm-server-cont-batching,true
llm-server-mlock,true
llm-server-no-mmap,true
email-from,from@domain.com email-from,from@domain.com
email-server,mail.domain.com email-server,mail.domain.com

1 name value
4 sites_root /tmp
5 llm-key gsk_ none
6 llm-model llm-url openai/gpt-oss-20b http://localhost:8080/v1
7 llm-url llm-model https://api.groq.com/openai/v1/chat/completions botserver-stack/data/llm/DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf
llm-url http://localhost:8080/v1
llm-model ./botserver-stack/llm/data/DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf
embedding-url http://localhost:8082
8 embedding-model-path embedding-url ./botserver-stack/llm/data/bge-small-en-v1.5-f32.gguf http://localhost:8082
9 llm-server embedding-model-path false botserver-stack/data/llm/bge-small-en-v1.5-f32.gguf
10 llm-server-path llm-server ~/llama.cpp false
11 email-from llm-server-path from@domain.com botserver-stack/bin/llm/
12 email-server llm-server-model mail.domain.com botserver-stack/data/llm/DeepSeek-R1-Distill-Qwen-1.5B-Q3_K_M.gguf
13 email-port llm-server-host 587 0.0.0.0
14 email-user llm-server-port user@domain.com 8080
15 email-pass llm-server-gpu-layers 35
16 llm-server-n-moe 4
17 llm-server-ctx-size 2048
18 llm-server-parallel 4
19 llm-server-cont-batching true
20 llm-server-mlock true
21 llm-server-no-mmap true
22 email-from from@domain.com
23 email-server mail.domain.com
24 email-port 587
25 email-user user@domain.com
26 custom-server email-pass localhost
27 custom-port custom-server 5432 localhost
28 custom-database custom-port mycustomdb 5432