mubit-sdk 0.13.1

//! MuBit Learn — self-improving agentic memory.
//!
//! Unlike Python/JS which can monkey-patch LLM clients, the Rust module
//! provides explicit `enrich_messages()` and `record()` methods that the
//! user calls around their LLM invocation. `enrich_messages()` recalls and
//! injects relevant lessons; it does NOT reinforce on its own (a successful
//! call says nothing about whether the recalled memory helped). Call
//! [`LearnSession::feedback`] to credit the memories that actually helped.
//!
//! # Example
//!
//! ```rust,no_run
//! use mubit_sdk::learn::{LearnSession, LearnConfig};
//! use serde_json::json;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let session = LearnSession::new(
//!     LearnConfig::from_env()
//!         .agent_id("my-agent")
//!         .inject_lessons(true)
//!         .auto_reflect(true)
//! ).await;
//!
//! let messages = vec![
//!     json!({"role": "user", "content": "Fix the auth bug"}),
//! ];
//!
//! // Enrich messages with lessons before LLM call
//! let enriched = session.enrich_messages(&messages).await;
//!
//! // ... make your LLM call with `enriched` ...
//!
//! // Record the interaction after LLM call
//! session.record("The fix involves refreshing the token", "gpt-4o", 1500.0).await;
//!
//! // Credit the entries that were recalled for this call (closes the loop)
//! session.feedback(1.0).await;
//!
//! // End run (triggers reflection)
//! session.end().await;
//! # Ok(())
//! # }
//! ```

mod cache;
pub mod injection;
mod run_manager;

use std::sync::Arc;
use std::time::Duration;

use serde_json::Value;

use cache::LessonCache;
use injection::{extract_query, inject_context_openai};
use run_manager::RunManagerInner;

/// Configuration for a learn session.
#[derive(Clone, Debug)]
pub struct LearnConfig {
    pub api_key: String,
    pub endpoint: String,
    pub agent_id: String,
    pub user_id: String,
    pub session_id: Option<String>,
    pub inject_lessons: bool,
    pub injection_position: String,
    pub max_token_budget: u32,
    pub entry_types: Vec<String>,
    pub context_sections: Vec<String>,
    pub lane: String,
    pub auto_reflect: bool,
    pub reflect_after_n_calls: Option<u64>,
    pub cache_ttl_seconds: f64,
    pub cache_max_entries: usize,
    pub fail_open: bool,
    /// Pre-LLM lesson-injection timeout (seconds, hot path). Must exceed server
    /// context-assembly latency or injection silently no-ops under `fail_open`.
    /// Env override: `MUBIT_LEARN_CONTEXT_TIMEOUT`.
    pub context_fetch_timeout: f64,
    /// Include mental_model entries in context retrieval (highest priority).
    pub include_mental_models: bool,
}

impl Default for LearnConfig {
    fn default() -> Self {
        Self {
            api_key: String::new(),
            endpoint: "http://127.0.0.1:3000".to_string(),
            agent_id: "auto".to_string(),
            user_id: String::new(),
            session_id: None,
            inject_lessons: true,
            injection_position: "system".to_string(),
            max_token_budget: 2048,
            entry_types: Vec::new(),
            context_sections: Vec::new(),
            lane: String::new(),
            auto_reflect: true,
            reflect_after_n_calls: None,
            cache_ttl_seconds: 30.0,
            cache_max_entries: 100,
            fail_open: true,
            context_fetch_timeout: 5.0,
            include_mental_models: true,
        }
    }
}

impl LearnConfig {
    /// Create a config from environment variables.
    pub fn from_env() -> Self {
        let api_key = std::env::var("MUBIT_API_KEY").unwrap_or_default();
        let endpoint = std::env::var("MUBIT_ENDPOINT")
            .unwrap_or_else(|_| "http://127.0.0.1:3000".to_string());
        let mut cfg = Self {
            api_key,
            endpoint: endpoint.trim_end_matches('/').to_string(),
            ..Default::default()
        };
        if let Ok(t) = std::env::var("MUBIT_LEARN_CONTEXT_TIMEOUT") {
            if let Ok(secs) = t.parse::<f64>() {
                cfg.context_fetch_timeout = secs;
            }
        }
        cfg
    }

    pub fn api_key(mut self, key: impl Into<String>) -> Self {
        self.api_key = key.into();
        self
    }

    pub fn endpoint(mut self, ep: impl Into<String>) -> Self {
        self.endpoint = ep.into().trim_end_matches('/').to_string();
        self
    }

    pub fn agent_id(mut self, id: impl Into<String>) -> Self {
        self.agent_id = id.into();
        self
    }

    pub fn user_id(mut self, id: impl Into<String>) -> Self {
        self.user_id = id.into();
        self
    }

    pub fn session_id(mut self, id: impl Into<String>) -> Self {
        self.session_id = Some(id.into());
        self
    }

    pub fn inject_lessons(mut self, val: bool) -> Self {
        self.inject_lessons = val;
        self
    }

    pub fn injection_position(mut self, pos: impl Into<String>) -> Self {
        self.injection_position = pos.into();
        self
    }

    pub fn max_token_budget(mut self, budget: u32) -> Self {
        self.max_token_budget = budget;
        self
    }

    pub fn entry_types(mut self, types: Vec<String>) -> Self {
        self.entry_types = types;
        self
    }

    pub fn context_sections(mut self, sections: Vec<String>) -> Self {
        self.context_sections = sections;
        self
    }

    pub fn lane(mut self, lane: impl Into<String>) -> Self {
        self.lane = lane.into();
        self
    }

    pub fn auto_reflect(mut self, val: bool) -> Self {
        self.auto_reflect = val;
        self
    }

    pub fn reflect_after_n_calls(mut self, n: u64) -> Self {
        self.reflect_after_n_calls = Some(n);
        self
    }

    pub fn cache_ttl_seconds(mut self, ttl: f64) -> Self {
        self.cache_ttl_seconds = ttl;
        self
    }

    pub fn fail_open(mut self, val: bool) -> Self {
        self.fail_open = val;
        self
    }

    pub fn context_fetch_timeout(mut self, secs: f64) -> Self {
        self.context_fetch_timeout = secs;
        self
    }
}

/// A learn session that provides closed-loop memory for LLM interactions.
///
/// Call [`enrich_messages`] before your LLM call and [`record`] after.
/// Call [`end`] when your run is complete to trigger reflection.
pub struct LearnSession {
    inner: Arc<RunManagerInner>,
    cache: Arc<LessonCache>,
}

impl LearnSession {
    /// Create a new learn session.
    pub async fn new(config: LearnConfig) -> Self {
        let ttl = Duration::from_secs_f64(config.cache_ttl_seconds);
        let max_entries = config.cache_max_entries;

        Self {
            inner: Arc::new(RunManagerInner::new(config)),
            cache: Arc::new(LessonCache::new(ttl, max_entries)),
        }
    }

    /// The auto-generated or configured session ID for this run.
    pub fn session_id(&self) -> &str {
        &self.inner.session_id
    }

    /// Number of LLM calls recorded so far.
    pub fn call_count(&self) -> u64 {
        self.inner.call_count()
    }

    /// Enrich a message array with relevant lessons from MuBit.
    ///
    /// Fetches context via `get_context()` and injects it into the messages.
    /// Returns the enriched messages array (original is not modified).
    ///
    /// If injection is disabled or context retrieval fails (with `fail_open`),
    /// returns the original messages unchanged.
    pub async fn enrich_messages(&self, messages: &[Value]) -> Vec<Value> {
        if !self.inner.config.inject_lessons {
            return messages.to_vec();
        }

        let query = extract_query(messages, 200);
        if query.is_empty() {
            return messages.to_vec();
        }

        // One recall serves both injection and (later) feedback attribution:
        // fetch the block + the source entry IDs together, stash the IDs so
        // feedback() can credit exactly the entries that were recalled.
        let (context_block, ids) =
            if let Some((cached, cached_ids)) = self.cache.get(&self.inner.session_id, &query).await {
                (cached, cached_ids)
            } else {
                match self.inner.get_context_with_ids_http(&query).await {
                    Ok((block, ids)) => {
                        self.cache
                            .set(&self.inner.session_id, &query, block.clone(), ids.clone())
                            .await;
                        (block, ids)
                    }
                    Err(_) if self.inner.config.fail_open => (String::new(), Vec::new()),
                    Err(e) => {
                        eprintln!("mubit.learn: context retrieval failed: {e}");
                        (String::new(), Vec::new())
                    }
                }
            };

        self.inner.set_recalled_ids(ids);

        if context_block.is_empty() {
            messages.to_vec()
        } else {
            inject_context_openai(
                messages,
                &context_block,
                &self.inner.config.injection_position,
            )
        }
    }

    /// Record that an LLM call happened. Increments the call counter (which may
    /// trigger periodic reflection). It deliberately does NOT reinforce recalled
    /// entries — a successful call says nothing about whether the recalled
    /// memory helped, so auto-crediting would poison the store. Use
    /// [`feedback`](Self::feedback) to credit the memories that actually helped.
    pub async fn record(&self, _response_text: &str, _model: &str, _latency_ms: f64) {
        self.inner.increment();
    }

    /// Credit the entries recalled for the most recent call with a reward in
    /// `[-1, 1]` (clamped). This is the one-line way to close the loop: positive
    /// reinforces the recalled memories, negative penalizes them. No-op if no
    /// entries were recalled.
    pub async fn feedback(&self, signal: f32) {
        self.feedback_entries(signal, None, false).await;
    }

    /// Like [`feedback`](Self::feedback) but with explicit `entry_ids` and the
    /// option to mark the credited entries `verified_in_production` (a strong
    /// trust boost, only ever set explicitly). Passing `None` for `entry_ids`
    /// credits the entries recalled for the most recent call.
    pub async fn feedback_entries(
        &self,
        signal: f32,
        entry_ids: Option<Vec<String>>,
        verified_in_production: bool,
    ) {
        let ids: Vec<String> = entry_ids
            .unwrap_or_else(|| self.inner.last_recalled_ids())
            .into_iter()
            .filter(|s| !s.is_empty())
            .collect();
        if ids.is_empty() {
            return;
        }
        let (outcome, signal) = normalized_outcome(signal);
        let _ = self
            .inner
            .record_outcome_http(outcome, signal, &ids, verified_in_production)
            .await;
    }

    /// The entry IDs recalled for the most recent [`enrich_messages`] call.
    pub fn last_recalled_ids(&self) -> Vec<String> {
        self.inner.last_recalled_ids()
    }

    /// End the session. Triggers reflection if `auto_reflect` is enabled.
    pub async fn end(&self) {
        self.inner.end().await;
    }
}

/// Clamp a raw reward to `[-1, 1]` and derive the outcome label. Pure helper so
/// the feedback signal logic can be unit-tested without a server.
pub(crate) fn normalized_outcome(signal: f32) -> (&'static str, f32) {
    let s = signal.clamp(-1.0, 1.0);
    (if s >= 0.0 { "success" } else { "failure" }, s)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn config_default_context_timeout() {
        assert_eq!(LearnConfig::default().context_fetch_timeout, 5.0);
    }

    #[test]
    fn config_builder_sets_context_timeout() {
        assert_eq!(
            LearnConfig::default().context_fetch_timeout(3.0).context_fetch_timeout,
            3.0
        );
    }

    #[test]
    fn normalized_outcome_clamps_and_labels() {
        assert_eq!(normalized_outcome(2.5), ("success", 1.0));
        assert_eq!(normalized_outcome(0.0), ("success", 0.0));
        assert_eq!(normalized_outcome(-0.4), ("failure", -0.4));
        assert_eq!(normalized_outcome(-9.0), ("failure", -1.0));
    }

    #[test]
    fn injection_wraps_block_and_extracts_query() {
        let msgs = vec![
            serde_json::json!({"role": "system", "content": "You are helpful."}),
            serde_json::json!({"role": "user", "content": "what timeout?"}),
        ];
        let out = inject_context_openai(&msgs, "lesson A", "system");
        let joined = serde_json::to_string(&out).unwrap();
        assert!(joined.contains("<memory_context>"));
        assert!(joined.contains("You are helpful."));
        assert_eq!(extract_query(&msgs, 200), "what timeout?");
    }

    #[tokio::test]
    async fn feedback_is_noop_without_recalled_ids() {
        // Empty api_key + no recalled ids: must not panic and must not attempt I/O.
        let session =
            LearnSession::new(LearnConfig::default().session_id("t").auto_reflect(false)).await;
        session.feedback(1.0).await; // no recalled ids -> no-op
        assert!(session.last_recalled_ids().is_empty());
    }
}