roboticus-agent 0.11.4

//! Topic segmentation for session conversation threads.
//!
//! Detects topic changes within a session so the context builder can
//! prioritize messages from the current topic over stale threads. This
//! prevents the model from confusing concurrent conversation topics
//! (e.g. workspace cleanup vs. security monitoring vs. Go migration)
//! in long-running sessions.
//!
//! ## Architecture
//!
//! - **Detection**: Compares the user's new message against the most recent
//!   assistant response via cosine similarity on embeddings. A drop below
//!   the `TOPIC_CONTINUITY_THRESHOLD` signals a topic shift.
//! - **Tagging**: Each message gets a `topic_tag` (e.g. `"topic-3"`) stored
//!   in the database. Sequential messages on the same topic share a tag.
//! - **Context assembly**: `build_context_with_budget` uses topic tags to
//!   ensure the current topic's messages are prioritized within the token
//!   budget, with older topics getting a compressed summary instead of
//!   full message history.

/// Cosine similarity threshold below which a new message is considered
/// a topic change from the previous conversation thread.
///
/// Typical values:
/// - 0.85+: very similar (same topic, follow-up question)
/// - 0.65-0.85: related but drifting (adjacent topic)
/// - < 0.65: different topic entirely
pub const TOPIC_CONTINUITY_THRESHOLD: f64 = 0.65;

/// Maximum number of sequential turns to check for topic continuity.
/// We compare against the most recent N user messages, not just the last one,
/// to handle cases where the user interleaves topics quickly.
const LOOKBACK_TURNS: usize = 3;

/// Represents a detected topic boundary in a conversation.
#[derive(Debug, Clone)]
pub struct TopicSegment {
    /// Opaque tag identifying this topic thread (e.g. "topic-1", "topic-2").
    pub tag: String,
    /// Whether this message starts a new topic (vs continuing the previous).
    pub is_new_topic: bool,
    /// Similarity score against the previous topic (0.0 if first message).
    pub continuity_score: f64,
}

/// Detect whether a new user message is a continuation of the current topic
/// or a shift to a new one.
///
/// Uses embedding cosine similarity against recent assistant responses.
/// Returns the topic tag to assign to this message.
///
/// `current_tag`: the topic tag of the most recent message (e.g. "topic-2").
///   Pass `None` for the first message in a session.
/// `next_tag_number`: the next available topic number (e.g. 3 if "topic-2" is current).
/// `new_message_embedding`: embedding of the user's new message.
/// `recent_embeddings`: embeddings of the most recent `LOOKBACK_TURNS` messages
///   (from newest to oldest).
pub fn detect_topic_shift(
    current_tag: Option<&str>,
    next_tag_number: u32,
    new_message_embedding: &[f32],
    recent_embeddings: &[Vec<f32>],
) -> TopicSegment {
    if recent_embeddings.is_empty() || current_tag.is_none() {
        // First message or no history — start topic-1.
        return TopicSegment {
            tag: format!("topic-{next_tag_number}"),
            is_new_topic: true,
            continuity_score: 0.0,
        };
    }

    // Compare against the most recent messages and take the MAX similarity.
    // This handles interleaved topics where the user might briefly address
    // one topic before returning to another.
    let max_similarity = recent_embeddings
        .iter()
        .take(LOOKBACK_TURNS)
        .map(|emb| cosine_similarity(new_message_embedding, emb))
        .fold(0.0_f64, f64::max);

    if max_similarity >= TOPIC_CONTINUITY_THRESHOLD {
        // Continuing the current topic.
        TopicSegment {
            tag: current_tag.unwrap_or("topic-1").to_string(),
            is_new_topic: false,
            continuity_score: max_similarity,
        }
    } else {
        // Topic shift detected.
        TopicSegment {
            tag: format!("topic-{next_tag_number}"),
            is_new_topic: true,
            continuity_score: max_similarity,
        }
    }
}

/// Cosine similarity between two embedding vectors.
fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
    if a.len() != b.len() || a.is_empty() {
        return 0.0;
    }
    let (mut dot, mut norm_a, mut norm_b) = (0.0f64, 0.0f64, 0.0f64);
    for (x, y) in a.iter().zip(b.iter()) {
        let (x, y) = (*x as f64, *y as f64);
        dot += x * y;
        norm_a += x * x;
        norm_b += y * y;
    }
    let denom = norm_a.sqrt() * norm_b.sqrt();
    if denom < 1e-10 { 0.0 } else { dot / denom }
}

/// Given a list of messages with topic tags, partition them into
/// current-topic messages and off-topic messages.
///
/// Returns `(current_topic_messages, off_topic_messages)` where both
/// preserve chronological order.
pub fn partition_by_topic<'a, T>(
    messages: &'a [T],
    current_tag: &str,
    get_tag: impl Fn(&T) -> Option<&str>,
) -> (Vec<&'a T>, Vec<&'a T>) {
    let mut current = Vec::new();
    let mut other = Vec::new();
    for msg in messages {
        if get_tag(msg).is_some_and(|t| t == current_tag) {
            current.push(msg);
        } else {
            other.push(msg);
        }
    }
    (current, other)
}

/// Generate a compressed summary line for an off-topic message block.
/// Used in context assembly to represent trimmed topics as a single system note
/// rather than including all their messages verbatim.
pub fn summarize_topic_block(
    topic_tag: &str,
    message_count: usize,
    first_user_msg: &str,
) -> String {
    let snippet: String = first_user_msg.chars().take(80).collect();
    format!("[Earlier topic ({topic_tag}, {message_count} messages): \"{snippet}...\"]")
}

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

    fn make_embedding(seed: f32) -> Vec<f32> {
        // Simple deterministic embedding for testing
        (0..8).map(|i| (seed + i as f32 * 0.1).sin()).collect()
    }

    #[test]
    fn first_message_starts_new_topic() {
        let emb = make_embedding(1.0);
        let result = detect_topic_shift(None, 1, &emb, &[]);
        assert!(result.is_new_topic);
        assert_eq!(result.tag, "topic-1");
    }

    #[test]
    fn similar_message_continues_topic() {
        let emb1 = make_embedding(1.0);
        let emb2 = make_embedding(1.01); // very similar
        let result = detect_topic_shift(Some("topic-1"), 2, &emb2, &[emb1]);
        assert!(!result.is_new_topic);
        assert_eq!(result.tag, "topic-1");
        assert!(result.continuity_score > TOPIC_CONTINUITY_THRESHOLD);
    }

    #[test]
    fn dissimilar_message_starts_new_topic() {
        let emb1 = make_embedding(1.0);
        let emb2 = make_embedding(100.0); // very different
        let result = detect_topic_shift(Some("topic-1"), 2, &emb2, &[emb1]);
        assert!(result.is_new_topic);
        assert_eq!(result.tag, "topic-2");
        assert!(result.continuity_score < TOPIC_CONTINUITY_THRESHOLD);
    }

    #[test]
    fn partition_separates_topics() {
        let msgs = vec![
            ("hello", Some("topic-1")),
            ("world", Some("topic-1")),
            ("new thing", Some("topic-2")),
            ("back to hello", Some("topic-1")),
        ];
        let (current, other) = partition_by_topic(&msgs, "topic-2", |m| m.1);
        assert_eq!(current.len(), 1);
        assert_eq!(other.len(), 3);
    }

    #[test]
    fn summarize_topic_block_formats_correctly() {
        let summary = summarize_topic_block(
            "topic-1",
            5,
            "cleanup Duncan's workspace and remove redundant items",
        );
        assert!(summary.contains("topic-1"));
        assert!(summary.contains("5 messages"));
        assert!(summary.contains("cleanup"));
    }

    #[test]
    fn cosine_similarity_identical_vectors() {
        let v = vec![1.0, 2.0, 3.0];
        let sim = cosine_similarity(&v, &v);
        assert!((sim - 1.0).abs() < 1e-6);
    }

    #[test]
    fn cosine_similarity_orthogonal_vectors() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![0.0, 1.0, 0.0];
        let sim = cosine_similarity(&a, &b);
        assert!(sim.abs() < 1e-6);
    }
}