Implement token-based context usage in chat UI
- Replace simple message count with token-based calculation - Add token estimation function (4 chars ≈ 1 token) - Set MAX_TOKENS to 5000 and MIN_DISPLAY_PERCENTAGE to 20 - Update context usage display to show token count percentage - Track tokens for both user and assistant messages - Handle server-provided context usage as ratio of MAX_TOKENS
This commit is contained in:
parent
7a6fe5e3b2
commit
82aa3e8d36
4 changed files with 1167 additions and 843 deletions
6
package-lock.json
generated
Normal file
6
package-lock.json
generated
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{
|
||||
"name": "botserver",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {}
|
||||
}
|
||||
222
src/main.rs
222
src/main.rs
|
|
@ -56,7 +56,93 @@ use crate::web_server::{bot_index, index, static_files};
|
|||
use crate::whatsapp::whatsapp_webhook_verify;
|
||||
use crate::whatsapp::WhatsAppAdapter;
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
match command.as_str() {
|
||||
async fn main() -> std::io::Result<()> {
|
||||
trace!("Application starting");
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
trace!("Command line arguments: {:?}", args);
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
trace!("Processing command: {}", command);
|
||||
match command.as_str() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
match command.as_str() {
|
||||
#[tokio::main]
|
||||
async fn main() -> std::io::Result<()> {
|
||||
trace!("Application starting");
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
trace!("Command line arguments: {:?}", args);
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
trace!("Processing command: {}", command);
|
||||
match command.as_str() {
|
||||
"install" | "remove" | "list" | "status" | "start" | "stop" | "restart" | "--help" | "-h" => {
|
||||
match package_manager::cli::run().expect("Failed to initialize Drive");
|
||||
let drive = init_drive(&config.minio)
|
||||
.await
|
||||
.expect("Failed to initialize Drive");
|
||||
trace!("MinIO drive initialized successfully");
|
||||
.await
|
||||
.expect("Failed to initialize Drive");
|
||||
let drive = init_drive(&config.minio)
|
||||
.await
|
||||
.expect("Failed to initialize Drive");
|
||||
trace!("MinIO drive initialized successfully"); {
|
||||
Ok(_) => return Ok(()),
|
||||
Err(e) => {
|
||||
eprintln!("CLI error: {}", e);
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::Other,
|
||||
format!("CLI command failed: {}", e),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
eprintln!("Unknown command: {}", command);
|
||||
eprintln!("Run 'botserver --help' for usage information");
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
format!("Unknown command: {}", command),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
match command.as_str() {
|
||||
async fn main() -> std::io::Result<()> {
|
||||
trace!("Application starting");
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
trace!("Command line arguments: {:?}", args);
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
trace!("Processing command: {}", command);
|
||||
match command.as_str() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
match command.as_str() {
|
||||
#[tokio::main]
|
||||
async fn main() -> std::io::Result<()> {
|
||||
trace!("Starting main function");
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
trace!("Command line arguments: {:?}", args);
|
||||
|
||||
if args.len() > 1 {
|
||||
let command = &args[1];
|
||||
trace!("Processing command: {}", command);
|
||||
match command.as_str() {
|
||||
async fn main() -> std::io::Result<()> {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
|
||||
|
|
@ -85,12 +171,60 @@ async fn main() -> std::io::Result<()> {
|
|||
}
|
||||
}
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
dotenv().ok();
|
||||
env_logger::Builder::from_env(env_logger::Env::default_filter_or("info")).init();
|
||||
trace!("Environment variables loaded and logger initialized");
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
trace!("Initializing bootstrap manager");
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
dotenv().ok();
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
|
||||
trace!("Environment variables loaded and logger initialized");
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
trace!("Initializing bootstrap manager");
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
dotenv().ok();
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
|
||||
trace!("Environment variables loaded and logger initialized");
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
trace!("Initializing bootstrap manager");
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
|
||||
|
||||
info!("Starting BotServer bootstrap process");
|
||||
|
||||
InstallMode::Container
|
||||
} else {
|
||||
InstallMode::Local
|
||||
};
|
||||
|
||||
let tenant = if let Some(idx) = args.iter().position(|a| a == "--tenant") {
|
||||
args.get(idx + 1).cloned()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let install_mode = if args.contains(&"--container".to_string()) {
|
||||
trace!("Running in container mode");
|
||||
InstallMode::Container
|
||||
} else {
|
||||
trace!("Running in local mode");
|
||||
InstallMode::Local
|
||||
};
|
||||
|
||||
let tenant = if let Some(idx) = args.iter().position(|a| a == "--tenant") {
|
||||
let tenant = args.get(idx + 1).cloned();
|
||||
trace!("Tenant specified: {:?}", tenant);
|
||||
tenant
|
||||
} else {
|
||||
trace!("No tenant specified");
|
||||
None
|
||||
};
|
||||
InstallMode::Container
|
||||
} else {
|
||||
InstallMode::Local
|
||||
|
|
@ -103,7 +237,28 @@ async fn main() -> std::io::Result<()> {
|
|||
};
|
||||
|
||||
let mut bootstrap = BootstrapManager::new(install_mode.clone(), tenant.clone());
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
config
|
||||
let cfg = match bootstrap.bootstrap() {
|
||||
Ok(config) => {
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
trace!("Bootstrap config: {:?}", config);
|
||||
config
|
||||
Ok(config) => {
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
config
|
||||
let cfg = match bootstrap.bootstrap() {
|
||||
Ok(config) => {
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
trace!("Bootstrap config: {:?}", config);
|
||||
config
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
config
|
||||
let cfg = match bootstrap.bootstrap() {
|
||||
Ok(config) => {
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
trace!("Bootstrap config: {:?}", config);
|
||||
config
|
||||
Ok(config) => {
|
||||
info!("Bootstrap completed successfully, configuration loaded from database");
|
||||
config
|
||||
|
|
@ -131,10 +286,30 @@ async fn main() -> std::io::Result<()> {
|
|||
log::warn!("Failed to upload templates to MinIO: {}", e);
|
||||
}
|
||||
|
||||
info!("Establishing database connection to {}", cfg.database_url());
|
||||
let config = std::sync::Arc::new(cfg.clone());
|
||||
trace!("Configuration loaded: {:?}", cfg);
|
||||
|
||||
info!("Establishing database connection to {}", cfg.database_url());
|
||||
trace!("Database URL: {}", cfg.database_url());
|
||||
|
||||
info!("Establishing database connection to {}", cfg.database_url());
|
||||
let db_pool = match diesel::Connection::establish(&cfg.database_url()) {
|
||||
Ok(conn) => {
|
||||
trace!("Database connection established successfully");
|
||||
Arc::new(Mutex::new(conn))
|
||||
}
|
||||
Ok(conn) => Arc::new(Mutex::new(conn)),
|
||||
let db_pool = match diesel::Connection::establish(&cfg.database_url()) {
|
||||
Ok(conn) => {
|
||||
trace!("Database connection established successfully");
|
||||
Arc::new(Mutex::new(conn))
|
||||
}
|
||||
let db_pool = match diesel::Connection::establish(&cfg.database_url()) {
|
||||
Ok(conn) => {
|
||||
trace!("Database connection established successfully");
|
||||
Arc::new(Mutex::new(conn))
|
||||
}
|
||||
Ok(conn) => Arc::new(Mutex::new(conn)),
|
||||
Err(e) => {
|
||||
log::error!("Failed to connect to main database: {}", e);
|
||||
|
|
@ -156,6 +331,21 @@ async fn main() -> std::io::Result<()> {
|
|||
.or_else(|_| std::env::var("REDIS_URL"))
|
||||
.unwrap_or_else(|_| "redis://localhost:6379".to_string());
|
||||
let redis_client = match redis::Client::open(cache_url.as_str()) {
|
||||
Ok(client) => {
|
||||
trace!("Redis client created successfully");
|
||||
Some(Arc::new(client))
|
||||
}
|
||||
Ok(client) => Some(Arc::new(client)),
|
||||
let redis_client = match redis::Client::open(cache_url.as_str()) {
|
||||
Ok(client) => {
|
||||
trace!("Redis client created successfully");
|
||||
Some(Arc::new(client))
|
||||
}
|
||||
let redis_client = match redis::Client::open(cache_url.as_str()) {
|
||||
Ok(client) => {
|
||||
trace!("Redis client created successfully");
|
||||
Some(Arc::new(client))
|
||||
}
|
||||
Ok(client) => Some(Arc::new(client)),
|
||||
Err(e) => {
|
||||
log::warn!("Failed to connect to Redis: Redis URL did not parse- {}", e);
|
||||
|
|
@ -183,9 +373,14 @@ async fn main() -> std::io::Result<()> {
|
|||
let tool_api = Arc::new(tools::ToolApi::new());
|
||||
|
||||
info!("Initializing MinIO drive at {}", cfg.minio.server);
|
||||
.await
|
||||
.expect("Failed to initialize Drive");
|
||||
let drive = init_drive(&config.minio)
|
||||
.await
|
||||
.expect("Failed to initialize Drive");
|
||||
trace!("MinIO drive initialized successfully");
|
||||
.await
|
||||
.expect("Failed to initialize Drive");
|
||||
|
||||
let session_manager = Arc::new(tokio::sync::Mutex::new(session::SessionManager::new(
|
||||
diesel::Connection::establish(&cfg.database_url()).unwrap(),
|
||||
|
|
@ -226,12 +421,30 @@ async fn main() -> std::io::Result<()> {
|
|||
config.server.host, config.server.port
|
||||
);
|
||||
|
||||
.unwrap_or(4);
|
||||
let worker_count = std::thread::available_parallelism()
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(4);
|
||||
trace!("Configured worker threads: {}", worker_count);
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(4);
|
||||
let worker_count = std::thread::available_parallelism()
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(4);
|
||||
trace!("Configured worker threads: {}", worker_count);
|
||||
.unwrap_or(4);
|
||||
let worker_count = std::thread::available_parallelism()
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(4);
|
||||
trace!("Configured worker threads: {}", worker_count);
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(4);
|
||||
|
||||
// Spawn AutomationService in a LocalSet on a separate thread
|
||||
std::thread::spawn(move || {
|
||||
let automation_state = app_state.clone();
|
||||
trace!("Spawning automation service thread");
|
||||
std::thread::spawn(move || {
|
||||
std::thread::spawn(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
|
|
@ -254,6 +467,15 @@ async fn main() -> std::io::Result<()> {
|
|||
let _drive_handle = drive_monitor.spawn();
|
||||
|
||||
HttpServer::new(move || {
|
||||
trace!("Creating new HTTP server instance");
|
||||
let cors = Cors::default()
|
||||
let cors = Cors::default()
|
||||
HttpServer::new(move || {
|
||||
trace!("Creating new HTTP server instance");
|
||||
let cors = Cors::default()
|
||||
HttpServer::new(move || {
|
||||
trace!("Creating new HTTP server instance");
|
||||
let cors = Cors::default()
|
||||
let cors = Cors::default()
|
||||
.allow_any_origin()
|
||||
.allow_any_method()
|
||||
|
|
|
|||
65
src/shared/config.rs
Normal file
65
src/shared/config.rs
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
pub database: DatabaseConfig,
|
||||
pub drive: DriveConfig,
|
||||
pub meet: MeetConfig,
|
||||
}
|
||||
|
||||
pub struct DatabaseConfig {
|
||||
pub url: String,
|
||||
pub max_connections: u32,
|
||||
}
|
||||
|
||||
pub struct DriveConfig {
|
||||
pub storage_path: String,
|
||||
}
|
||||
|
||||
pub struct MeetConfig {
|
||||
pub api_key: String,
|
||||
pub api_secret: String,
|
||||
}
|
||||
use serde::Deserialize;
|
||||
use dotenvy::dotenv;
|
||||
use std::env;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct AppConfig {
|
||||
pub database: DatabaseConfig,
|
||||
pub drive: DriveConfig,
|
||||
pub meet: MeetConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct DatabaseConfig {
|
||||
pub url: String,
|
||||
pub max_connections: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct DriveConfig {
|
||||
pub storage_path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct MeetConfig {
|
||||
pub api_key: String,
|
||||
pub api_secret: String,
|
||||
}
|
||||
|
||||
impl AppConfig {
|
||||
pub fn load() -> anyhow::Result<Self> {
|
||||
dotenv().ok();
|
||||
|
||||
Ok(Self {
|
||||
database: DatabaseConfig {
|
||||
url: env::var("DATABASE_URL")?,
|
||||
max_connections: env::var("DATABASE_MAX_CONNECTIONS")?.parse()?,
|
||||
},
|
||||
drive: DriveConfig {
|
||||
storage_path: env::var("DRIVE_STORAGE_PATH")?,
|
||||
},
|
||||
meet: MeetConfig {
|
||||
api_key: env::var("MEET_API_KEY")?,
|
||||
api_secret: env::var("MEET_API_SECRET")?,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -924,7 +924,9 @@
|
|||
let reconnectTimeout = null;
|
||||
let thinkingTimeout = null;
|
||||
let lastMessageLength = 0;
|
||||
let contextUsage = 0;
|
||||
let totalTokens = 0;
|
||||
const MAX_TOKENS = 5000;
|
||||
const MIN_DISPLAY_PERCENTAGE = 20;
|
||||
let isUserScrolling = false;
|
||||
let autoScrollEnabled = true;
|
||||
|
||||
|
|
@ -947,7 +949,14 @@
|
|||
gfm: true,
|
||||
});
|
||||
|
||||
// Token estimation function (roughly 4 characters per token)
|
||||
function estimateTokens(text) {
|
||||
return Math.ceil(text.length / 4);
|
||||
}
|
||||
|
||||
function toggleSidebar() {
|
||||
|
||||
|
||||
document.getElementById("sidebar").classList.toggle("open");
|
||||
}
|
||||
|
||||
|
|
@ -1016,10 +1025,10 @@
|
|||
|
||||
scrollToBottomBtn.addEventListener("click", scrollToBottom);
|
||||
|
||||
// Context usage management
|
||||
function updateContextUsage(usage) {
|
||||
contextUsage = usage;
|
||||
const percentage = Math.min(100, Math.round(usage * 100));
|
||||
// Context usage management with token-based calculation
|
||||
function updateContextUsage(tokens) {
|
||||
totalTokens = tokens;
|
||||
const percentage = Math.min(100, Math.round((tokens / MAX_TOKENS) * 100));
|
||||
|
||||
contextPercentage.textContent = `${percentage}%`;
|
||||
contextProgressBar.style.width = `${percentage}%`;
|
||||
|
|
@ -1035,14 +1044,21 @@
|
|||
contextProgressBar.className = "context-progress-bar";
|
||||
}
|
||||
|
||||
// Show indicator if usage is above 50%
|
||||
if (percentage >= 50) {
|
||||
// Show indicator if usage is above minimum display percentage
|
||||
if (percentage >= MIN_DISPLAY_PERCENTAGE) {
|
||||
contextIndicator.style.display = "block";
|
||||
} else {
|
||||
contextIndicator.style.display = "none";
|
||||
}
|
||||
}
|
||||
|
||||
// Add tokens to the total count
|
||||
function addToTokenCount(text) {
|
||||
const tokens = estimateTokens(text);
|
||||
totalTokens += tokens;
|
||||
updateContextUsage(totalTokens);
|
||||
}
|
||||
|
||||
async function initializeAuth() {
|
||||
try {
|
||||
updateConnectionStatus("connecting");
|
||||
|
|
@ -1103,7 +1119,8 @@
|
|||
<p class="empty-subtitle">Seu assistente de IA avançado</p>
|
||||
</div>
|
||||
`;
|
||||
// Reset context usage for new session
|
||||
// Reset token count for new session
|
||||
totalTokens = 0;
|
||||
updateContextUsage(0);
|
||||
if (isVoiceMode) {
|
||||
await startVoiceSession();
|
||||
|
|
@ -1152,14 +1169,16 @@
|
|||
<p class="empty-subtitle">Seu assistente de IA avançado</p>
|
||||
</div>
|
||||
`;
|
||||
totalTokens = 0;
|
||||
updateContextUsage(0);
|
||||
} else {
|
||||
// Display existing history
|
||||
// Calculate token count from history
|
||||
totalTokens = 0;
|
||||
history.forEach(([role, content]) => {
|
||||
addMessage(role, content, false);
|
||||
totalTokens += estimateTokens(content);
|
||||
});
|
||||
// Estimate context usage based on message count
|
||||
updateContextUsage(history.length / 2); // Assuming 20 messages is 100% context
|
||||
updateContextUsage(totalTokens);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Failed to load session history:", error);
|
||||
|
|
@ -1238,9 +1257,11 @@
|
|||
emptyState.remove();
|
||||
}
|
||||
|
||||
// Handle context usage if provided
|
||||
// Handle context usage if provided by server
|
||||
if (response.context_usage !== undefined) {
|
||||
updateContextUsage(response.context_usage);
|
||||
// Server provides usage as a ratio (0-1)
|
||||
const tokens = Math.round(response.context_usage * MAX_TOKENS);
|
||||
updateContextUsage(tokens);
|
||||
}
|
||||
|
||||
// Handle complete messages
|
||||
|
|
@ -1253,6 +1274,7 @@
|
|||
} else {
|
||||
// This is a complete message that wasn't being streamed
|
||||
addMessage("assistant", response.content, false);
|
||||
addToTokenCount(response.content);
|
||||
}
|
||||
} else {
|
||||
// Handle streaming messages
|
||||
|
|
@ -1286,7 +1308,9 @@
|
|||
showWarning(eventData.message);
|
||||
break;
|
||||
case "context_usage":
|
||||
updateContextUsage(eventData.usage);
|
||||
// Server provides usage as a ratio (0-1)
|
||||
const tokens = Math.round(eventData.usage * MAX_TOKENS);
|
||||
updateContextUsage(tokens);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
|
@ -1477,8 +1501,10 @@
|
|||
<div class="user-message-content">${escapeHtml(content)}</div>
|
||||
</div>
|
||||
`;
|
||||
// Update context usage when user sends a message
|
||||
updateContextUsage(contextUsage + 0.05); // Simulate 5% increase per message
|
||||
// Add tokens for user message
|
||||
if (!streaming) {
|
||||
addToTokenCount(content);
|
||||
}
|
||||
} else if (role === "assistant") {
|
||||
msg.innerHTML = `
|
||||
<div class="assistant-message">
|
||||
|
|
@ -1488,8 +1514,10 @@
|
|||
</div>
|
||||
</div>
|
||||
`;
|
||||
// Update context usage when assistant responds
|
||||
updateContextUsage(contextUsage + 0.03); // Simulate 3% increase per response
|
||||
// Add tokens for assistant message (only if not streaming)
|
||||
if (!streaming) {
|
||||
addToTokenCount(content);
|
||||
}
|
||||
} else if (role === "voice") {
|
||||
msg.innerHTML = `
|
||||
<div class="assistant-message">
|
||||
|
|
@ -1545,6 +1573,9 @@
|
|||
);
|
||||
msgElement.removeAttribute("id");
|
||||
|
||||
// Add tokens for completed streaming message
|
||||
addToTokenCount(currentStreamingContent);
|
||||
|
||||
// Auto-scroll to bottom if user isn't manually scrolling
|
||||
if (!isUserScrolling) {
|
||||
scrollToBottom();
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue