diff --git a/runtime/csharp/Prompty.Core.Tests/CompactionTests.cs b/runtime/csharp/Prompty.Core.Tests/CompactionTests.cs new file mode 100644 index 00000000..79685346 --- /dev/null +++ b/runtime/csharp/Prompty.Core.Tests/CompactionTests.cs @@ -0,0 +1,270 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Prompty.Core; + +namespace Prompty.Core.Tests; + +public class CompactionTests +{ + private static Message TextMsg(string role, string text) => + new() { Role = role, Parts = [new TextPart { Value = text }] }; + + private static Message ToolCallMsg(string name, string args) => + new() + { + Role = "assistant", + Parts = [new TextPart { Value = "" }], + Metadata = new Dictionary + { + ["tool_calls"] = new List + { + new() { Id = "tc1", Name = name, Arguments = args } + } + } + }; + + // ----------------------------------------------------------------------- + // FormatDroppedMessages + // ----------------------------------------------------------------------- + + [Fact] + public void FormatDroppedMessages_FormatsReadableText() + { + var dropped = new List + { + TextMsg("user", "What is 2+2?"), + TextMsg("assistant", "It is 4."), + }; + + var result = ContextWindow.FormatDroppedMessages(dropped); + + Assert.Contains("[user]: What is 2+2?", result); + Assert.Contains("[assistant]: It is 4.", result); + } + + [Fact] + public void FormatDroppedMessages_ShowsToolCalls() + { + var dropped = new List + { + ToolCallMsg("get_weather", "{\"city\":\"Seattle\"}") + }; + + var result = ContextWindow.FormatDroppedMessages(dropped); + + Assert.Contains("[assistant]: Called: get_weather(", result); + Assert.Contains("Seattle", result); + } + + [Fact] + public void FormatDroppedMessages_EmptyForEmptyList() + { + var result = ContextWindow.FormatDroppedMessages(new List()); + Assert.Equal("", result); + } + + // ----------------------------------------------------------------------- + // CompactionStrategy factory methods + // ----------------------------------------------------------------------- + + [Fact] + public void FromFunction_CreatesStrategy() + { + var strategy = CompactionStrategy.FromFunction( + msgs => Task.FromResult("summary")); + Assert.NotNull(strategy); + } + + [Fact] + public void FromPrompty_CreatesStrategy() + { + var strategy = CompactionStrategy.FromPrompty("some/path.prompty"); + Assert.NotNull(strategy); + } + + [Fact] + public void FromPrompty_ThrowsOnEmpty() + { + Assert.Throws(() => + CompactionStrategy.FromPrompty("")); + } + + [Fact] + public void FromFunction_ThrowsOnNull() + { + Assert.Throws(() => + CompactionStrategy.FromFunction(null!)); + } + + // ----------------------------------------------------------------------- + // ApplyCompactionAsync — function path + // ----------------------------------------------------------------------- + + [Fact] + public async Task CompactionFunction_ReplacesSummary() + { + // Arrange: build messages that simulate a trimmed conversation + // with a default summary already in place + var messages = new List + { + TextMsg("system", "You are a helpful assistant."), + TextMsg("user", "[Context summary: User asked: old question]"), + TextMsg("user", "latest question"), + }; + + var dropped = new List + { + TextMsg("user", "old question"), + TextMsg("assistant", "old answer"), + }; + + var compaction = CompactionStrategy.FromFunction( + msgs => Task.FromResult("Compacted: discussed old question and got old answer")); + + // Act + await Pipeline.ApplyCompactionAsync(compaction, dropped, messages, null); + + // Assert: the summary message was replaced + Assert.Equal("Compacted: discussed old question and got old answer", messages[1].Text); + } + + [Fact] + public async Task CompactionFailure_PreservesDefault() + { + var messages = new List + { + TextMsg("system", "system"), + TextMsg("user", "[Context summary: original summary]"), + TextMsg("user", "latest"), + }; + + var dropped = new List { TextMsg("user", "old") }; + + var compaction = CompactionStrategy.FromFunction( + _ => throw new InvalidOperationException("LLM unavailable")); + + var events = new List<(AgentEventType Type, Dictionary Data)>(); + EventCallback onEvent = (type, data) => events.Add((type, data)); + + // Act + await Pipeline.ApplyCompactionAsync(compaction, dropped, messages, onEvent); + + // Assert: original summary preserved + Assert.Equal("[Context summary: original summary]", messages[1].Text); + + // Assert: failure event emitted + Assert.Contains(events, e => e.Type == AgentEventType.CompactionFailed); + var failedEvent = events.First(e => e.Type == AgentEventType.CompactionFailed); + Assert.Contains("LLM unavailable", failedEvent.Data["reason"]?.ToString()); + } + + [Fact] + public async Task CompactionEmpty_EmitsFailedEvent() + { + var messages = new List + { + TextMsg("system", "system"), + TextMsg("user", "[Context summary: original]"), + TextMsg("user", "latest"), + }; + + var dropped = new List { TextMsg("user", "old") }; + + var compaction = CompactionStrategy.FromFunction( + _ => Task.FromResult("")); + + var events = new List<(AgentEventType Type, Dictionary Data)>(); + EventCallback onEvent = (type, data) => events.Add((type, data)); + + await Pipeline.ApplyCompactionAsync(compaction, dropped, messages, onEvent); + + // Original preserved + Assert.Equal("[Context summary: original]", messages[1].Text); + + // CompactionFailed with "empty result" + var failedEvent = events.First(e => e.Type == AgentEventType.CompactionFailed); + Assert.Equal("empty result", failedEvent.Data["reason"]?.ToString()); + } + + // ----------------------------------------------------------------------- + // Event emission + // ----------------------------------------------------------------------- + + [Fact] + public async Task CompactionEvents_AreEmitted() + { + var messages = new List + { + TextMsg("system", "system"), + TextMsg("user", "[Context summary: original]"), + TextMsg("user", "latest"), + }; + + var dropped = new List { TextMsg("user", "old") }; + + var compaction = CompactionStrategy.FromFunction( + _ => Task.FromResult("New compacted summary")); + + var events = new List<(AgentEventType Type, Dictionary Data)>(); + EventCallback onEvent = (type, data) => events.Add((type, data)); + + await Pipeline.ApplyCompactionAsync(compaction, dropped, messages, onEvent); + + // Both start and complete events emitted + Assert.Contains(events, e => e.Type == AgentEventType.CompactionStart); + Assert.Contains(events, e => e.Type == AgentEventType.CompactionComplete); + + var startEvent = events.First(e => e.Type == AgentEventType.CompactionStart); + Assert.Equal(1, startEvent.Data["dropped_count"]); + + var completeEvent = events.First(e => e.Type == AgentEventType.CompactionComplete); + Assert.Equal("New compacted summary".Length, completeEvent.Data["summary_length"]); + } + + // ----------------------------------------------------------------------- + // ReplaceSummaryMessage + // ----------------------------------------------------------------------- + + [Fact] + public void ReplaceSummaryMessage_ReplacesFirstMatch() + { + var messages = new List + { + TextMsg("system", "system prompt"), + TextMsg("user", "[Context summary: old info]"), + TextMsg("user", "new question"), + }; + + Pipeline.ReplaceSummaryMessage(messages, "Better summary"); + + Assert.Equal("Better summary", messages[1].Text); + Assert.Equal("user", messages[1].Role); + } + + [Fact] + public void ReplaceSummaryMessage_NoMatchLeavesListUnchanged() + { + var messages = new List + { + TextMsg("system", "system prompt"), + TextMsg("user", "just a normal message"), + }; + + Pipeline.ReplaceSummaryMessage(messages, "Better summary"); + + Assert.Equal("just a normal message", messages[1].Text); + } + + // ----------------------------------------------------------------------- + // Null compaction preserves existing behavior + // ----------------------------------------------------------------------- + + [Fact] + public void NullCompaction_DefaultParameterValue() + { + // Verify the parameter defaults to null — this is a compile-time check. + // If this compiles, the default is correct. + // We just verify the method signature accepts no compaction argument. + Assert.Null((CompactionStrategy?)null); + } +} diff --git a/runtime/csharp/Prompty.Core/AgentEvents.cs b/runtime/csharp/Prompty.Core/AgentEvents.cs index 32a95a0f..f2a272cd 100644 --- a/runtime/csharp/Prompty.Core/AgentEvents.cs +++ b/runtime/csharp/Prompty.Core/AgentEvents.cs @@ -13,7 +13,10 @@ public enum AgentEventType MessagesUpdated, Done, Error, - Cancelled + Cancelled, + CompactionStart, + CompactionComplete, + CompactionFailed } /// Callback for agent loop events. diff --git a/runtime/csharp/Prompty.Core/CompactionStrategy.cs b/runtime/csharp/Prompty.Core/CompactionStrategy.cs new file mode 100644 index 00000000..5297b2c0 --- /dev/null +++ b/runtime/csharp/Prompty.Core/CompactionStrategy.cs @@ -0,0 +1,50 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Prompty.Core; + +/// +/// Context compaction strategy for replacing default dropped-message summaries +/// with higher-quality LLM-powered or function-based summaries. +/// +public abstract class CompactionStrategy +{ + /// Produce a compact summary from dropped messages. + public abstract Task CompactAsync(IReadOnlyList dropped); + + /// Create compaction from a .prompty file path. + public static CompactionStrategy FromPrompty(string path) => new PromptyCompaction(path); + + /// Create compaction from a function. + public static CompactionStrategy FromFunction(Func, Task> fn) => new FunctionCompaction(fn); +} + +internal sealed class PromptyCompaction : CompactionStrategy +{ + private readonly string _path; + + public PromptyCompaction(string path) + { + ArgumentException.ThrowIfNullOrWhiteSpace(path); + _path = path; + } + + public override async Task CompactAsync(IReadOnlyList dropped) + { + var text = ContextWindow.FormatDroppedMessages(dropped.ToList()); + var result = await Pipeline.InvokeAsync(_path, new Dictionary { ["messages"] = text }); + return result?.ToString() ?? ""; + } +} + +internal sealed class FunctionCompaction : CompactionStrategy +{ + private readonly Func, Task> _fn; + + public FunctionCompaction(Func, Task> fn) + { + ArgumentNullException.ThrowIfNull(fn); + _fn = fn; + } + + public override Task CompactAsync(IReadOnlyList dropped) => _fn(dropped); +} diff --git a/runtime/csharp/Prompty.Core/ContextWindow.cs b/runtime/csharp/Prompty.Core/ContextWindow.cs index c43bd4fb..250fc6f7 100644 --- a/runtime/csharp/Prompty.Core/ContextWindow.cs +++ b/runtime/csharp/Prompty.Core/ContextWindow.cs @@ -119,6 +119,32 @@ public static (int DroppedCount, List Dropped) TrimToContextWindow( return (droppedCount, dropped); } + /// + /// Format dropped messages as readable text for compaction input. + /// Each message is rendered as "[role]: text" with tool calls shown as "Called: name(args)". + /// + public static string FormatDroppedMessages(List messages) + { + var lines = new List(); + foreach (var msg in messages) + { + if (msg.Metadata.TryGetValue("tool_calls", out var toolCalls) && toolCalls is System.Collections.IEnumerable tcList) + { + foreach (var tc in tcList) + { + var name = tc?.GetType().GetProperty("Name")?.GetValue(tc)?.ToString() ?? "unknown"; + var args = tc?.GetType().GetProperty("Arguments")?.GetValue(tc)?.ToString() ?? ""; + lines.Add($"[{msg.Role}]: Called: {name}({args})"); + } + } + + var text = msg.Text.Trim(); + if (text.Length > 0) + lines.Add($"[{msg.Role}]: {text}"); + } + return string.Join("\n", lines); + } + private static string Truncate(string text, int maxLen = 200) => text.Length <= maxLen ? text : text[..maxLen] + "…"; } diff --git a/runtime/csharp/Prompty.Core/Pipeline.cs b/runtime/csharp/Prompty.Core/Pipeline.cs index 6db585c0..174bf065 100644 --- a/runtime/csharp/Prompty.Core/Pipeline.cs +++ b/runtime/csharp/Prompty.Core/Pipeline.cs @@ -232,7 +232,8 @@ public static async Task TurnAsync( Guardrails? guardrails = null, Steering? steering = null, bool parallelToolCalls = false, - int maxLlmRetries = 3) + int maxLlmRetries = 3, + CompactionStrategy? compaction = null) { var label = turnNumber.HasValue ? $"turn {turnNumber.Value}" : "turn"; return await Trace.TraceAsync($"prompty.turn", async (emit) => @@ -292,11 +293,16 @@ public static async Task TurnAsync( // Context window trimming if (contextBudget is not null) { - var (droppedCount, _) = ContextWindow.TrimToContextWindow(messages, contextBudget.Value); + var (droppedCount, droppedMessages) = ContextWindow.TrimToContextWindow(messages, contextBudget.Value); if (droppedCount > 0) { AgentEvents.EmitEvent(onEvent, AgentEventType.MessagesUpdated, new Dictionary { ["source"] = "context_trim", ["dropped"] = droppedCount }); + + if (compaction is not null) + { + await ApplyCompactionAsync(compaction, droppedMessages, messages, onEvent); + } } } @@ -536,11 +542,12 @@ public static async Task TurnAsync( Guardrails? guardrails = null, Steering? steering = null, bool parallelToolCalls = false, - int maxLlmRetries = 3) + int maxLlmRetries = 3, + CompactionStrategy? compaction = null) { var agent = PromptyLoader.Load(path); return await TurnAsync(agent, inputs, tools, maxIterations, raw, turnNumber, onEvent, - cancellationToken, contextBudget, guardrails, steering, parallelToolCalls, maxLlmRetries); + cancellationToken, contextBudget, guardrails, steering, parallelToolCalls, maxLlmRetries, compaction); } // ----------------------------------------------------------------------- @@ -581,10 +588,11 @@ public static async Task TurnAsync( Guardrails? guardrails = null, Steering? steering = null, bool parallelToolCalls = false, - int maxLlmRetries = 3) + int maxLlmRetries = 3, + CompactionStrategy? compaction = null) { var result = await TurnAsync(agent, inputs, tools, maxIterations, raw, turnNumber, onEvent, - cancellationToken, contextBudget, guardrails, steering, parallelToolCalls, maxLlmRetries); + cancellationToken, contextBudget, guardrails, steering, parallelToolCalls, maxLlmRetries, compaction); return PromptyCast.Cast(result); } @@ -604,13 +612,71 @@ public static async Task TurnAsync( Guardrails? guardrails = null, Steering? steering = null, bool parallelToolCalls = false, - int maxLlmRetries = 3) + int maxLlmRetries = 3, + CompactionStrategy? compaction = null) { var result = await TurnAsync(path, inputs, tools, maxIterations, raw, turnNumber, onEvent, - cancellationToken, contextBudget, guardrails, steering, parallelToolCalls, maxLlmRetries); + cancellationToken, contextBudget, guardrails, steering, parallelToolCalls, maxLlmRetries, compaction); return PromptyCast.Cast(result); } + // ----------------------------------------------------------------------- + // Compaction Helper + // ----------------------------------------------------------------------- + + /// + /// Apply a compaction strategy to replace the default dropped-message summary. + /// On failure, the existing SummarizeDropped summary is preserved. + /// + internal static async Task ApplyCompactionAsync( + CompactionStrategy compaction, + List dropped, + List messages, + EventCallback? onEvent) + { + AgentEvents.EmitEvent(onEvent, AgentEventType.CompactionStart, + new Dictionary { ["dropped_count"] = dropped.Count }); + try + { + var summary = await compaction.CompactAsync(dropped); + if (!string.IsNullOrWhiteSpace(summary)) + { + ReplaceSummaryMessage(messages, summary); + AgentEvents.EmitEvent(onEvent, AgentEventType.CompactionComplete, + new Dictionary { ["summary_length"] = summary.Length }); + } + else + { + AgentEvents.EmitEvent(onEvent, AgentEventType.CompactionFailed, + new Dictionary { ["reason"] = "empty result" }); + } + } + catch (Exception ex) + { + AgentEvents.EmitEvent(onEvent, AgentEventType.CompactionFailed, + new Dictionary { ["reason"] = ex.Message }); + } + } + + /// + /// Replace the first "[Context summary:" message with a compaction-produced summary. + /// + internal static void ReplaceSummaryMessage(List messages, string newSummary) + { + for (int i = 0; i < messages.Count; i++) + { + if (messages[i].Text.StartsWith("[Context summary:")) + { + messages[i] = new Message + { + Role = messages[i].Role, + Parts = [new TextPart { Value = newSummary }] + }; + return; + } + } + } + // ----------------------------------------------------------------------- // LLM Retry Helper (§9.10) // ----------------------------------------------------------------------- diff --git a/runtime/python/prompty/prompty/__init__.py b/runtime/python/prompty/prompty/__init__.py index b1eb6fd8..5334da2b 100644 --- a/runtime/python/prompty/prompty/__init__.py +++ b/runtime/python/prompty/prompty/__init__.py @@ -104,6 +104,7 @@ "CancelledError", "ExecuteError", "estimate_chars", + "format_dropped_messages", "summarize_dropped", "trim_to_context_window", "GuardrailError", @@ -127,7 +128,7 @@ from .core.agent_events import AgentEvent, EventCallback, emit_event from .core.cancellation import CancellationToken, CancelledError from .core.connections import clear_connections, get_connection, register_connection -from .core.context import estimate_chars, summarize_dropped, trim_to_context_window +from .core.context import estimate_chars, format_dropped_messages, summarize_dropped, trim_to_context_window from .core.guardrails import GuardrailError, GuardrailResult, Guardrails # Loader diff --git a/runtime/python/prompty/prompty/core/__init__.py b/runtime/python/prompty/prompty/core/__init__.py index 8f03fa8e..752680f1 100644 --- a/runtime/python/prompty/prompty/core/__init__.py +++ b/runtime/python/prompty/prompty/core/__init__.py @@ -5,7 +5,7 @@ from .agent_events import AgentEvent, EventCallback, emit_event from .cancellation import CancellationToken, CancelledError from .connections import clear_connections, get_connection, register_connection -from .context import estimate_chars, summarize_dropped, trim_to_context_window +from .context import estimate_chars, format_dropped_messages, summarize_dropped, trim_to_context_window from .discovery import ( InvokerError, clear_cache, diff --git a/runtime/python/prompty/prompty/core/context.py b/runtime/python/prompty/prompty/core/context.py index 55b9ccf4..d50ed9df 100644 --- a/runtime/python/prompty/prompty/core/context.py +++ b/runtime/python/prompty/prompty/core/context.py @@ -10,11 +10,36 @@ __all__ = [ "estimate_chars", + "format_dropped_messages", "summarize_dropped", "trim_to_context_window", ] +def format_dropped_messages(messages: list[Message]) -> str: + """Format dropped messages as readable text for compaction prompts. + + Each message is rendered as ``[role]: text`` with tool calls shown + as ``Called: name(args)``. + """ + import json + + lines: list[str] = [] + for msg in messages: + text = msg.text.strip() if msg.text else "" + if text: + lines.append(f"[{msg.role}]: {text}") + tool_calls = msg.metadata.get("tool_calls") + if tool_calls and isinstance(tool_calls, list): + for tc in tool_calls: + name = tc.get("name", tc.get("function", {}).get("name", "?")) + args = tc.get("arguments", tc.get("function", {}).get("arguments", "")) + if isinstance(args, dict): + args = json.dumps(args) + lines.append(f"Called: {name}({args})") + return "\n".join(lines) + + def estimate_chars(messages: list[Message]) -> int: """Estimate the character cost of a message list. diff --git a/runtime/python/prompty/prompty/core/pipeline.py b/runtime/python/prompty/prompty/core/pipeline.py index 7fd99541..c3f61ede 100644 --- a/runtime/python/prompty/prompty/core/pipeline.py +++ b/runtime/python/prompty/prompty/core/pipeline.py @@ -24,11 +24,13 @@ from __future__ import annotations import asyncio +import inspect import json import random import time from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Any from ..model import Prompty @@ -77,6 +79,10 @@ "_expand_thread_markers", "_dict_to_message", "_dict_content_to_part", + # Compaction helpers (used by tests) + "_apply_compaction", + "_apply_compaction_async", + "_replace_summary_message", ] @@ -873,6 +879,7 @@ def turn( on_event: EventCallback | None = None, cancel: CancellationToken | None = None, context_budget: int | None = None, + compaction: str | Path | Callable[..., Any] | None = None, guardrails: Guardrails | None = None, steering: Steering | None = None, parallel_tool_calls: bool = False, @@ -909,6 +916,11 @@ def turn( Optional cancellation token for cooperative cancellation. context_budget: Character budget for context window management. ``None`` = no trimming. + compaction: + Optional compaction strategy for when messages are trimmed. A string + or :class:`~pathlib.Path` is treated as a ``.prompty`` file to invoke. + A callable receives ``(dropped_messages)`` and returns a summary + string. ``None`` (default) keeps the built-in summary. guardrails: Optional validation hooks (input, output, tool). steering: @@ -982,10 +994,12 @@ def turn( emit_event(on_event, "status", {"message": f"Injected {len(pending)} steering message(s)"}) if context_budget is not None: - dropped_count, _ = trim_to_context_window(messages, context_budget) + dropped_count, dropped = trim_to_context_window(messages, context_budget) if dropped_count > 0: emit_event(on_event, "messages_updated", {"messages": messages}) emit_event(on_event, "status", {"message": f"Trimmed {dropped_count} messages for context budget"}) + if compaction is not None: + _apply_compaction(compaction, dropped, messages, on_event, agent) if guardrails is not None: gr_input = guardrails.check_input(messages) @@ -1120,6 +1134,7 @@ async def turn_async( on_event: EventCallback | None = None, cancel: CancellationToken | None = None, context_budget: int | None = None, + compaction: str | Path | Callable[..., Any] | None = None, guardrails: Guardrails | None = None, steering: Steering | None = None, parallel_tool_calls: bool = False, @@ -1176,10 +1191,12 @@ async def turn_async( emit_event(on_event, "status", {"message": f"Injected {len(pending)} steering message(s)"}) if context_budget is not None: - dropped_count, _ = trim_to_context_window(messages, context_budget) + dropped_count, dropped = trim_to_context_window(messages, context_budget) if dropped_count > 0: emit_event(on_event, "messages_updated", {"messages": messages}) emit_event(on_event, "status", {"message": f"Trimmed {dropped_count} messages for context budget"}) + if compaction is not None: + await _apply_compaction_async(compaction, dropped, messages, on_event, agent) if guardrails is not None: gr_input = guardrails.check_input(messages) @@ -1303,6 +1320,88 @@ async def turn_async( return processed_result +# --------------------------------------------------------------------------- +# Compaction helpers +# --------------------------------------------------------------------------- + + +def _replace_summary_message(messages: list[Message], summary: str) -> None: + """Replace the synthetic summary user message with the compacted version.""" + for i, msg in enumerate(messages): + if msg.role == "user" and msg.text and "[Context summary:" in msg.text: + messages[i] = Message( + role="user", + parts=[TextPart(value=f"[Context summary: {summary}]")], + ) + return + + +def _apply_compaction( + compaction: str | Path | Callable[..., Any], + dropped: list[Message], + messages: list[Message], + on_event: EventCallback | None, + agent: Any, +) -> None: + """Replace the default summary message with a compacted one.""" + from .context import format_dropped_messages + + emit_event(on_event, "compaction_start", {"dropped_count": len(dropped)}) + try: + if isinstance(compaction, (str, Path)): + text = format_dropped_messages(dropped) + summary = invoke(str(compaction), inputs={"messages": text}) + if not isinstance(summary, str): + summary = str(summary) + elif callable(compaction): + result = compaction(dropped) + summary = str(result) if result is not None else "" + else: + return + + if summary and summary.strip(): + _replace_summary_message(messages, summary) + emit_event(on_event, "compaction_complete", {"summary_length": len(summary)}) + else: + emit_event(on_event, "compaction_failed", {"reason": "empty result"}) + except Exception as exc: + emit_event(on_event, "compaction_failed", {"reason": str(exc)}) + + +async def _apply_compaction_async( + compaction: str | Path | Callable[..., Any], + dropped: list[Message], + messages: list[Message], + on_event: EventCallback | None, + agent: Any, +) -> None: + """Async version of _apply_compaction.""" + from .context import format_dropped_messages + + emit_event(on_event, "compaction_start", {"dropped_count": len(dropped)}) + try: + if isinstance(compaction, (str, Path)): + text = format_dropped_messages(dropped) + summary = await invoke_async(str(compaction), inputs={"messages": text}) + if not isinstance(summary, str): + summary = str(summary) + elif callable(compaction): + result = compaction(dropped) + if inspect.isawaitable(result): + result = await result + summary = str(result) if result is not None else "" + else: + return + + if summary and summary.strip(): + _replace_summary_message(messages, summary) + emit_event(on_event, "compaction_complete", {"summary_length": len(summary)}) + else: + emit_event(on_event, "compaction_failed", {"reason": "empty result"}) + except Exception as exc: + emit_event(on_event, "compaction_failed", {"reason": str(exc)}) + + def _resolve_bindings( agent: Any, fn_name: str, diff --git a/runtime/python/prompty/tests/test_agent_extensions.py b/runtime/python/prompty/tests/test_agent_extensions.py index f679b561..965acefe 100644 --- a/runtime/python/prompty/tests/test_agent_extensions.py +++ b/runtime/python/prompty/tests/test_agent_extensions.py @@ -749,3 +749,205 @@ def executor_with_cancel(a, msgs): # Should have made at least 2 executor calls before cancellation assert iteration_count >= 2 + + +# ========================================================================= +# Context Compaction +# ========================================================================= + + +class TestContextCompaction: + """Tests for context compaction via the compaction parameter.""" + + def test_compaction_function_replaces_summary(self): + """When a compaction function is provided and trimming occurs, the summary is replaced.""" + msgs = [ + Message(role="system", parts=[TextPart(value="You are a helper.")]), + Message(role="user", parts=[TextPart(value="First question " * 100)]), + Message(role="assistant", parts=[TextPart(value="First answer " * 100)]), + Message(role="user", parts=[TextPart(value="Second question " * 100)]), + Message(role="assistant", parts=[TextPart(value="Second answer " * 100)]), + Message(role="user", parts=[TextPart(value="Current question")]), + ] + + def my_compactor(dropped: list[Message]) -> str: + return f"COMPACTED: {len(dropped)} messages summarized" + + dropped_count, dropped = trim_to_context_window(msgs, 500) + assert dropped_count > 0 + + from prompty.core.pipeline import _replace_summary_message + + summary = my_compactor(dropped) + _replace_summary_message(msgs, summary) + + summary_msgs = [m for m in msgs if m.role == "user" and "[Context summary:" in (m.text or "")] + assert len(summary_msgs) == 1 + assert "COMPACTED:" in (summary_msgs[0].text or "") + + def test_compaction_function_failure_preserves_default(self): + """When compaction fails, the default summarize_dropped result stays.""" + from prompty.core.pipeline import _apply_compaction + + def failing_compactor(dropped: list[Message]) -> str: + raise RuntimeError("Compaction failed!") + + events: list[tuple[str, Any]] = [] + + def on_event(event_type: str, data: Any) -> None: + events.append((event_type, data)) + + msgs = [ + Message(role="system", parts=[TextPart(value="System.")]), + Message(role="user", parts=[TextPart(value="Q " * 200)]), + Message(role="assistant", parts=[TextPart(value="A " * 200)]), + Message(role="user", parts=[TextPart(value="Current")]), + ] + dropped_count, dropped = trim_to_context_window(msgs, 300) + if dropped_count > 0: + _apply_compaction(failing_compactor, dropped, msgs, on_event, None) + + assert any(e[0] == "compaction_failed" for e in events) + + def test_compaction_none_leaves_default(self): + """When compaction is None, default behavior is preserved.""" + msgs = [ + Message(role="system", parts=[TextPart(value="System.")]), + Message(role="user", parts=[TextPart(value="Q " * 200)]), + Message(role="user", parts=[TextPart(value="Current")]), + ] + dropped_count, dropped = trim_to_context_window(msgs, 200) + # Default summary should exist when messages were dropped + if dropped_count > 0: + summary_msgs = [m for m in msgs if m.role == "user" and "[Context summary:" in (m.text or "")] + assert len(summary_msgs) == 1 + + def test_format_dropped_messages(self): + """format_dropped_messages produces readable text from dropped messages.""" + from prompty.core.context import format_dropped_messages + + msgs = [ + Message(role="user", parts=[TextPart(value="What is 2+2?")]), + Message( + role="assistant", + parts=[TextPart(value="4")], + metadata={"tool_calls": [{"name": "calculator", "arguments": '{"expr": "2+2"}'}]}, + ), + ] + text = format_dropped_messages(msgs) + assert "[user]: What is 2+2?" in text + assert "[assistant]: 4" in text + assert "Called: calculator" in text + + def test_compaction_events_emitted(self): + """Verify compaction_start, compaction_complete events are emitted.""" + from prompty.core.pipeline import _apply_compaction + + events: list[tuple[str, Any]] = [] + + def on_event(event_type: str, data: Any) -> None: + events.append((event_type, data)) + + def good_compactor(dropped: list[Message]) -> str: + return "Summary of dropped messages" + + msgs = [ + Message(role="system", parts=[TextPart(value="Sys.")]), + Message(role="user", parts=[TextPart(value="Old " * 200)]), + Message(role="assistant", parts=[TextPart(value="Old " * 200)]), + Message(role="user", parts=[TextPart(value="Current")]), + ] + dropped_count, dropped = trim_to_context_window(msgs, 300) + if dropped_count > 0: + _apply_compaction(good_compactor, dropped, msgs, on_event, None) + + event_types = [e[0] for e in events] + assert "compaction_start" in event_types + assert "compaction_complete" in event_types + + def test_compaction_empty_result_emits_failed(self): + """When compaction returns an empty string, compaction_failed is emitted.""" + from prompty.core.pipeline import _apply_compaction + + events: list[tuple[str, Any]] = [] + + def on_event(event_type: str, data: Any) -> None: + events.append((event_type, data)) + + def empty_compactor(dropped: list[Message]) -> str: + return "" + + msgs = [ + Message(role="system", parts=[TextPart(value="Sys.")]), + Message(role="user", parts=[TextPart(value="Old " * 200)]), + Message(role="assistant", parts=[TextPart(value="Old " * 200)]), + Message(role="user", parts=[TextPart(value="Current")]), + ] + dropped_count, dropped = trim_to_context_window(msgs, 300) + if dropped_count > 0: + _apply_compaction(empty_compactor, dropped, msgs, on_event, None) + + event_types = [e[0] for e in events] + assert "compaction_failed" in event_types + + def test_replace_summary_message_no_op_when_missing(self): + """_replace_summary_message is a no-op when no summary message exists.""" + from prompty.core.pipeline import _replace_summary_message + + msgs = [ + Message(role="system", parts=[TextPart(value="System.")]), + Message(role="user", parts=[TextPart(value="Hello")]), + ] + original = [m.text for m in msgs] + _replace_summary_message(msgs, "replacement") + assert [m.text for m in msgs] == original + + @patch(f"{_PIPELINE}._invoke_executor") + @patch(f"{_PIPELINE}.prepare", return_value=[Message(role="user", parts=[TextPart(value="hi")])]) + @patch(f"{_PIPELINE}.process", return_value="result") + def test_turn_compaction_integrated(self, mock_process, mock_prepare, mock_exec): + """turn() accepts compaction parameter and applies it during trimming.""" + mock_exec.return_value = _mock_final_response("result") + agent = _make_agent() + events: list[tuple[str, Any]] = [] + + def on_event(event_type: str, data: Any) -> None: + events.append((event_type, data)) + + def my_compactor(dropped: list[Message]) -> str: + return "compacted summary" + + # With a very small budget no trimming occurs on a single message, so + # compaction won't be triggered. This test verifies the parameter is + # accepted without error and the function runs to completion. + result = turn( + agent, + {}, + tools={}, + on_event=on_event, + context_budget=100000, + compaction=my_compactor, + ) + assert result == "result" + + @pytest.mark.asyncio + @patch(f"{_PIPELINE}._invoke_executor_async") + @patch( + f"{_PIPELINE}.prepare_async", + return_value=[Message(role="user", parts=[TextPart(value="hi")])], + ) + @patch(f"{_PIPELINE}.process_async", return_value="async result") + async def test_turn_async_compaction_integrated(self, mock_process, mock_prepare, mock_exec): + """turn_async() accepts compaction parameter.""" + mock_exec.return_value = _mock_final_response("async result") + agent = _make_agent() + + result = await turn_async( + agent, + {}, + tools={}, + on_event=lambda t, d: None, + context_budget=100000, + compaction=lambda dropped: "async compacted", + ) + assert result == "async result" diff --git a/runtime/rust/prompty/src/context.rs b/runtime/rust/prompty/src/context.rs index c10e5057..7635b095 100644 --- a/runtime/rust/prompty/src/context.rs +++ b/runtime/rust/prompty/src/context.rs @@ -89,7 +89,7 @@ pub fn summarize_dropped(messages: &[Message]) -> String { /// Trim messages to fit within a character budget. /// -/// Returns `(dropped_count, trimmed_messages)`. +/// Returns `(dropped_messages, trimmed_messages)`. /// /// Behavior (matches TypeScript): /// 1. System messages at the start are always preserved @@ -99,10 +99,13 @@ pub fn summarize_dropped(messages: &[Message]) -> String { /// after the system messages /// /// The summary budget is `min(5000, 5% of budget_chars)`. -pub fn trim_to_context_window(messages: &[Message], budget_chars: usize) -> (usize, Vec) { +pub fn trim_to_context_window( + messages: &[Message], + budget_chars: usize, +) -> (Vec, Vec) { let current = estimate_chars(messages); if current <= budget_chars { - return (0, messages.to_vec()); + return (vec![], messages.to_vec()); } // Split: leading system messages + rest @@ -115,7 +118,7 @@ pub fn trim_to_context_window(messages: &[Message], budget_chars: usize) -> (usi // Keep at least 2 non-system messages if rest.len() <= 2 { - return (0, messages.to_vec()); + return (vec![], messages.to_vec()); } let system_chars = estimate_chars(system_msgs); @@ -133,7 +136,7 @@ pub fn trim_to_context_window(messages: &[Message], budget_chars: usize) -> (usi } if drop_count == 0 { - return (0, messages.to_vec()); + return (vec![], messages.to_vec()); } // Build the trimmed message list @@ -151,7 +154,41 @@ pub fn trim_to_context_window(messages: &[Message], budget_chars: usize) -> (usi result.push(summary_msg); result.extend_from_slice(kept); - (drop_count, result) + (dropped.to_vec(), result) +} + +/// Format dropped messages for use in compaction prompts. +/// +/// Produces a `[role]: text` line per message. Tool calls are shown as +/// `Called: name(args)`. +pub fn format_dropped_messages(messages: &[Message]) -> String { + let mut lines: Vec = Vec::new(); + for msg in messages { + let role = msg.role.to_string(); + let text = msg.text_content(); + + // Check for tool calls in metadata + if let Some(tc_val) = msg.metadata.get("tool_calls") { + if let Some(arr) = tc_val.as_array() { + for tc in arr { + let name = tc.get("name").and_then(|v| v.as_str()).unwrap_or("unknown"); + let args = tc.get("arguments").and_then(|v| v.as_str()).unwrap_or("{}"); + lines.push(format!("[{role}]: Called: {name}({args})")); + } + } + } + + if !text.is_empty() { + lines.push(format!("[{role}]: {text}")); + } else if lines.is_empty() + || !lines + .last() + .is_some_and(|l| l.starts_with(&format!("[{role}]"))) + { + lines.push(format!("[{role} message]")); + } + } + lines.join("\n") } // --------------------------------------------------------------------------- @@ -219,7 +256,7 @@ mod tests { fn test_trim_under_budget() { let msgs = vec![msg(Role::System, "sys"), msg(Role::User, "hi")]; let (dropped, result) = trim_to_context_window(&msgs, 100_000); - assert_eq!(dropped, 0); + assert!(dropped.is_empty()); assert_eq!(result.len(), 2); } @@ -234,7 +271,7 @@ mod tests { ]; // Budget that can't fit all messages let (dropped, result) = trim_to_context_window(&msgs, 500); - assert!(dropped > 0); + assert!(!dropped.is_empty()); // System message preserved assert_eq!(result[0].role, Role::System); // Summary message inserted @@ -269,7 +306,51 @@ mod tests { ]; // Even with tiny budget, keep at least 2 non-system messages let (dropped, result) = trim_to_context_window(&msgs, 10); - assert_eq!(dropped, 0); + assert!(dropped.is_empty()); assert_eq!(result.len(), 3); } + + #[test] + fn test_trim_returns_dropped_messages() { + let msgs = vec![ + msg(Role::System, "sys"), + msg(Role::User, &"A".repeat(1000)), + msg(Role::User, &"B".repeat(1000)), + msg(Role::User, &"C".repeat(100)), + msg(Role::User, &"D".repeat(100)), + ]; + let (dropped, _result) = trim_to_context_window(&msgs, 500); + assert!(!dropped.is_empty()); + // Dropped messages should be the oldest non-system ones + assert_eq!(dropped[0].role, Role::User); + assert!(dropped[0].text_content().starts_with('A')); + } + + #[test] + fn test_format_dropped_messages_basic() { + let msgs = vec![ + msg(Role::User, "Hello there"), + msg(Role::Assistant, "Hi! How can I help?"), + ]; + let formatted = format_dropped_messages(&msgs); + assert!(formatted.contains("[user]: Hello there")); + assert!(formatted.contains("[assistant]: Hi! How can I help?")); + } + + #[test] + fn test_format_dropped_messages_with_tool_calls() { + let mut m = msg(Role::Assistant, ""); + m.metadata.insert( + "tool_calls".into(), + serde_json::json!([{"name": "get_weather", "arguments": "{\"city\":\"NY\"}"}]), + ); + let formatted = format_dropped_messages(&[m]); + assert!(formatted.contains("Called: get_weather")); + assert!(formatted.contains("NY")); + } + + #[test] + fn test_format_dropped_messages_empty() { + assert_eq!(format_dropped_messages(&[]), ""); + } } diff --git a/runtime/rust/prompty/src/lib.rs b/runtime/rust/prompty/src/lib.rs index 9456d7f2..4c34fdc5 100644 --- a/runtime/rust/prompty/src/lib.rs +++ b/runtime/rust/prompty/src/lib.rs @@ -43,6 +43,7 @@ pub mod loader; pub mod model; pub mod parsers; pub mod pipeline; +pub mod prelude; pub mod registry; pub mod renderers; pub mod steering; @@ -53,7 +54,9 @@ pub mod types; // Re-export core types for convenience pub use connections::{clear_connections, has_connection, register_connection, with_connection}; -pub use context::{estimate_chars, summarize_dropped, trim_to_context_window}; +pub use context::{ + estimate_chars, format_dropped_messages, summarize_dropped, trim_to_context_window, +}; pub use guardrails::{ GuardrailError, GuardrailPhase, GuardrailResult, Guardrails, InputGuardrail, OutputGuardrail, ToolGuardrail, @@ -62,9 +65,9 @@ pub use interfaces::{ExecuteError, Executor, InvokerError, Parser, Processor, Re pub use loader::{LoadError, load, load_async, load_from_string}; pub use model::Prompty; pub use pipeline::{ - AgentEvent, AsyncToolFn, EventCallback, ToolFn, ToolHandler, TurnOptions, - invoke as invoke_agent, invoke_from_path, prepare, process, register_defaults, render, run, - turn, turn_from_path, validate_inputs, + AgentEvent, AsyncToolFn, Compaction, CompactionFn, EventCallback, ToolFn, ToolHandler, + TurnOptions, TurnOptionsBuilder, invoke as invoke_agent, invoke_from_path, prepare, process, + register_defaults, render, run, turn, turn_from_path, validate_inputs, }; pub use registry::{ clear_cache, has_executor, has_parser, has_processor, has_renderer, invoke_executor, diff --git a/runtime/rust/prompty/src/pipeline.rs b/runtime/rust/prompty/src/pipeline.rs index 8825f338..27158de9 100644 --- a/runtime/rust/prompty/src/pipeline.rs +++ b/runtime/rust/prompty/src/pipeline.rs @@ -10,7 +10,9 @@ use std::collections::HashMap; use std::future::Future; use std::path::Path; +use std::path::PathBuf; use std::pin::Pin; +use std::sync::Arc; use serde_json::{Value, json}; @@ -612,6 +614,38 @@ pub enum ToolHandler { Async(AsyncToolFn), } +// --------------------------------------------------------------------------- +// Compaction +// --------------------------------------------------------------------------- + +/// Type alias for compaction functions. +/// +/// Receives dropped messages and returns a summary string (or an error). +pub type CompactionFn = Arc< + dyn Fn( + &[Message], + ) -> Pin< + Box< + dyn Future>> + + Send, + >, + > + Send + + Sync, +>; + +/// Context compaction strategy for replacing low-signal summaries. +/// +/// When messages are trimmed by `trim_to_context_window`, the default summary +/// is a simple concatenation of truncated message texts. With compaction, the +/// summary can be replaced by an LLM-powered (Prompty file) or function-based +/// higher-quality summary. +pub enum Compaction { + /// Path to a `.prompty` file that summarizes dropped messages. + Prompty(PathBuf), + /// Custom async function that receives dropped messages and returns a summary. + Function(CompactionFn), +} + // --------------------------------------------------------------------------- // TurnOptions // --------------------------------------------------------------------------- @@ -641,6 +675,9 @@ pub struct TurnOptions { pub validator: Option Result<(), String> + Send + Sync>>, /// Maximum retries for LLM calls with exponential backoff (§9.10, default: 3). pub max_llm_retries: usize, + /// Context compaction strategy. When set and messages are trimmed, replaces + /// the default `summarize_dropped()` summary with a higher-quality one. + pub compaction: Option, } impl Default for TurnOptions { @@ -657,6 +694,7 @@ impl Default for TurnOptions { parallel_tool_calls: false, validator: None, max_llm_retries: 3, + compaction: None, } } } @@ -686,12 +724,170 @@ impl TurnOptions { .map(|c| c.load(std::sync::atomic::Ordering::Relaxed)) .unwrap_or(false) } + + /// Create a builder starting from defaults. + pub fn builder() -> TurnOptionsBuilder { + TurnOptionsBuilder { + opts: TurnOptions::default(), + } + } +} + +/// Builder for [`TurnOptions`] with fluent API. +/// +/// ```rust +/// use prompty::TurnOptions; +/// +/// let opts = TurnOptions::builder() +/// .max_iterations(5) +/// .context_budget(50_000) +/// .build(); +/// assert_eq!(opts.max_iterations, 5); +/// ``` +pub struct TurnOptionsBuilder { + opts: TurnOptions, +} + +impl TurnOptionsBuilder { + pub fn max_iterations(mut self, n: usize) -> Self { + self.opts.max_iterations = n; + self + } + + pub fn raw(mut self, raw: bool) -> Self { + self.opts.raw = raw; + self + } + + pub fn tools(mut self, tools: HashMap) -> Self { + self.opts.tools = tools; + self + } + + pub fn tool(mut self, name: impl Into, handler: ToolHandler) -> Self { + self.opts.tools.insert(name.into(), handler); + self + } + + pub fn on_event(mut self, cb: EventCallback) -> Self { + self.opts.on_event = Some(cb); + self + } + + pub fn cancelled(mut self, token: std::sync::Arc) -> Self { + self.opts.cancelled = Some(token); + self + } + + pub fn context_budget(mut self, budget: usize) -> Self { + self.opts.context_budget = Some(budget); + self + } + + pub fn guardrails(mut self, g: crate::guardrails::Guardrails) -> Self { + self.opts.guardrails = Some(g); + self + } + + pub fn steering(mut self, s: crate::steering::Steering) -> Self { + self.opts.steering = Some(s); + self + } + + pub fn parallel_tool_calls(mut self, parallel: bool) -> Self { + self.opts.parallel_tool_calls = parallel; + self + } + + #[allow(clippy::type_complexity)] + pub fn validator( + mut self, + v: Box Result<(), String> + Send + Sync>, + ) -> Self { + self.opts.validator = Some(v); + self + } + + pub fn max_llm_retries(mut self, n: usize) -> Self { + self.opts.max_llm_retries = n; + self + } + + pub fn compaction(mut self, c: Compaction) -> Self { + self.opts.compaction = Some(c); + self + } + + /// Consume the builder and return the configured [`TurnOptions`]. + pub fn build(self) -> TurnOptions { + self.opts + } } // --------------------------------------------------------------------------- // turn — conversational round-trip with optional tool calling // --------------------------------------------------------------------------- +// --------------------------------------------------------------------------- +// Compaction helpers +// --------------------------------------------------------------------------- + +/// Replace the synthetic summary message in `messages` with a compacted summary. +/// +/// Scans for the first `User` message whose text starts with `[Context summary:` +/// and replaces it. +fn replace_summary_message(messages: &mut [Message], summary: &str) { + for msg in messages.iter_mut() { + if msg.role == Role::User && msg.text_content().starts_with("[Context summary:") { + *msg = Message::text(Role::User, format!("[Context summary: {summary}]")); + return; + } + } +} + +/// Run the compaction strategy on dropped messages, replacing the default summary +/// in `messages` on success. On failure the existing summary is preserved. +pub async fn apply_compaction( + compaction: &Compaction, + dropped: &[Message], + messages: &mut [Message], + span: &crate::tracing::SpanEmitter, +) { + span.emit("compaction_start", &json!({"dropped_count": dropped.len()})); + + let result = match compaction { + Compaction::Prompty(path) => { + let text = crate::context::format_dropped_messages(dropped); + let mut inputs = serde_json::Map::new(); + inputs.insert("messages".into(), serde_json::Value::String(text)); + match crate::invoke_from_path(path, Some(&serde_json::Value::Object(inputs))).await { + Ok(val) => Ok(val.as_str().unwrap_or("").to_string()), + Err(e) => Err(format!("{e}")), + } + } + Compaction::Function(f) => match f(dropped).await { + Ok(s) => Ok(s), + Err(e) => Err(format!("{e}")), + }, + }; + + match result { + Ok(summary) if !summary.trim().is_empty() => { + replace_summary_message(messages, &summary); + span.emit( + "compaction_complete", + &json!({"summary_length": summary.len()}), + ); + } + Ok(_) => { + span.emit("compaction_failed", &json!({"reason": "empty result"})); + } + Err(reason) => { + span.emit("compaction_failed", &json!({"reason": reason})); + } + } +} + /// One conversational round-trip: prepare → [agent loop with tool calls] → process. /// /// Without tools, this is equivalent to `invoke` (simple mode). With tools, @@ -731,10 +927,16 @@ pub async fn turn( // Context trimming if let Some(budget) = opts.context_budget { let (dropped, trimmed) = crate::context::trim_to_context_window(&messages, budget); - if dropped > 0 { - span.emit("context_trimmed", &json!(dropped)); + if !dropped.is_empty() { + span.emit("context_trimmed", &json!(dropped.len())); + messages = trimmed; + // Apply compaction if configured + if let Some(ref compaction) = opts.compaction { + apply_compaction(compaction, &dropped, &mut messages, &span).await; + } + } else { + messages = trimmed; } - messages = trimmed; } // Input guardrail @@ -879,9 +1081,13 @@ pub async fn turn( // Context trimming if let Some(budget) = opts.context_budget { let (dropped, trimmed) = crate::context::trim_to_context_window(&messages, budget); - if dropped > 0 { - iter_span.emit("context_trimmed", &json!(dropped)); + if !dropped.is_empty() { + iter_span.emit("context_trimmed", &json!(dropped.len())); messages = trimmed; + // Apply compaction if configured + if let Some(ref compaction) = opts.compaction { + apply_compaction(compaction, &dropped, &mut messages, &iter_span).await; + } opts.emit(AgentEvent::MessagesUpdated { messages: messages.clone(), }); @@ -2451,4 +2657,79 @@ mod tests { *captured ); } + + // ----------------------------------------------------------------------- + // TurnOptionsBuilder tests + // ----------------------------------------------------------------------- + + #[test] + fn test_builder_defaults() { + let opts = TurnOptions::builder().build(); + assert_eq!(opts.max_iterations, 10); + assert_eq!(opts.max_llm_retries, 3); + assert!(!opts.raw); + assert!(!opts.parallel_tool_calls); + assert!(opts.context_budget.is_none()); + assert!(opts.compaction.is_none()); + assert!(opts.on_event.is_none()); + assert!(opts.cancelled.is_none()); + assert!(opts.guardrails.is_none()); + assert!(opts.steering.is_none()); + assert!(opts.validator.is_none()); + assert!(opts.tools.is_empty()); + } + + #[test] + fn test_builder_chaining() { + let opts = TurnOptions::builder() + .max_iterations(5) + .context_budget(50_000) + .max_llm_retries(5) + .parallel_tool_calls(true) + .raw(true) + .build(); + assert_eq!(opts.max_iterations, 5); + assert_eq!(opts.context_budget, Some(50_000)); + assert_eq!(opts.max_llm_retries, 5); + assert!(opts.parallel_tool_calls); + assert!(opts.raw); + } + + #[test] + fn test_builder_tool_method() { + let handler = ToolHandler::Sync(Box::new(|_args| Ok("result".to_string()))); + let opts = TurnOptions::builder().tool("my_tool", handler).build(); + assert!(opts.tools.contains_key("my_tool")); + assert_eq!(opts.tools.len(), 1); + } + + #[test] + fn test_builder_multiple_tools() { + let h1 = ToolHandler::Sync(Box::new(|_| Ok("a".to_string()))); + let h2 = ToolHandler::Sync(Box::new(|_| Ok("b".to_string()))); + let opts = TurnOptions::builder() + .tool("tool_a", h1) + .tool("tool_b", h2) + .build(); + assert_eq!(opts.tools.len(), 2); + assert!(opts.tools.contains_key("tool_a")); + assert!(opts.tools.contains_key("tool_b")); + } + + #[test] + fn test_builder_compaction() { + let opts = TurnOptions::builder() + .compaction(Compaction::Prompty("summarize.prompty".into())) + .build(); + assert!(opts.compaction.is_some()); + } + + #[test] + fn test_builder_cancelled_token() { + let token = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let opts = TurnOptions::builder().cancelled(token.clone()).build(); + assert!(!opts.is_cancelled()); + token.store(true, std::sync::atomic::Ordering::Relaxed); + assert!(opts.is_cancelled()); + } } diff --git a/runtime/rust/prompty/src/prelude.rs b/runtime/rust/prompty/src/prelude.rs new file mode 100644 index 00000000..4dd46c35 --- /dev/null +++ b/runtime/rust/prompty/src/prelude.rs @@ -0,0 +1,20 @@ +//! Convenience re-exports for common Prompty types. +//! +//! ```rust +//! use prompty::prelude::*; +//! ``` + +pub use crate::connections::{register_connection, with_connection}; +pub use crate::interfaces::{ExecuteError, Executor, Processor}; +pub use crate::loader::load; +pub use crate::model::Prompty; +pub use crate::pipeline::{ + AgentEvent, Compaction, EventCallback, ToolHandler, TurnOptions, TurnOptionsBuilder, + invoke as invoke_agent, prepare, run, turn, +}; +pub use crate::registry::{register_executor, register_processor}; +pub use crate::tool_dispatch::{register_tool, register_tool_handler}; +pub use crate::tracing::{Tracer, console_tracer, trace, trace_async}; +pub use crate::types::{ + ContentPart, Message, PromptyStream, Role, StreamChunk, TextPart, ToolCall, +}; diff --git a/runtime/rust/prompty/src/tracing/mod.rs b/runtime/rust/prompty/src/tracing/mod.rs index 1dfcbfea..bf21a301 100644 --- a/runtime/rust/prompty/src/tracing/mod.rs +++ b/runtime/rust/prompty/src/tracing/mod.rs @@ -14,4 +14,6 @@ pub use console::console_tracer; #[cfg(feature = "otel")] pub use otel::{init_otel_stdout, otel_tracer}; pub use prompty_tracer::PromptyTracer; -pub use tracer::{Tracer, sanitize_value, trace, trace_async, trace_span, trace_span_async}; +pub use tracer::{ + SpanEmitter, Tracer, sanitize_value, trace, trace_async, trace_span, trace_span_async, +}; diff --git a/runtime/rust/prompty/tests/compaction_tests.rs b/runtime/rust/prompty/tests/compaction_tests.rs new file mode 100644 index 00000000..e4646796 --- /dev/null +++ b/runtime/rust/prompty/tests/compaction_tests.rs @@ -0,0 +1,256 @@ +//! Tests for context compaction feature. + +use std::sync::Arc; + +use prompty::context::{format_dropped_messages, trim_to_context_window}; +use prompty::types::{Message, Role}; +use prompty::{Compaction, CompactionFn, TurnOptions}; + +// --------------------------------------------------------------------------- +// format_dropped_messages +// --------------------------------------------------------------------------- + +#[test] +fn test_format_dropped_messages_basic() { + let msgs = vec![ + Message::text(Role::User, "Hello there"), + Message::text(Role::Assistant, "Hi! How can I help?"), + ]; + let formatted = format_dropped_messages(&msgs); + assert!(formatted.contains("[user]: Hello there")); + assert!(formatted.contains("[assistant]: Hi! How can I help?")); +} + +#[test] +fn test_format_dropped_messages_with_tool_calls() { + let mut m = Message::text(Role::Assistant, ""); + m.metadata.insert( + "tool_calls".into(), + serde_json::json!([{"name": "get_weather", "arguments": "{\"city\":\"NY\"}"}]), + ); + let formatted = format_dropped_messages(&[m]); + assert!( + formatted.contains("Called: get_weather"), + "should contain tool call name" + ); + assert!( + formatted.contains("NY"), + "should contain tool call arguments" + ); +} + +#[test] +fn test_format_dropped_messages_empty() { + assert_eq!(format_dropped_messages(&[]), ""); +} + +#[test] +fn test_format_dropped_messages_multiple_tool_calls() { + let mut m = Message::text(Role::Assistant, ""); + m.metadata.insert( + "tool_calls".into(), + serde_json::json!([ + {"name": "get_weather", "arguments": "{\"city\":\"NY\"}"}, + {"name": "search", "arguments": "{\"q\":\"hello\"}"}, + ]), + ); + let formatted = format_dropped_messages(&[m]); + assert!(formatted.contains("Called: get_weather")); + assert!(formatted.contains("Called: search")); +} + +// --------------------------------------------------------------------------- +// trim_to_context_window returns dropped messages +// --------------------------------------------------------------------------- + +#[test] +fn test_trim_returns_dropped_messages() { + let msgs = vec![ + Message::text(Role::System, "sys"), + Message::text(Role::User, &"A".repeat(1000)), + Message::text(Role::User, &"B".repeat(1000)), + Message::text(Role::User, &"C".repeat(100)), + Message::text(Role::User, &"D".repeat(100)), + ]; + let (dropped, trimmed) = trim_to_context_window(&msgs, 500); + assert!(!dropped.is_empty(), "should have dropped messages"); + // Dropped messages are the oldest non-system ones + assert_eq!(dropped[0].role, Role::User); + assert!(dropped[0].text_content().starts_with('A')); + // Trimmed should contain summary message + assert!(trimmed[1].text_content().contains("Context summary")); +} + +#[test] +fn test_trim_under_budget_returns_empty_dropped() { + let msgs = vec![ + Message::text(Role::System, "sys"), + Message::text(Role::User, "hi"), + ]; + let (dropped, result) = trim_to_context_window(&msgs, 100_000); + assert!(dropped.is_empty()); + assert_eq!(result.len(), 2); +} + +// --------------------------------------------------------------------------- +// TurnOptions default +// --------------------------------------------------------------------------- + +#[test] +fn test_compaction_none_default() { + let opts = TurnOptions::default(); + assert!(opts.compaction.is_none()); +} + +// --------------------------------------------------------------------------- +// Compaction::Function replaces summary +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_compaction_function_replaces_summary() { + // Build messages that will exceed budget and be trimmed + let mut messages = vec![ + Message::text(Role::System, "sys"), + // The summary message (simulating what trim_to_context_window produces) + Message::text( + Role::User, + "[Context summary: [user]: old content\n... (1 messages omitted)]", + ), + Message::text(Role::User, "recent message"), + ]; + + let compaction_fn: CompactionFn = Arc::new(|_dropped: &[Message]| { + Box::pin(async { Ok("LLM-powered summary of the conversation".to_string()) }) + }); + + let dropped = vec![Message::text(Role::User, "old content that was dropped")]; + + let span = prompty::tracing::Tracer::start("test"); + prompty::pipeline::apply_compaction( + &Compaction::Function(compaction_fn), + &dropped, + &mut messages, + &span, + ) + .await; + span.end(); + + // The summary message should be replaced + let summary_msg = &messages[1]; + assert_eq!(summary_msg.role, Role::User); + assert!( + summary_msg + .text_content() + .contains("LLM-powered summary of the conversation"), + "summary should be replaced: got {}", + summary_msg.text_content() + ); +} + +// --------------------------------------------------------------------------- +// Compaction failure preserves default summary +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_compaction_failure_preserves_default() { + let original_summary = "[Context summary: [user]: old stuff\n... (1 messages omitted)]"; + let mut messages = vec![ + Message::text(Role::System, "sys"), + Message::text(Role::User, original_summary), + Message::text(Role::User, "recent"), + ]; + + let compaction_fn: CompactionFn = Arc::new(|_dropped: &[Message]| { + Box::pin(async { Err("model unavailable".to_string().into()) }) + }); + + let dropped = vec![Message::text(Role::User, "dropped content")]; + + let span = prompty::tracing::Tracer::start("test"); + prompty::pipeline::apply_compaction( + &Compaction::Function(compaction_fn), + &dropped, + &mut messages, + &span, + ) + .await; + span.end(); + + // Original summary should be preserved + assert_eq!(messages[1].text_content(), original_summary); +} + +// --------------------------------------------------------------------------- +// Compaction returning empty string preserves default +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_compaction_empty_result_preserves_default() { + let original_summary = "[Context summary: [user]: stuff\n... (1 messages omitted)]"; + let mut messages = vec![ + Message::text(Role::System, "sys"), + Message::text(Role::User, original_summary), + Message::text(Role::User, "recent"), + ]; + + let compaction_fn: CompactionFn = Arc::new(|_dropped: &[Message]| { + Box::pin(async { Ok(" ".to_string()) }) // whitespace-only + }); + + let dropped = vec![Message::text(Role::User, "dropped")]; + + let span = prompty::tracing::Tracer::start("test"); + prompty::pipeline::apply_compaction( + &Compaction::Function(compaction_fn), + &dropped, + &mut messages, + &span, + ) + .await; + span.end(); + + assert_eq!(messages[1].text_content(), original_summary); +} + +// --------------------------------------------------------------------------- +// Compaction function receives the dropped messages +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_compaction_receives_dropped_messages() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let count = Arc::new(AtomicUsize::new(0)); + let count_clone = count.clone(); + + let compaction_fn: CompactionFn = Arc::new(move |dropped: &[Message]| { + count_clone.store(dropped.len(), Ordering::SeqCst); + Box::pin(async { Ok("summary".to_string()) }) + }); + + let mut messages = vec![ + Message::text(Role::System, "sys"), + Message::text( + Role::User, + "[Context summary: old\n... (2 messages omitted)]", + ), + Message::text(Role::User, "recent"), + ]; + + let dropped = vec![ + Message::text(Role::User, "msg1"), + Message::text(Role::User, "msg2"), + ]; + + let span = prompty::tracing::Tracer::start("test"); + prompty::pipeline::apply_compaction( + &Compaction::Function(compaction_fn), + &dropped, + &mut messages, + &span, + ) + .await; + span.end(); + + assert_eq!(count.load(Ordering::SeqCst), 2); +} diff --git a/runtime/typescript/packages/core/src/core/agent-events.ts b/runtime/typescript/packages/core/src/core/agent-events.ts index 324d8e0f..9e57170a 100644 --- a/runtime/typescript/packages/core/src/core/agent-events.ts +++ b/runtime/typescript/packages/core/src/core/agent-events.ts @@ -13,7 +13,10 @@ export type AgentEventType = | "messages_updated" | "done" | "error" - | "cancelled"; + | "cancelled" + | "compaction_start" + | "compaction_complete" + | "compaction_failed"; /** Callback signature for agent loop events. */ export type EventCallback = (eventType: AgentEventType, data: Record) => void; diff --git a/runtime/typescript/packages/core/src/core/context.ts b/runtime/typescript/packages/core/src/core/context.ts index c14fba25..ad5dbc85 100644 --- a/runtime/typescript/packages/core/src/core/context.ts +++ b/runtime/typescript/packages/core/src/core/context.ts @@ -67,6 +67,31 @@ export function summarizeDropped(messages: Message[]): string { return result.trimEnd() + "]"; } +/** + * Format dropped messages as a human-readable text block. + * Each message is rendered as `[role]: content`, with tool calls shown as + * `Called: name(args)`. + */ +export function formatDroppedMessages(messages: Message[]): string { + const lines: string[] = []; + for (const msg of messages) { + const msgText = msg.text.trim(); + if (msgText) { + lines.push(`[${msg.role}]: ${msgText}`); + } + const toolCalls = msg.metadata?.tool_calls; + if (Array.isArray(toolCalls)) { + for (const tc of toolCalls) { + const tcObj = tc as Record; + const name = (tcObj.name as string) ?? ((tcObj.function as Record)?.name ?? "?"); + const args = (tcObj.arguments as string) ?? ((tcObj.function as Record)?.arguments ?? ""); + lines.push(`Called: ${name}(${args})`); + } + } + } + return lines.join("\n"); +} + /** * Trim messages in-place to fit within a character budget. * Returns [droppedCount, droppedMessages]. diff --git a/runtime/typescript/packages/core/src/core/index.ts b/runtime/typescript/packages/core/src/core/index.ts index 3230d31e..fbda71f6 100644 --- a/runtime/typescript/packages/core/src/core/index.ts +++ b/runtime/typescript/packages/core/src/core/index.ts @@ -32,7 +32,7 @@ export { } from "./tool-dispatch.js"; export { type AgentEventType, type EventCallback, emitEvent } from "./agent-events.js"; export { CancelledError, checkCancellation } from "./cancellation.js"; -export { estimateChars, summarizeDropped, trimToContextWindow } from "./context.js"; +export { estimateChars, summarizeDropped, trimToContextWindow, formatDroppedMessages } from "./context.js"; export { type GuardrailResult, GuardrailError, diff --git a/runtime/typescript/packages/core/src/core/pipeline.ts b/runtime/typescript/packages/core/src/core/pipeline.ts index 05d9f2ba..a7505908 100644 --- a/runtime/typescript/packages/core/src/core/pipeline.ts +++ b/runtime/typescript/packages/core/src/core/pipeline.ts @@ -46,7 +46,7 @@ import { load } from "./loader.js"; import { dispatchTool, resilientJsonParse } from "./tool-dispatch.js"; import { type EventCallback, emitEvent } from "./agent-events.js"; import { CancelledError, checkCancellation } from "./cancellation.js"; -import { trimToContextWindow } from "./context.js"; +import { trimToContextWindow, formatDroppedMessages } from "./context.js"; import { GuardrailError, Guardrails } from "./guardrails.js"; import { Steering } from "./steering.js"; import { cast } from "./structured.js"; @@ -548,6 +548,63 @@ export interface TurnOptions { parallelToolCalls?: boolean; /** Maximum LLM call retries on transient failure in agent loop (§9.10, default: 3). */ maxLlmRetries?: number; + /** Context compaction: path to .prompty file or function returning a summary of dropped messages. */ + compaction?: string | ((messages: Message[]) => string | Promise); +} + +// --------------------------------------------------------------------------- +// Context Compaction +// --------------------------------------------------------------------------- + +/** + * Replace the default "[Context summary: …]" message with a compaction-produced summary. + * Finds the summary message by looking for the "[Context summary:" marker and replaces it. + */ +function replaceSummaryMessage(messages: Message[], summary: string): void { + const idx = messages.findIndex( + (m) => + m.role === "user" && + m.parts.some( + (p) => p.kind === "text" && (p as { value: string }).value.includes("[Context summary:"), + ), + ); + if (idx >= 0) { + messages[idx] = new Message("user", [text(`[Context summary: ${summary}]`)]); + } +} + +/** + * Apply context compaction to replace the default summary with a higher-quality one. + * On failure, the default summary already in `messages` is preserved. + */ +async function applyCompaction( + compaction: string | ((messages: Message[]) => string | Promise), + dropped: Message[], + messages: Message[], + onEvent: EventCallback | undefined, +): Promise { + emitEvent(onEvent, "compaction_start", { dropped_count: dropped.length }); + try { + let summary: string; + if (typeof compaction === "string") { + const formatted = formatDroppedMessages(dropped); + const result = await invoke(compaction, { messages: formatted }); + summary = typeof result === "string" ? result : String(result); + } else { + const result = compaction(dropped); + summary = + typeof result === "string" ? result : String(await (result as Promise)); + } + + if (summary && summary.trim()) { + replaceSummaryMessage(messages, summary); + emitEvent(onEvent, "compaction_complete", { summary_length: summary.length }); + } else { + emitEvent(onEvent, "compaction_failed", { reason: "empty result" }); + } + } catch (err) { + emitEvent(onEvent, "compaction_failed", { reason: String(err) }); + } } /** @@ -626,9 +683,12 @@ export async function turn( // §13.3 — Trim context window if (options?.contextBudget !== undefined) { - const [droppedCount] = trimToContextWindow(messages, options.contextBudget); + const [droppedCount, droppedMsgs] = trimToContextWindow(messages, options.contextBudget); if (droppedCount > 0) { emitEvent(onEvent, "messages_updated", { messages }); + if (options.compaction != null) { + await applyCompaction(options.compaction, droppedMsgs, messages, onEvent); + } } } @@ -715,10 +775,13 @@ export async function turn( // §13.3 — Trim context window if (contextBudget !== undefined) { - const [droppedCount] = trimToContextWindow(messages, contextBudget); + const [droppedCount, droppedMsgs] = trimToContextWindow(messages, contextBudget); if (droppedCount > 0) { emitEvent(onEvent, "messages_updated", { messages }); emitEvent(onEvent, "status", { message: `Trimmed ${droppedCount} messages for context budget` }); + if (options?.compaction != null) { + await applyCompaction(options.compaction, droppedMsgs, messages, onEvent); + } } } diff --git a/runtime/typescript/packages/core/src/index.ts b/runtime/typescript/packages/core/src/index.ts index 4ff0aa23..4e453b9f 100644 --- a/runtime/typescript/packages/core/src/index.ts +++ b/runtime/typescript/packages/core/src/index.ts @@ -80,6 +80,7 @@ export { estimateChars, summarizeDropped, trimToContextWindow, + formatDroppedMessages, type GuardrailResult, GuardrailError, type InputGuardrail, diff --git a/runtime/typescript/packages/core/tests/compaction.test.ts b/runtime/typescript/packages/core/tests/compaction.test.ts new file mode 100644 index 00000000..04db3742 --- /dev/null +++ b/runtime/typescript/packages/core/tests/compaction.test.ts @@ -0,0 +1,335 @@ +import { describe, it, expect, beforeEach, vi } from "vitest"; +import { Message, text } from "../src/core/types.js"; +import { formatDroppedMessages, trimToContextWindow } from "../src/core/context.js"; +import { turn } from "../src/core/pipeline.js"; +import { emitEvent, type AgentEventType, type EventCallback } from "../src/core/agent-events.js"; +import { + registerRenderer, + registerParser, + registerExecutor, + registerProcessor, +} from "../src/core/registry.js"; +import type { Renderer, Parser, Executor, Processor } from "../src/core/interfaces.js"; +import { Prompty } from "../src/model/prompty.js"; + +// =========================================================================== +// formatDroppedMessages +// =========================================================================== + +describe("formatDroppedMessages", () => { + it("formats user and assistant messages as readable text", () => { + const msgs = [ + new Message("user", [text("What is AI?")]), + new Message("assistant", [text("AI is artificial intelligence.")]), + ]; + const result = formatDroppedMessages(msgs); + expect(result).toContain("[user]: What is AI?"); + expect(result).toContain("[assistant]: AI is artificial intelligence."); + }); + + it("includes tool calls as Called: name(args)", () => { + const msgs = [ + new Message("assistant", [text("Let me check.")], { + tool_calls: [ + { name: "get_weather", arguments: '{"city":"Seattle"}' }, + ], + }), + ]; + const result = formatDroppedMessages(msgs); + expect(result).toContain("[assistant]: Let me check."); + expect(result).toContain('Called: get_weather({"city":"Seattle"})'); + }); + + it("handles function-style tool calls", () => { + const msgs = [ + new Message("assistant", [text("Checking...")], { + tool_calls: [ + { function: { name: "lookup", arguments: '{"q":"test"}' } }, + ], + }), + ]; + const result = formatDroppedMessages(msgs); + expect(result).toContain('Called: lookup({"q":"test"})'); + }); + + it("returns empty string for empty array", () => { + expect(formatDroppedMessages([])).toBe(""); + }); + + it("skips messages with empty text but includes tool calls", () => { + const msgs = [ + new Message("assistant", [], { + tool_calls: [{ name: "fn", arguments: "" }], + }), + ]; + const result = formatDroppedMessages(msgs); + expect(result).toBe("Called: fn()"); + }); +}); + +// =========================================================================== +// Context Compaction in turn() +// =========================================================================== + +// Stubs for turn() integration tests +class StubRenderer implements Renderer { + render(_agent: Prompty, _inputs: Record): string { + return "system:\nYou are helpful.\n\nuser:\nHello"; + } +} + +class StubParser implements Parser { + parse(_agent: Prompty, rendered: string): Message[] { + // Simple two-message output + return [ + new Message("system", [text("You are helpful.")]), + new Message("user", [text("Hello")]), + ]; + } +} + +class StubProcessor implements Processor { + process(_agent: Prompty, response: unknown): unknown { + const resp = response as { choices: { message: { content: string } }[] }; + return resp.choices?.[0]?.message?.content ?? response; + } +} + +function makeStubExecutor(fn: (agent: Prompty, messages: Message[]) => Promise): Executor { + return { execute: fn }; +} + +function makeFinalResponse(content: string) { + return { + choices: [{ + finish_reason: "stop", + message: { role: "assistant", content, tool_calls: null }, + }], + }; +} + +function makeTestAgent(): Prompty { + const agent = new Prompty({ + name: "compact-test", + instructions: "Hello {{name}}", + }); + agent.template = { format: { kind: "compact-stub" }, parser: { kind: "compact-stub" } } as any; + (agent as any).model = { provider: "compact-stub" }; + return agent; +} + +describe("Context Compaction", () => { + let capturedMessages: Message[]; + + beforeEach(() => { + capturedMessages = []; + + registerRenderer("compact-stub", new StubRenderer()); + registerParser("compact-stub", new StubParser()); + registerProcessor("compact-stub", new StubProcessor()); + registerExecutor( + "compact-stub", + makeStubExecutor(async (_agent, messages) => { + capturedMessages = [...messages]; + return makeFinalResponse("The answer is 42"); + }), + ); + }); + + it("function compaction replaces default summary", async () => { + // Build a parser that returns many messages to exceed the budget + const longMessages = [ + new Message("system", [text("System prompt")]), + new Message("user", [text("q1 " + "x".repeat(500))]), + new Message("assistant", [text("a1 " + "x".repeat(500))]), + new Message("user", [text("q2 " + "x".repeat(500))]), + new Message("assistant", [text("a2 " + "x".repeat(500))]), + new Message("user", [text("Final question")]), + ]; + registerParser("compact-stub", { + parse: () => [...longMessages], + }); + + const compactionFn = vi.fn((_dropped: Message[]) => "LLM-quality summary of prior conversation"); + + const agent = makeTestAgent(); + await turn(agent, { name: "Alice" }, { + contextBudget: 300, + compaction: compactionFn, + }); + + expect(compactionFn).toHaveBeenCalled(); + // The summary message in capturedMessages should contain our custom summary + const summaryMsg = capturedMessages.find( + (m) => m.role === "user" && m.text.includes("[Context summary:"), + ); + expect(summaryMsg).toBeDefined(); + expect(summaryMsg!.text).toContain("LLM-quality summary of prior conversation"); + }); + + it("async function compaction works", async () => { + const longMessages = [ + new Message("system", [text("System prompt")]), + new Message("user", [text("q1 " + "x".repeat(500))]), + new Message("assistant", [text("a1 " + "x".repeat(500))]), + new Message("user", [text("q2 " + "x".repeat(500))]), + new Message("assistant", [text("a2 " + "x".repeat(500))]), + new Message("user", [text("Final question")]), + ]; + registerParser("compact-stub", { + parse: () => [...longMessages], + }); + + const compactionFn = vi.fn(async (_dropped: Message[]) => { + return "Async compacted summary"; + }); + + const agent = makeTestAgent(); + await turn(agent, { name: "Alice" }, { + contextBudget: 300, + compaction: compactionFn, + }); + + expect(compactionFn).toHaveBeenCalled(); + const summaryMsg = capturedMessages.find( + (m) => m.role === "user" && m.text.includes("[Context summary:"), + ); + expect(summaryMsg).toBeDefined(); + expect(summaryMsg!.text).toContain("Async compacted summary"); + }); + + it("compaction failure preserves default summary", async () => { + const longMessages = [ + new Message("system", [text("System prompt")]), + new Message("user", [text("q1 " + "x".repeat(500))]), + new Message("assistant", [text("a1 " + "x".repeat(500))]), + new Message("user", [text("q2 " + "x".repeat(500))]), + new Message("assistant", [text("a2 " + "x".repeat(500))]), + new Message("user", [text("Final question")]), + ]; + registerParser("compact-stub", { + parse: () => [...longMessages], + }); + + const events: { type: AgentEventType; data: Record }[] = []; + const onEvent: EventCallback = (t, d) => events.push({ type: t, data: d }); + + const compactionFn = vi.fn(() => { + throw new Error("Compaction service down"); + }); + + const agent = makeTestAgent(); + await turn(agent, { name: "Alice" }, { + contextBudget: 300, + compaction: compactionFn, + onEvent, + }); + + // Default summary should still be present (not replaced) + const summaryMsg = capturedMessages.find( + (m) => m.role === "user" && m.text.includes("[Context summary:"), + ); + expect(summaryMsg).toBeDefined(); + // Should NOT contain our custom text since it failed + expect(summaryMsg!.text).not.toContain("Compaction service down"); + + // compaction_failed event should have been emitted + const failedEvent = events.find((e) => e.type === "compaction_failed"); + expect(failedEvent).toBeDefined(); + expect(failedEvent!.data.reason).toContain("Compaction service down"); + }); + + it("compaction events are emitted on success", async () => { + const longMessages = [ + new Message("system", [text("System prompt")]), + new Message("user", [text("q1 " + "x".repeat(500))]), + new Message("assistant", [text("a1 " + "x".repeat(500))]), + new Message("user", [text("q2 " + "x".repeat(500))]), + new Message("assistant", [text("a2 " + "x".repeat(500))]), + new Message("user", [text("Final question")]), + ]; + registerParser("compact-stub", { + parse: () => [...longMessages], + }); + + const events: { type: AgentEventType; data: Record }[] = []; + const onEvent: EventCallback = (t, d) => events.push({ type: t, data: d }); + + const agent = makeTestAgent(); + await turn(agent, { name: "Alice" }, { + contextBudget: 300, + compaction: () => "Compact summary here", + onEvent, + }); + + const startEvent = events.find((e) => e.type === "compaction_start"); + expect(startEvent).toBeDefined(); + expect(startEvent!.data.dropped_count).toBeGreaterThan(0); + + const completeEvent = events.find((e) => e.type === "compaction_complete"); + expect(completeEvent).toBeDefined(); + expect(completeEvent!.data.summary_length).toBe("Compact summary here".length); + }); + + it("no compaction when compaction option is undefined", async () => { + const longMessages = [ + new Message("system", [text("System prompt")]), + new Message("user", [text("q1 " + "x".repeat(500))]), + new Message("assistant", [text("a1 " + "x".repeat(500))]), + new Message("user", [text("q2 " + "x".repeat(500))]), + new Message("assistant", [text("a2 " + "x".repeat(500))]), + new Message("user", [text("Final question")]), + ]; + registerParser("compact-stub", { + parse: () => [...longMessages], + }); + + const events: { type: AgentEventType; data: Record }[] = []; + const onEvent: EventCallback = (t, d) => events.push({ type: t, data: d }); + + const agent = makeTestAgent(); + await turn(agent, { name: "Alice" }, { + contextBudget: 300, + onEvent, + // no compaction option + }); + + // Default summarizeDropped summary should be present + const summaryMsg = capturedMessages.find( + (m) => m.role === "user" && m.text.includes("[Context summary:"), + ); + expect(summaryMsg).toBeDefined(); + + // No compaction events + expect(events.find((e) => e.type === "compaction_start")).toBeUndefined(); + expect(events.find((e) => e.type === "compaction_complete")).toBeUndefined(); + }); + + it("empty compaction result emits compaction_failed", async () => { + const longMessages = [ + new Message("system", [text("System prompt")]), + new Message("user", [text("q1 " + "x".repeat(500))]), + new Message("assistant", [text("a1 " + "x".repeat(500))]), + new Message("user", [text("q2 " + "x".repeat(500))]), + new Message("assistant", [text("a2 " + "x".repeat(500))]), + new Message("user", [text("Final question")]), + ]; + registerParser("compact-stub", { + parse: () => [...longMessages], + }); + + const events: { type: AgentEventType; data: Record }[] = []; + const onEvent: EventCallback = (t, d) => events.push({ type: t, data: d }); + + const agent = makeTestAgent(); + await turn(agent, { name: "Alice" }, { + contextBudget: 300, + compaction: () => " ", + onEvent, + }); + + const failedEvent = events.find((e) => e.type === "compaction_failed"); + expect(failedEvent).toBeDefined(); + expect(failedEvent!.data.reason).toBe("empty result"); + }); +}); diff --git a/web/src/content/docs/core-concepts/agent-extensions.mdx b/web/src/content/docs/core-concepts/agent-extensions.mdx index f20480b3..e37affcc 100644 --- a/web/src/content/docs/core-concepts/agent-extensions.mdx +++ b/web/src/content/docs/core-concepts/agent-extensions.mdx @@ -300,6 +300,142 @@ let result = prompty::turn(&agent, Some(&inputs), Some(options)).await?; model's response. +### Context Compaction + +The default summary produced by trimming is a simple concatenation of message +previews — functional but low-signal. **Context compaction** replaces it with a +higher-quality summary using either a `.prompty` file or a custom function. + + + +```python +# Option 1: Use a .prompty file for compaction +result = turn( + agent, + inputs={"question": "Continue our discussion"}, + tools=tools, + context_budget=50_000, + compaction="summarize.prompty", +) + +# Option 2: Use a custom function +def my_compactor(dropped_messages): + """Summarize dropped messages however you like.""" + return f"Previously discussed: {len(dropped_messages)} exchanges" + +result = turn( + agent, + inputs={"question": "Continue our discussion"}, + tools=tools, + context_budget=50_000, + compaction=my_compactor, +) +``` + + +```typescript +// Option 1: Use a .prompty file for compaction +const result = await turn(agent, inputs, { + tools, + contextBudget: 50_000, + compaction: "summarize.prompty", +}); + +// Option 2: Use a custom function +const result = await turn(agent, inputs, { + tools, + contextBudget: 50_000, + compaction: (dropped) => + `Previously discussed: ${dropped.length} exchanges`, +}); +``` + + +```csharp +// Option 1: Use a .prompty file for compaction +var result = await Pipeline.TurnAsync( + agent, inputs, tools, + contextBudget: 50_000, + compaction: "summarize.prompty" +); + +// Option 2: Use a custom function +var result = await Pipeline.TurnAsync( + agent, inputs, tools, + contextBudget: 50_000, + compaction: async (dropped) => + $"Previously discussed: {dropped.Count} exchanges" +); +``` + + +```rust +use prompty::{TurnOptions, Compaction}; + +// Option 1: Use a .prompty file for compaction +let options = TurnOptions { + context_budget: Some(50_000), + compaction: Some(Compaction::Prompty("summarize.prompty".into())), + ..Default::default() +}; + +// Option 2: Use a custom function +let options = TurnOptions { + context_budget: Some(50_000), + compaction: Some(Compaction::Function(Arc::new(|dropped| { + Box::pin(async move { + Ok(format!("Previously discussed: {} exchanges", dropped.len())) + }) + }))), + ..Default::default() +}; +``` + + + +#### The Compaction `.prompty` File + +A compaction prompt receives the dropped messages as formatted text and should +return a concise summary preserving key facts, decisions, and context: + +```yaml +--- +name: context-compactor +model: + id: gpt-4o-mini + provider: openai +inputSchema: + properties: + messages: + kind: string + description: Formatted text of dropped conversation messages +template: + format: + kind: jinja2 + parser: + kind: prompty +--- +system: +Summarize these conversation messages into a concise summary. +Preserve key facts, decisions, tool results, and context needed +to continue the conversation coherently. Be brief but complete. + +user: +{{messages}} +``` + + + + + --- ## Guardrails