Skip to main content

formal_ai/summarization/
dialog.rs

1//! Dialog-aware helpers for the summarization pipeline.
2//!
3//! [`DialogTurn`] models a single user/assistant turn. The
4//! `formalize_dialog` → `summarize_dialog` → `generate_chat_title` chain runs
5//! the same formalize → summarize → deformalize pipeline as the rest of the
6//! module, with a role-aware bias so user turns dominate the output when the
7//! caller asks for a short summary or a chat title.
8
9use super::{
10    deformalize, formalize, summarize, to_topic, Statement, SummarizationConfig, SummarizationMode,
11};
12
13/// A single dialog turn passed to [`summarize_dialog`] /
14/// [`generate_chat_title`]. The role is informational only — the summarizer
15/// uses the text content.
16#[derive(Debug, Clone)]
17pub struct DialogTurn {
18    pub role: String,
19    pub text: String,
20}
21
22impl DialogTurn {
23    /// Build a turn from explicit role + text.
24    #[must_use]
25    pub fn new(role: impl Into<String>, text: impl Into<String>) -> Self {
26        Self {
27            role: role.into(),
28            text: text.into(),
29        }
30    }
31
32    /// Convenience constructor for user turns.
33    #[must_use]
34    pub fn user(text: impl Into<String>) -> Self {
35        Self::new("user", text)
36    }
37
38    /// Convenience constructor for assistant turns.
39    #[must_use]
40    pub fn assistant(text: impl Into<String>) -> Self {
41        Self::new("assistant", text)
42    }
43}
44
45/// Convert dialog turns into [`Statement`]s with role-aware weighting.
46///
47/// Each turn's text is formalized individually; user turns are weighted
48/// higher than assistant turns so a `Short` summary keeps the user's
49/// original questions when both sides are long.
50#[must_use]
51pub fn formalize_dialog(turns: &[DialogTurn]) -> Vec<Statement> {
52    let mut out = Vec::new();
53    for turn in turns {
54        let bias: i16 = match turn.role.as_str() {
55            "user" => 20,
56            "assistant" => -10,
57            _ => 0,
58        };
59        for mut stmt in formalize(&turn.text) {
60            let bumped = i16::from(stmt.weight).saturating_add(bias).clamp(0, 100);
61            stmt.weight = u8::try_from(bumped).unwrap_or(0);
62            out.push(stmt);
63        }
64    }
65    out
66}
67
68/// Summarize a dialog. The output preserves the order of the highest-weight
69/// statements (user questions first when they tie with assistant prose) and
70/// passes through [`deformalize`] for display.
71#[must_use]
72pub fn summarize_dialog(turns: &[DialogTurn], config: &SummarizationConfig) -> String {
73    let statements = formalize_dialog(turns);
74    if config.mode == SummarizationMode::Topic {
75        let highest = statements.iter().max_by_key(|s| s.weight);
76        return highest
77            .map(|s| to_topic("", std::slice::from_ref(s)))
78            .unwrap_or_default();
79    }
80    if statements.is_empty() {
81        return String::new();
82    }
83    let summarized = summarize(&statements, config);
84    deformalize(&summarized)
85}
86
87/// Generate a 1–5 word chat title from a dialog.
88///
89/// Equivalent to running [`summarize_dialog`] in `Topic` mode but spelled
90/// out so the call site reads as `generate_chat_title(turns, "en")` instead
91/// of building a config.
92#[must_use]
93pub fn generate_chat_title(turns: &[DialogTurn], language: &str) -> String {
94    let config = SummarizationConfig::default()
95        .with_mode(SummarizationMode::Topic)
96        .with_language(language);
97    summarize_dialog(turns, &config)
98}