st/proxy/

memory.rs

//! 🧠 Memory Proxy - Scoped conversation history for LLMs
//!
//! This module adds memory capabilities to the LLM proxy, allowing for
//! persistent, scoped conversation history.
//!
//! "A proxy that remembers is a proxy that cares!" - The Cheet 😺

use crate::proxy::{LlmMessage, LlmProxy, LlmRequest, LlmResponse, LlmRole};
use anyhow::Result;
use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

/// 🧠 Scoped memory for a conversation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversationScope {
    pub id: String,
    pub messages: Vec<LlmMessage>,
    pub last_updated: chrono::DateTime<Utc>,
}

/// 🗄️ Persistent memory storage for the proxy
pub struct ProxyMemory {
    storage_path: PathBuf,
    scopes: HashMap<String, ConversationScope>,
    /// If true, skip all disk I/O operations
    in_memory_only: bool,
}

impl ProxyMemory {
    pub fn new() -> Result<Self> {
        let home = std::env::var("HOME").unwrap_or_else(|_| ".".to_string());
        let storage_path = Path::new(&home).join(".st").join("proxy_memory.json");

        if let Some(parent) = storage_path.parent() {
            fs::create_dir_all(parent)?;
        }

        let mut memory = Self {
            storage_path,
            scopes: HashMap::new(),
            in_memory_only: false,
        };

        memory.load()?;
        Ok(memory)
    }

    /// Create an in-memory only instance that doesn't persist to disk
    /// Used as a fallback when filesystem access fails
    pub fn in_memory_only() -> Self {
        Self {
            storage_path: PathBuf::new(), // Empty path, won't be used
            scopes: HashMap::new(),
            in_memory_only: true,
        }
    }

    pub fn get_scope(&self, scope_id: &str) -> Option<&ConversationScope> {
        self.scopes.get(scope_id)
    }

    pub fn update_scope(&mut self, scope_id: &str, messages: Vec<LlmMessage>) -> Result<()> {
        let scope = self
            .scopes
            .entry(scope_id.to_string())
            .or_insert_with(|| ConversationScope {
                id: scope_id.to_string(),
                messages: Vec::new(),
                last_updated: Utc::now(),
            });

        scope.messages.extend(messages);
        scope.last_updated = Utc::now();

        // Limit history to last 20 messages to keep it manageable
        if scope.messages.len() > 20 {
            scope.messages = scope.messages.split_off(scope.messages.len() - 20);
        }

        self.save()?;
        Ok(())
    }

    pub fn clear_scope(&mut self, scope_id: &str) -> Result<()> {
        self.scopes.remove(scope_id);
        self.save()?;
        Ok(())
    }

    fn load(&mut self) -> Result<()> {
        // Skip loading if in memory-only mode
        if self.in_memory_only {
            return Ok(());
        }
        if self.storage_path.exists() {
            let content = fs::read_to_string(&self.storage_path)?;
            self.scopes = serde_json::from_str(&content).unwrap_or_default();
        }
        Ok(())
    }

    fn save(&self) -> Result<()> {
        // Skip saving if in memory-only mode
        if self.in_memory_only {
            return Ok(());
        }
        let content = serde_json::to_string_pretty(&self.scopes)?;
        fs::write(&self.storage_path, content)?;
        Ok(())
    }
}

/// 🛠️ Enhanced proxy with memory support
pub struct MemoryProxy {
    pub inner: LlmProxy,
    pub memory: ProxyMemory,
}

impl MemoryProxy {
    pub fn new() -> Result<Self> {
        Ok(Self {
            inner: LlmProxy::default(),
            memory: ProxyMemory::new()?,
        })
    }

    /// Create a new MemoryProxy with auto-detection of local LLMs (Ollama, LM Studio)
    pub async fn with_local_detection() -> Result<Self> {
        Ok(Self {
            inner: LlmProxy::with_local_detection().await,
            memory: ProxyMemory::new()?,
        })
    }

    pub async fn complete_with_memory(
        &mut self,
        provider_name: &str,
        scope_id: &str,
        mut request: LlmRequest,
    ) -> Result<LlmResponse> {
        // 1. Retrieve history from scope
        if let Some(scope) = self.memory.get_scope(scope_id) {
            // Prepend history to current messages (after system message if present)
            let mut new_messages = Vec::new();

            // Keep existing system message at the top
            if let Some(system_msg) = request
                .messages
                .iter()
                .find(|m| m.role == LlmRole::System)
                .cloned()
            {
                new_messages.push(system_msg);
            }

            // Add history
            for msg in &scope.messages {
                if msg.role != LlmRole::System {
                    new_messages.push(msg.clone());
                }
            }

            // Add current user message(s)
            for msg in request.messages {
                if msg.role != LlmRole::System {
                    new_messages.push(msg);
                }
            }

            request.messages = new_messages;
        }

        // 2. Call the inner proxy
        let response = self.inner.complete(provider_name, request.clone()).await?;

        // 3. Update memory with the new exchange
        let mut new_history = Vec::new();
        // Add the last user message
        if let Some(last_user_msg) = request
            .messages
            .iter()
            .rev()
            .find(|m| m.role == LlmRole::User)
        {
            new_history.push(last_user_msg.clone());
        }
        // Add the assistant response
        new_history.push(LlmMessage {
            role: LlmRole::Assistant,
            content: response.content.clone(),
        });

        self.memory.update_scope(scope_id, new_history)?;

        Ok(response)
    }
}