Added short-term memory

Cargo.toml

@@ -15,4 +15,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 
 # Парсинг аргументов CLI (флаги --session-id и т.д.)
-clap = { version = "4.5", features = ["derive"] }
+clap = { version = "4.5", features = ["derive"] }
+
+# Асинхронная работа с SQLite
+sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "sqlite"] }

readme.txt

@@ -0,0 +1,13 @@
+Как запускать и протестировать:
+Запуск с дефолтными параметрами:
+
+Bash
+cargo run
+Запуск с указанием конкретной сессии и своего URL (например, vLLM / llama.cpp / Ollama):
+
+Bash
+cargo run -- --session-id 12345 --api-url "http://localhost:8000/v1/chat/completions" --model "qwen2.5-coder"
+Справка по доступным флагам:
+
+Bash
+cargo run -- --help

src/main.rs

@@ -1,44 +1,129 @@
 use clap::Parser;
+use serde::{Deserialize, Serialize};
 use serde_json::json;
+use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
+use sqlx::{Pool, Sqlite};
 use std::io::{self, Write};
+use std::str::FromStr;
 use tokio::io::{AsyncBufReadExt, BufReader};
 
-/// CLI-утилита для взаимодействия с LLM в рамках указанной сессии
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
-    /// Идентификатор сессии/беседы (на будущее для VK peer_id)
     #[arg(short, long, default_value = "default_session")]
     session_id: String,
 
-    /// URL OpenAI-совместимого эндпоинта (vLLM / llama.cpp / Ollama)
     #[arg(
         short,
         long,
-        default_value = "http://localhost:11434/v1/chat/completions"
+        default_value = "http://192.168.0.50:6969/v1/chat/completions"
     )]
     api_url: String,
 
-    /// Название модели (для Ollama/vLLM)
+    #[arg(short, long, default_value = "qwen2.5-coder")]
-    #[arg(m_flag, short, long, default_value = "qwen2.5-coder")]
     model: String,
+
+    /// Путь к файлу базы данных SQLite
+    #[arg(long, default_value = "sqlite:chat_history.db?mode=rwc")]
+    db_url: String,
 }
 
-/// Функция отправки Stateless-запроса к LLM
+#[derive(Serialize, Deserialize, Debug, Clone)]
-async fn send_llm_request(
+struct ChatMessage {
+    role: String,
+    content: String,
+}
+
+/// Инициализация БД: создание таблицы, индексов и включение режима WAL
+async fn init_db(db_url: &str) -> Result<Pool<Sqlite>, Box<dyn std::error::Error>> {
+    let options = SqliteConnectOptions::from_str(db_url)?
+        .create_if_missing(true)
+        .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal); // WAL режим для высокой параллельности
+
+    let pool = SqlitePoolOptions::new()
+        .max_connections(10) // До 10 параллельных соединений к БД
+        .connect_with(options)
+        .await?;
+
+    // Создаем таблицу и индекс, если их нет
+    sqlx::query(
+        r#"
+        CREATE TABLE IF NOT EXISTS messages (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            session_id TEXT NOT NULL,
+            role TEXT NOT NULL,
+            content TEXT NOT NULL,
+            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
+        );
+        CREATE INDEX IF NOT EXISTS idx_messages_session_id_id 
+        ON messages (session_id, id DESC);
+        "#,
+    )
+    .execute(&pool)
+    .await?;
+
+    Ok(pool)
+}
+
+/// Сохранение сообщения в историю
+async fn save_message(
+    pool: &Pool<Sqlite>,
+    session_id: &str,
+    role: &str,
+    content: &str,
+) -> Result<(), sqlx::Error> {
+    sqlx::query("INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)")
+        .bind(session_id)
+        .bind(role)
+        .bind(content)
+        .execute(pool)
+        .await?;
+
+    Ok(())
+}
+
+/// Выборка последних N (5) сообщений контекста в хронологическом порядке
+async fn get_recent_history(
+    pool: &Pool<Sqlite>,
+    session_id: &str,
+    limit: i64,
+) -> Result<Vec<ChatMessage>, sqlx::Error> {
+    let rows = sqlx::query_as::<_, (String, String)>(
+        r#"
+        SELECT role, content 
+        FROM (
+            SELECT role, content, id
+            FROM messages
+            WHERE session_id = ?
+            ORDER BY id DESC
+            LIMIT ?
+        ) 
+        ORDER BY id ASC
+        "#,
+    )
+    .bind(session_id)
+    .bind(limit * 2)
+    .fetch_all(pool)
+    .await?;
+
+    let history = rows
+        .into_iter()
+        .map(|(role, content)| ChatMessage { role, content })
+        .collect();
+
+    Ok(history)
+}
+
+/// Отправка запроса с учетом накопленного контекста
+async fn send_llm_request_with_history(
     client: &reqwest::Client,
     api_url: &str,
     model: &str,
-    user_message: &str,
+    history: &[ChatMessage],
 ) -> Result<String, Box<dyn std::error::Error>> {
     let payload = json!({
         "model": model,
-        "messages": [
+        "messages": history, // Передаем весь сохраненный контекст
-            {
-                "role": "user",
-                "content": user_message
-            }
-        ],
         "temperature": 0.7
     });
 
@@ -64,13 +149,23 @@ async fn main() {
     let args = Args::parse();
     let client = reqwest::Client::new();
 
+    // 1. Инициализируем БД
+    let db_pool = match init_db(&args.db_url).await {
+        Ok(pool) => pool,
+        Err(e) => {
+            eprintln!("Ошибка инициализации БД: {}", e);
+            return;
+        }
+    };
+
     println!("==================================================");
-    println!("  LLM CLI Agent (Stage 1)");
+    println!("  LLM CLI Agent with Memory (Stage 2)");
     println!("  Session ID: {}", args.session_id);
     println!("  API URL:    {}", args.api_url);
     println!("  Model:      {}", args.model);
+    println!("  Database:   SQLite (WAL Mode)");
     println!("==================================================");
-    println!("Введите сообщение и нажмите Enter. Для выхода наберите 'exit' или 'quit'.\n");
+    println!("Введите сообщение. Для выхода наберите 'exit' или 'quit'.\n");
 
     let stdin = tokio::io::stdin();
     let mut reader = BufReader::new(stdin);
@@ -83,7 +178,7 @@ async fn main() {
         let bytes_read = reader.read_line(&mut input).await.unwrap_or(0);
 
         if bytes_read == 0 {
-            break; // EOF (Ctrl+D)
+            break;
         }
 
         let trimmed_input = input.trim();
@@ -92,23 +187,49 @@ async fn main() {
             continue;
         }
 
-        if trimmed_input.eq_ignore_ascii_case("exit")
+        if trimmed_input.eq_ignore_ascii_case("exit") || trimmed_input.eq_ignore_ascii_case("quit")
-            || trimmed_input.eq_ignore_ascii_case("quit")
         {
             println!("Завершение работы.");
             break;
         }
 
-        print!("[Ожидание ответа от LLM...]\r");
+        // 2. Сохраняем сообщение пользователя в БД
+        if let Err(e) = save_message(&db_pool, &args.session_id, "user", trimmed_input).await {
+            eprintln!("[Ошибка записи в БД]: {}", e);
+            continue;
+        }
+
+        // 3. Достаем последние 5 сообщений из БД (включая только что сохраненное)
+        let history = match get_recent_history(&db_pool, &args.session_id, 5).await {
+            Ok(h) => h,
+            Err(e) => {
+                eprintln!("[Ошибка чтения из БД]: {}", e);
+                continue;
+            }
+        };
+
+        print!(
+            "[Запрос к LLM с контекстом ({} сообщ.)...]\r",
+            history.len()
+        );
         io::stdout().flush().unwrap();
 
-        match send_llm_request(&client, &args.api_url, &args.model, trimmed_input).await {
+        // 4. Отправляем контекст в LLM
+        match send_llm_request_with_history(&client, &args.api_url, &args.model, &history).await {
             Ok(reply) => {
-                println!("\r[LLM]: {}\n", reply);
+                // \r - в начало, \x1B[2K - очистить всю текущую строку терминала
+                print!("\r\x1B[2K");
+                println!("[LLM]: {}\n", reply);
+
+                if let Err(e) = save_message(&db_pool, &args.session_id, "assistant", &reply).await
+                {
+                    eprintln!("[Ошибка записи ответа в БД]: {}", e);
+                }
             }
             Err(e) => {
-                eprintln!("\r[Ошибка]: {}\n", e);
+                print!("\r\x1B[2K");
+                eprintln!("[Ошибка LLM]: {}\n", e);
             }
         }
     }
-}
+}