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": [],
"env": {
"RUST_LOG": "debug"
"RUST_LOG": "info"
},
"cwd": "${workspaceFolder}"
},

View file

View file

@ -1,7 +1,8 @@
use crate::config::AppConfig;
use crate::package_manager::{InstallMode, PackageManager};
use crate::shared::utils::establish_pg_connection;
use anyhow::Result;
use diesel::{connection::SimpleConnection, RunQueryDsl, Connection, QueryableByName};
use diesel::{connection::SimpleConnection, RunQueryDsl, QueryableByName, Connection};
use dotenvy::dotenv;
use log::{debug, error, info, trace};
use aws_sdk_s3::Client;
@ -140,10 +141,7 @@ impl BootstrapManager {
if pm.is_installed(component.name) {
pm.start(component.name)?;
} else {
let database_url = std::env::var("DATABASE_URL")
.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 mut conn = establish_pg_connection()?;
let default_bot_id: uuid::Uuid = diesel::sql_query("SELECT id FROM bots LIMIT 1")
.load::<BotIdRow>(&mut conn)
.map(|rows| rows.first().map(|r| r.id).unwrap_or_else(|| uuid::Uuid::new_v4()))
@ -189,7 +187,7 @@ impl BootstrapManager {
return Ok(config);
}
match diesel::PgConnection::establish(&database_url) {
match establish_pg_connection() {
Ok(mut conn) => {
if let Err(e) = self.apply_migrations(&mut conn) {
log::warn!("Failed to apply migrations: {}", e);
@ -197,7 +195,7 @@ impl BootstrapManager {
return Ok(AppConfig::from_database(&mut conn));
}
Err(e) => {
log::warn!("Failed to connect to legacy database: {}", e);
log::warn!("Failed to connect to database: {}", e);
return Ok(AppConfig::from_env());
}
}
@ -205,7 +203,7 @@ impl BootstrapManager {
}
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();
for component in required_components {
@ -260,8 +258,7 @@ impl BootstrapManager {
futures::executor::block_on(pm.install(component))?;
if component == "tables" {
let database_url = std::env::var("DATABASE_URL").unwrap();
let mut conn = diesel::PgConnection::establish(&database_url)
let mut conn = establish_pg_connection()
.map_err(|e| anyhow::anyhow!("Failed to connect to database: {}", e))?;
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<()> {
use diesel::sql_types::{Text, Uuid as SqlUuid};
let database_url = std::env::var("DATABASE_URL")
.unwrap_or_else(|_| "postgres://gbuser:@localhost:5432/botserver".to_string());
let mut conn = diesel::pg::PgConnection::establish(&database_url)?;
let mut conn = establish_pg_connection()?;
// Ensure globally unique keys and update values atomically
let config_key = format!("{}_{}", bot_id, component);
@ -388,8 +383,7 @@ impl BootstrapManager {
}
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 = diesel::PgConnection::establish(&database_url)?;
let mut conn = establish_pg_connection()?;
self.create_bots_from_templates(&mut conn)?;
let templates_dir = Path::new("templates");
if !templates_dir.exists() {
@ -539,10 +533,8 @@ impl BootstrapManager {
let bytes = response.body.collect().await?.into_bytes();
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
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)));
// 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))?;
// 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);
info!("Successfully loaded config from CSV");
Ok(config)

View file

@ -439,6 +439,32 @@ impl ConfigManager {
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(
&self,
bot_id: &uuid::Uuid,

View file

@ -1,10 +1,12 @@
use actix_web::{post, web, HttpRequest, HttpResponse, Result};
use crate::config::{AppConfig, ConfigManager};
use dotenvy::dotenv;
use log::{error, info};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::env;
use tokio::time::{sleep, Duration};
use uuid::Uuid;
// OpenAI-compatible request/response structures
#[derive(Debug, Serialize, Deserialize)]
@ -62,13 +64,45 @@ pub async fn ensure_llama_servers_running() -> Result<(), Box<dyn std::error::Er
return Ok(());
}
// Get configuration from environment variables
let llm_url = env::var("LLM_URL").unwrap_or_else(|_| "http://localhost:8081".to_string());
let embedding_url =
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_model_path = env::var("LLM_MODEL_PATH").unwrap_or_else(|_| "".to_string());
let embedding_model_path = env::var("EMBEDDING_MODEL_PATH").unwrap_or_else(|_| "".to_string());
// Get configuration with fallback to default bot config
let default_bot_id = Uuid::parse_str("00000000-0000-0000-0000-000000000000").unwrap();
let config_manager = AppConfig::from_env().db_conn.map(ConfigManager::new);
let llm_url = match &config_manager {
Some(cm) => env::var("LLM_URL").unwrap_or_else(|_|
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!("📋 Configuration:");
@ -189,18 +223,49 @@ async fn start_llm_server(
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 &",
// 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) {
let mut cmd = tokio::process::Command::new("cmd");
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 ",
llama_cpp_path, model_path, port
"cd {} && .\\llama-server.exe {}",
llama_cpp_path, args
));
cmd.spawn()?;
} else {
let mut cmd = tokio::process::Command::new("sh");
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 &",
llama_cpp_path, model_path, port
"cd {} && ./llama-server {} &",
llama_cpp_path, args
));
cmd.spawn()?;
}

View file

@ -1,4 +1,6 @@
use crate::config::AIConfig;
use anyhow::{Context, Result};
use diesel::{Connection, PgConnection};
use futures_util::StreamExt;
use indicatif::{ProgressBar, ProgressStyle};
use log::trace;
@ -177,3 +179,12 @@ pub async fn call_llm(
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
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
sites_root,/tmp
llm-key,gsk_
llm-model,openai/gpt-oss-20b
llm-url,https://api.groq.com/openai/v1/chat/completions
llm-key,none
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-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-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-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