polyc_agent/

llm_summarizer.rs

//! LLM-backed anchored-iterative [`Summarizer`].
//!
//! The agent crate's [`StubSummarizer`](crate::StubSummarizer) keeps the
//! summarization data path live without a provider. This module is the
//! one-trait swap-in that lands actual compression on long conversations.
//!
//! # Design
//!
//! Anchored iterative summarization (cf. Factory's 36k-message engineering-
//! session eval): each compaction *merges* into the prior summary rather than
//! re-summarizing the whole transcript from scratch. The prior summary is the
//! persistent state that survives every subsequent compaction; the new
//! transcript chunk is the delta. This preserves identifiers, commitments,
//! decisions and errors across compactions instead of losing them as the
//! window slides past them.
//!
//! # Failure mode
//!
//! Provider errors and missing-output are handled fail-soft: the summarizer
//! returns the existing `prior_summary` unchanged and logs a `warn`. The next
//! compaction will retry on the new transcript, and crucially the anchor is
//! never lost. If we returned an empty string on failure the next
//! `reconstruct_history` would still find the *previous* `summary:{uuid}`
//! event in the journal — but a subsequent successful compaction would then
//! anchor against a stale prior, so failing soft to the anchor is the safe
//! shape.

use std::sync::Arc;

use async_trait::async_trait;
use polyc_llm::{
    CompletionRequest, Content as LlmContent, LlmProvider, Message as LlmMessage, Role,
    turn::collect_turn,
};

use crate::Summarizer;

/// System prompt used by [`LlmSummarizer`].
///
/// Pinned in source (not a config knob) so on-disk summaries are reproducible:
/// the prompt that produced a given summary is the prompt encoded in the
/// shipped binary at that revision. If we ever need to evolve this we'll bump
/// it explicitly and accept that older summaries were produced under a prior
/// prompt — that's a feature, not a bug, for the eval story.
const SYSTEM_PROMPT: &str = "You compress agent transcripts for long-running conversations. \
You receive (a) a PRIOR_SUMMARY representing the conversation so far, and \
(b) a TRANSCRIPT chunk that just happened. \
Produce a NEW_SUMMARY that fully replaces PRIOR_SUMMARY going forward: it \
must preserve every commitment, identifier, decision, error, and unresolved \
question from PRIOR_SUMMARY, then merge in the new content from TRANSCRIPT. \
Be terse. No preamble. Maximum 500 words. Do not invent facts.";

/// Provider-backed anchored-iterative [`Summarizer`].
///
/// Wraps any [`LlmProvider`] behind the trait the control plane already
/// consumes. The control plane builds one of these from whatever provider it
/// instantiated for turns and injects it via `AgentSvc::with_summarizer`.
///
/// # Fields
///
/// - `provider` — the same trait the turn loop uses; sharing one `Arc` keeps
///   pooled connections / auth state hot across turns *and* summarizations.
/// - `model` — kept separate from the turn's model so production can point
///   summarization at a cheaper / smaller model (e.g. flash-lite vs flash)
///   without coupling the two upgrade paths.
/// - `max_output_tokens` — cap on the generated summary. The prompt asks for
///   "≤500 words" but the provider is the final guard; the cap is here to
///   protect against a runaway provider regardless of the prompt.
pub struct LlmSummarizer<P: ?Sized> {
    /// The provider used to run the summarization completion.
    provider: Arc<P>,
    /// Model identifier sent to the provider for summarization calls.
    model: String,
    /// Hard cap on output tokens; the prompt also asks for ≤500 words.
    max_output_tokens: u64,
}

impl<P: ?Sized> LlmSummarizer<P> {
    /// Build a new [`LlmSummarizer`] over an existing provider.
    pub fn new(provider: Arc<P>, model: impl Into<String>, max_output_tokens: u64) -> Self {
        Self {
            provider,
            model: model.into(),
            max_output_tokens,
        }
    }
}

/// Render a transcript slice into the user-message body — one `role: text`
/// line per message. Non-text content is rendered as a stable marker so the
/// model sees the call/result happened without us fabricating its content.
fn render_transcript(transcript: &[LlmMessage]) -> String {
    let mut s = String::new();
    for msg in transcript {
        let role = match msg.role {
            Role::Assistant => "assistant",
            Role::Tool => "tool",
            Role::System => "system",
            // `Role` is `#[non_exhaustive]`; default any future variant to
            // `user` so the model still sees the message rather than us
            // refusing to render it.
            _ => "user",
        };
        for content in &msg.content {
            match content {
                LlmContent::Text(t) => {
                    s.push_str(role);
                    s.push_str(": ");
                    s.push_str(t);
                    s.push('\n');
                }
                LlmContent::ToolUse(tc) => {
                    s.push_str(role);
                    s.push_str(": [tool_call name=");
                    s.push_str(&tc.name);
                    s.push_str(" args=");
                    s.push_str(&tc.args_json);
                    s.push_str("]\n");
                }
                LlmContent::ToolResult(tr) => {
                    s.push_str(role);
                    s.push_str(": [tool_result for=");
                    s.push_str(&tr.tool_call_id);
                    s.push_str(" body=");
                    s.push_str(&tr.result_json);
                    s.push_str("]\n");
                }
                LlmContent::Image(_) => {
                    s.push_str(role);
                    s.push_str(": [image]\n");
                }
                // `Content` is `#[non_exhaustive]`; if a future variant lands
                // (audio, video, …) we render a placeholder rather than
                // refusing to summarize and stalling the journal.
                _ => {
                    s.push_str(role);
                    s.push_str(": [unknown]\n");
                }
            }
        }
    }
    s
}

#[async_trait]
impl<P> Summarizer for LlmSummarizer<P>
where
    P: LlmProvider + Send + Sync + ?Sized,
{
    #[tracing::instrument(
        skip_all,
        fields(
            model = %self.model,
            transcript_messages = transcript.len(),
            prior_summary_len = prior_summary.len(),
        ),
    )]
    async fn summarize(&self, prior_summary: &str, transcript: &[LlmMessage]) -> String {
        // Empty-input shortcut: nothing to compress and no anchor to preserve.
        if transcript.is_empty() && prior_summary.is_empty() {
            return String::new();
        }

        let prior_block = if prior_summary.is_empty() {
            "(none — first compaction)".to_owned()
        } else {
            prior_summary.to_owned()
        };
        let transcript_block = if transcript.is_empty() {
            "(empty)".to_owned()
        } else {
            render_transcript(transcript)
        };
        let user_text = format!("PRIOR_SUMMARY:\n{prior_block}\n\nTRANSCRIPT:\n{transcript_block}");

        let mut req = CompletionRequest::new(&self.model);
        req.system = Some(SYSTEM_PROMPT.to_owned());
        req.messages.push(LlmMessage::user(user_text));
        // Cap output to bound the journal write and protect against runaway
        // providers. `max_tokens` is u32 on the request; saturate the cast.
        req.max_tokens = Some(u32::try_from(self.max_output_tokens).unwrap_or(u32::MAX));
        // Low temperature for compaction: summarization is a deterministic
        // rewriting task, not a creative one.
        req.temperature = Some(0.2);

        match self.provider.complete(req).await {
            Ok(stream) => match collect_turn(stream).await {
                Ok(out) => {
                    let trimmed = out.text.trim();
                    if trimmed.is_empty() {
                        tracing::warn!("summarizer received empty output; keeping prior summary");
                        prior_summary.to_owned()
                    } else {
                        trimmed.to_owned()
                    }
                }
                Err(err) => {
                    tracing::warn!(error = %err, "summarizer stream error; keeping prior summary");
                    prior_summary.to_owned()
                }
            },
            Err(err) => {
                tracing::warn!(error = %err, "summarizer provider error; keeping prior summary");
                prior_summary.to_owned()
            }
        }
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use super::*;
    use async_trait::async_trait;
    use futures::{StreamExt, stream};
    use polyc_llm::{Chunk, CompletionRequest, LlmProvider, StopReason, error::DummyError};
    use std::sync::Mutex;

    /// Provider that returns a fixed canned text. Used to verify the
    /// summarizer wires through provider → collect_turn → string.
    struct CannedProvider {
        text: String,
        seen: Mutex<Option<CompletionRequest>>,
    }

    impl CannedProvider {
        fn new(text: &str) -> Self {
            Self {
                text: text.to_owned(),
                seen: Mutex::new(None),
            }
        }
    }

    #[async_trait]
    impl LlmProvider for CannedProvider {
        type Error = DummyError;

        async fn complete(
            &self,
            req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            *self.seen.lock().unwrap() = Some(req);
            let chunks = vec![
                Ok(Chunk::text_delta(self.text.clone())),
                Ok(Chunk::Stop(StopReason::EndTurn)),
            ];
            Ok(stream::iter(chunks).boxed())
        }
    }

    /// Provider that always returns a pre-stream error. Used to verify the
    /// fail-soft path returns `prior_summary` unchanged.
    struct ErroringProvider;

    #[async_trait]
    impl LlmProvider for ErroringProvider {
        type Error = DummyError;

        async fn complete(
            &self,
            _req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            Err(DummyError::Other("nope".to_owned()))
        }
    }

    #[tokio::test]
    async fn returns_provider_output_trimmed() {
        let provider = Arc::new(CannedProvider::new("  the new summary  "));
        let summ = LlmSummarizer::new(provider.clone(), "test-model", 1024);
        let prior = "prior anchor";
        let transcript = vec![LlmMessage::user("hi"), LlmMessage::assistant("hello")];
        let out = summ.summarize(prior, &transcript).await;
        assert_eq!(out, "the new summary");
        // The provider saw a request shaped as expected: a system prompt,
        // one user message, and the configured model + cap.
        let seen = provider.seen.lock().unwrap().clone().expect("request seen");
        assert_eq!(seen.model, "test-model");
        assert!(seen.system.is_some_and(|s| s.contains("PRIOR_SUMMARY")));
        assert_eq!(seen.messages.len(), 1);
        assert_eq!(seen.max_tokens, Some(1024));
        // The rendered user message embeds both PRIOR_SUMMARY and TRANSCRIPT.
        let user_text = match &seen.messages[0].content[0] {
            LlmContent::Text(t) => t.clone(),
            _ => panic!("expected text content"),
        };
        assert!(user_text.contains("PRIOR_SUMMARY:"));
        assert!(user_text.contains("prior anchor"));
        assert!(user_text.contains("TRANSCRIPT:"));
        assert!(user_text.contains("user: hi"));
        assert!(user_text.contains("assistant: hello"));
    }

    #[tokio::test]
    async fn first_compaction_uses_none_marker() {
        let provider = Arc::new(CannedProvider::new("first summary"));
        let summ = LlmSummarizer::new(provider.clone(), "test-model", 256);
        let out = summ
            .summarize("", &[LlmMessage::user("a"), LlmMessage::assistant("b")])
            .await;
        assert_eq!(out, "first summary");
        let seen = provider.seen.lock().unwrap().clone().expect("request seen");
        let user_text = match &seen.messages[0].content[0] {
            LlmContent::Text(t) => t.clone(),
            _ => panic!("expected text content"),
        };
        assert!(user_text.contains("(none — first compaction)"));
    }

    #[tokio::test]
    async fn provider_error_returns_prior_summary_unchanged() {
        let summ = LlmSummarizer::new(Arc::new(ErroringProvider), "test-model", 256);
        let prior = "this is the anchor";
        let out = summ.summarize(prior, &[LlmMessage::user("hi")]).await;
        assert_eq!(
            out, prior,
            "fail-soft: prior summary survives provider errors"
        );
    }

    #[tokio::test]
    async fn empty_inputs_short_circuit() {
        // No provider call should be made when both inputs are empty.
        let summ = LlmSummarizer::new(Arc::new(ErroringProvider), "test-model", 256);
        let out = summ.summarize("", &[]).await;
        assert!(out.is_empty());
    }

    #[tokio::test]
    async fn empty_output_falls_back_to_prior() {
        // A provider that yields no text deltas (just stop). Fail-soft keeps
        // the prior anchor rather than overwriting it with "".
        struct EmptyProvider;
        #[async_trait]
        impl LlmProvider for EmptyProvider {
            type Error = DummyError;
            async fn complete(
                &self,
                _req: CompletionRequest,
            ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
            {
                Ok(stream::iter(vec![Ok(Chunk::Stop(StopReason::EndTurn))]).boxed())
            }
        }
        let summ = LlmSummarizer::new(Arc::new(EmptyProvider), "test-model", 256);
        let out = summ.summarize("keep me", &[LlmMessage::user("x")]).await;
        assert_eq!(out, "keep me");
    }
}