consciousness/src/thought/context.rs

// context.rs — Context window management
//
// Token counting, conversation trimming, and error classification.
// Journal entries are loaded from the memory graph store, not from
// a flat file — the parse functions are gone.

use crate::agent::types::*;
use chrono::{DateTime, Utc};
use tiktoken_rs::CoreBPE;

/// A single journal entry with its timestamp and content.
#[derive(Debug, Clone)]
pub struct JournalEntry {
    pub timestamp: DateTime<Utc>,
    pub content: String,
}

/// Context window size in tokens (from config).
pub fn context_window() -> usize {
    crate::config::get().api_context_window
}

/// Context budget in tokens: 60% of the model's context window.
fn context_budget_tokens() -> usize {
    context_window() * 60 / 100
}

/// Dedup and trim conversation entries to fit within the context budget.
///
/// 1. Dedup: if the same memory key appears multiple times, keep only
///    the latest render (drop the earlier Memory entry and its
///    corresponding assistant tool_call message).
/// 2. Trim: drop oldest entries until the conversation fits, snapping
///    to user message boundaries.
pub fn trim_entries(
    context: &ContextState,
    entries: &[ConversationEntry],
    tokenizer: &CoreBPE,
) -> Vec<ConversationEntry> {
    let count = |s: &str| tokenizer.encode_with_special_tokens(s).len();

    // --- Phase 1: dedup memory entries by key (keep last) ---
    let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
    let mut drop_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();

    for (i, entry) in entries.iter().enumerate() {
        if let ConversationEntry::Memory { key, .. } = entry {
            if let Some(prev) = seen_keys.insert(key.as_str(), i) {
                drop_indices.insert(prev);
            }
        }
    }

    let deduped: Vec<ConversationEntry> = entries.iter().enumerate()
        .filter(|(i, _)| !drop_indices.contains(i))
        .map(|(_, e)| e.clone())
        .collect();

    // --- Phase 2: trim to fit context budget ---
    let max_tokens = context_budget_tokens();
    let identity_cost = count(&context.system_prompt)
        + context.personality.iter().map(|(_, c)| count(c)).sum::<usize>();
    let journal_cost: usize = context.journal.iter().map(|e| count(&e.content)).sum();
    let reserve = max_tokens / 4;
    let available = max_tokens
        .saturating_sub(identity_cost)
        .saturating_sub(journal_cost)
        .saturating_sub(reserve);

    let msg_costs: Vec<usize> = deduped.iter()
        .map(|e| msg_token_count(tokenizer, e.message())).collect();
    let total: usize = msg_costs.iter().sum();

    let mut skip = 0;
    let mut trimmed = total;
    while trimmed > available && skip < deduped.len() {
        trimmed -= msg_costs[skip];
        skip += 1;
    }

    // Walk forward to user message boundary
    while skip < deduped.len() && deduped[skip].message().role != Role::User {
        skip += 1;
    }

    deduped[skip..].to_vec()
}

/// Count the token footprint of a message using BPE tokenization.
pub fn msg_token_count(tokenizer: &CoreBPE, msg: &Message) -> usize {
    let count = |s: &str| tokenizer.encode_with_special_tokens(s).len();
    let content = msg.content.as_ref().map_or(0, |c| match c {
        MessageContent::Text(s) => count(s),
        MessageContent::Parts(parts) => parts.iter()
            .map(|p| match p {
                ContentPart::Text { text } => count(text),
                ContentPart::ImageUrl { .. } => 85,
            })
            .sum(),
    });
    let tools = msg.tool_calls.as_ref().map_or(0, |calls| {
        calls.iter()
            .map(|c| count(&c.function.arguments) + count(&c.function.name))
            .sum()
    });
    content + tools
}

/// Detect context window overflow errors from the API.
pub fn is_context_overflow(err: &anyhow::Error) -> bool {
    let msg = err.to_string().to_lowercase();
    msg.contains("context length")
        || msg.contains("token limit")
        || msg.contains("too many tokens")
        || msg.contains("maximum context")
        || msg.contains("prompt is too long")
        || msg.contains("request too large")
        || msg.contains("input validation error")
        || msg.contains("content length limit")
        || (msg.contains("400") && msg.contains("tokens"))
}

/// Detect model/provider errors delivered inside the SSE stream.
pub fn is_stream_error(err: &anyhow::Error) -> bool {
    err.to_string().contains("model stream error")
}