Implementation:Hpcaitech ColossalAI ConversationSummaryMemory

Knowledge Sources	Hpcaitech_ColossalAI
Domains	NLP, Conversational AI, Memory Management
Last Updated	2026-02-09 00:00 GMT

Overview

ConversationSummaryMemory is a chat memory class that maintains a running summary of the conversation history using an LLM-based summarizer, built on top of LangChain's BaseChatMemory and a custom SummarizerMixin.

Description

This module provides two classes: SummarizerMixin, a base mixin that uses an LLM chain to recursively generate conversation summaries, and ConversationSummaryMemory, which integrates that summarization capability into LangChain's chat memory system. The summary buffer is updated after each conversation turn by predicting a new summary from the most recent messages and the existing summary. The code is modified from LangChain's original implementation and is licensed under the MIT license.

Usage

Use ConversationSummaryMemory when building a ColossalQA conversational pipeline where you need to compress long conversation histories into a concise summary rather than storing the full message list, enabling more efficient use of the LLM's context window.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalQA/colossalqa/chain/memory/summary.py
• Lines: 1-103

Signature

class SummarizerMixin(BaseModel):
    human_prefix: str = "Human"
    ai_prefix: str = "Assistant"
    llm: BaseLanguageModel
    prompt: BasePromptTemplate = SUMMARY_PROMPT
    summary_message_cls: Type[BaseMessage] = SystemMessage
    llm_kwargs: Dict = {}

    def predict_new_summary(self, messages: List[BaseMessage], existing_summary: str, stop: List = []) -> str:
        ...

class ConversationSummaryMemory(BaseChatMemory, SummarizerMixin):
    buffer: str = ""
    memory_key: str = "history"

    @classmethod
    def from_messages(
        cls,
        llm: BaseLanguageModel,
        chat_memory: BaseChatMessageHistory,
        summarize_step: int = 2,
        **kwargs: Any,
    ) -> ConversationSummaryMemory:
        ...

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        ...

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        ...

    def clear(self) -> None:
        ...

Import

from colossalqa.chain.memory.summary import ConversationSummaryMemory, SummarizerMixin

I/O Contract

Inputs

Name	Type	Required	Description
llm	BaseLanguageModel	Yes	The language model used to generate conversation summaries
chat_memory	BaseChatMessageHistory	No	Chat message history object to initialize from (used in from_messages classmethod)
summarize_step	int	No	Number of messages to summarize at a time when initializing from existing messages (default: 2)
human_prefix	str	No	Prefix label for human messages (default: "Human")
ai_prefix	str	No	Prefix label for AI messages (default: "Assistant")
prompt	BasePromptTemplate	No	The prompt template for summarization (default: SUMMARY_PROMPT)
memory_key	str	No	Key name under which the summary buffer is stored (default: "history")

Outputs

Name	Type	Description
load_memory_variables return	Dict[str, Any]	A dictionary mapping the memory_key to the current summary buffer (as a string or list of messages)
predict_new_summary return	str	The newly generated summary string combining existing summary with new conversation lines

Usage Examples

from langchain.schema.language_model import BaseLanguageModel
from colossalqa.chain.memory.summary import ConversationSummaryMemory

# Initialize with an LLM
memory = ConversationSummaryMemory(llm=my_llm)

# Save a conversation turn
memory.save_context(
    inputs={"input": "What is ColossalAI?"},
    outputs={"output": "ColossalAI is a distributed deep learning framework."}
)

# Retrieve the summarized history
variables = memory.load_memory_variables({})
print(variables["history"])

# Initialize from existing message history
memory = ConversationSummaryMemory.from_messages(
    llm=my_llm,
    chat_memory=existing_chat_history,
    summarize_step=2,
)

Related Pages

• Environment:Hpcaitech_ColossalAI_ColossalQA_RAG_Environment