From 7f74c40cc37d5f69832b57eaa4c7912ce7e6120b Mon Sep 17 00:00:00 2001 From: isaacbmiller Date: Fri, 13 Mar 2026 15:50:11 -0500 Subject: [PATCH] feat: add raw history data model --- dspy/__init__.py | 2 +- dspy/adapters/__init__.py | 3 +- dspy/adapters/types/__init__.py | 4 +- dspy/adapters/types/history.py | 151 ++++++++++++++++++++++++++++++- tests/primitives/test_example.py | 22 +++++ 5 files changed, 177 insertions(+), 5 deletions(-) diff --git a/dspy/__init__.py b/dspy/__init__.py index cdeb4b777c..f1552fced7 100644 --- a/dspy/__init__.py +++ b/dspy/__init__.py @@ -9,7 +9,7 @@ from dspy.evaluate import Evaluate # isort: skip from dspy.clients import * # isort: skip -from dspy.adapters import Adapter, ChatAdapter, JSONAdapter, XMLAdapter, TwoStepAdapter, Image, Audio, File, History, Type, Tool, ToolCalls, Code, Reasoning # isort: skip +from dspy.adapters import Adapter, ChatAdapter, JSONAdapter, XMLAdapter, TwoStepAdapter, Image, Audio, File, History, HistoryCompaction, Type, Tool, ToolCalls, Code, Reasoning # isort: skip from dspy.utils.logging_utils import configure_dspy_loggers, disable_logging, enable_logging from dspy.utils.asyncify import asyncify from dspy.utils.syncify import syncify diff --git a/dspy/adapters/__init__.py b/dspy/adapters/__init__.py index c217d7260e..7e52706291 100644 --- a/dspy/adapters/__init__.py +++ b/dspy/adapters/__init__.py @@ -2,7 +2,7 @@ from dspy.adapters.chat_adapter import ChatAdapter from dspy.adapters.json_adapter import JSONAdapter from dspy.adapters.two_step_adapter import TwoStepAdapter -from dspy.adapters.types import Audio, Code, File, History, Image, Reasoning, Tool, ToolCalls, Type +from dspy.adapters.types import Audio, Code, File, History, HistoryCompaction, Image, Reasoning, Tool, ToolCalls, Type from dspy.adapters.xml_adapter import XMLAdapter __all__ = [ @@ -10,6 +10,7 @@ "ChatAdapter", "Type", "History", + "HistoryCompaction", "Image", "Audio", "File", diff --git a/dspy/adapters/types/__init__.py b/dspy/adapters/types/__init__.py index 5ec8043021..1ecc8181ca 100644 --- a/dspy/adapters/types/__init__.py +++ b/dspy/adapters/types/__init__.py @@ -2,9 +2,9 @@ from dspy.adapters.types.base_type import Type from dspy.adapters.types.code import Code from dspy.adapters.types.file import File -from dspy.adapters.types.history import History +from dspy.adapters.types.history import History, HistoryCompaction from dspy.adapters.types.image import Image from dspy.adapters.types.reasoning import Reasoning from dspy.adapters.types.tool import Tool, ToolCalls -__all__ = ["History", "Image", "Audio", "File", "Type", "Tool", "ToolCalls", "Code", "Reasoning"] +__all__ = ["History", "HistoryCompaction", "Image", "Audio", "File", "Type", "Tool", "ToolCalls", "Code", "Reasoning"] diff --git a/dspy/adapters/types/history.py b/dspy/adapters/types/history.py index 6dda4f9b7c..9211d1f8a5 100644 --- a/dspy/adapters/types/history.py +++ b/dspy/adapters/types/history.py @@ -1,7 +1,29 @@ -from typing import Any +import copy +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Literal import pydantic +if TYPE_CHECKING: + from dspy.adapters.base import Adapter + from dspy.clients.lm import LM + + +class HistoryCompaction(pydantic.BaseModel): + max_visible_tokens: int = pydantic.Field(gt=0) + keep_last_messages: int = pydantic.Field(gt=0) + + model_config = pydantic.ConfigDict( + frozen=True, + str_strip_whitespace=True, + validate_assignment=True, + extra="forbid", + ) + + +_DEFAULT_RAW_HISTORY_COMPACTION = HistoryCompaction(max_visible_tokens=8_000, keep_last_messages=8) +_SUMMARY_PREFIX = "Summary of earlier conversation:\n" + class History(pydantic.BaseModel): """Class representing the conversation history. @@ -59,6 +81,10 @@ class MySignature(dspy.Signature): """ messages: list[dict[str, Any]] + mode: Literal["demo", "raw"] = "demo" + compaction: HistoryCompaction | None = None + summary: str | None = None + compacted_count: int = 0 model_config = pydantic.ConfigDict( frozen=True, @@ -66,3 +92,126 @@ class MySignature(dspy.Signature): validate_assignment=True, extra="forbid", ) + + @classmethod + def demo(cls, messages: list[dict[str, Any]] | None = None) -> "History": + return cls(messages=messages or [], mode="demo") + + @classmethod + def raw( + cls, + messages: list[dict[str, Any]] | None = None, + compaction: HistoryCompaction | None = None, + ) -> "History": + compaction = compaction or _DEFAULT_RAW_HISTORY_COMPACTION + return cls(messages=messages or [], mode="raw", compaction=compaction) + + def visible_messages(self) -> list[dict[str, Any]]: + if self.compacted_count == 0: + return copy.deepcopy(self.messages) + + visible_messages = [] + if self.summary is not None: + visible_messages.append({"role": "user", "content": f"{_SUMMARY_PREFIX}{self.summary}"}) + visible_messages.extend(self.messages[self.compacted_count :]) + return copy.deepcopy(visible_messages) + + def with_messages( + self, + new_messages: list[dict[str, Any]], + *, + lm: "LM | None", + adapter: "Adapter | None", + ) -> "History": + updated = type(self)( + messages=[*self.messages, *new_messages], + mode=self.mode, + compaction=self.compaction, + summary=self.summary, + compacted_count=self.compacted_count, + ) + return updated._maybe_compact(lm=lm, adapter=adapter) + + def _maybe_compact(self, *, lm: "LM | None", adapter: "Adapter | None") -> "History": + if self.mode != "raw" or self.compaction is None: + return self + + visible_tokens = self._count_visible_tokens(self.visible_messages(), lm) + if visible_tokens is None or visible_tokens <= self.compaction.max_visible_tokens: + return self + + new_compacted_count = len(self.messages) - self.compaction.keep_last_messages + if new_compacted_count <= self.compacted_count: + return self + + hidden_messages: list[dict[str, Any]] = [] + if self.summary is not None: + hidden_messages.append({"role": "user", "content": f"{_SUMMARY_PREFIX}{self.summary}"}) + hidden_messages.extend(self.messages[self.compacted_count : new_compacted_count]) + + if not hidden_messages or lm is None or adapter is None: + return self + + summary = self._summarize_prefix(hidden_messages, lm=lm, adapter=adapter) + return type(self)( + messages=self.messages, + mode=self.mode, + compaction=self.compaction, + summary=summary, + compacted_count=new_compacted_count, + ) + + def _count_visible_tokens(self, messages: list[dict[str, Any]], lm: "LM | None") -> int | None: + if lm is None or getattr(lm, "model", None) is None: + return None + + try: + from litellm.utils import token_counter + + return token_counter(model=lm.model, messages=messages) + except Exception: + return None + + def _summarize_prefix( + self, + hidden_messages: list[dict[str, Any]], + *, + lm: "LM", + adapter: "Adapter", + ) -> str: + compacted_history = type(self)(messages=hidden_messages, mode="raw", compaction=None) + + import dspy + + with dspy.context(lm=lm, adapter=adapter, trace=[]): + return _get_history_summarizer()(history=compacted_history).summary + + @pydantic.model_serializer(mode="plain") + def serialize_model(self) -> dict[str, Any]: + data: dict[str, Any] = {"messages": self.messages} + if self.mode != "demo": + data["mode"] = self.mode + if self.mode == "raw" and self.compaction is not None: + data["compaction"] = self.compaction.model_dump() + if self.summary is not None: + data["summary"] = self.summary + if self.compacted_count: + data["compacted_count"] = self.compacted_count + return data + + +@lru_cache(maxsize=1) +def _get_history_summarizer(): + import dspy + + class SummarizeHistory(dspy.Signature): + """Summarize the earlier conversation faithfully for future continuation. + + Preserve concrete facts, tool results, unresolved questions, and any relevant state. + Keep the summary concise. + """ + + history: History = dspy.InputField() + summary: str = dspy.OutputField() + + return dspy.Predict(SummarizeHistory) diff --git a/tests/primitives/test_example.py b/tests/primitives/test_example.py index c569016dbc..a8b39056a2 100644 --- a/tests/primitives/test_example.py +++ b/tests/primitives/test_example.py @@ -154,3 +154,25 @@ def test_example_to_dict_with_history(): json_str = json.dumps(result) restored = json.loads(json_str) assert restored["history"]["messages"] == result["history"]["messages"] + + +def test_example_to_dict_with_raw_history_compaction(): + history = dspy.History( + messages=[ + {"role": "assistant", "content": "Earlier step."}, + {"role": "tool", "name": "search", "content": "Earlier result."}, + {"role": "assistant", "content": "Latest step."}, + ], + mode="raw", + compaction=dspy.HistoryCompaction(max_visible_tokens=32, keep_last_messages=1), + summary="Earlier work.", + compacted_count=2, + ) + example = Example(question="Test question", history=history, answer="Test answer") + + result = example.toDict() + + assert result["history"]["mode"] == "raw" + assert result["history"]["compaction"] == {"max_visible_tokens": 32, "keep_last_messages": 1} + assert result["history"]["summary"] == "Earlier work." + assert result["history"]["compacted_count"] == 2