Custom summarization prompt

README.md

@@ -24,8 +24,15 @@ Pre-built Docker images are available from:
 
 **Quick Start (Development Mode)**:
 ```bash
-# Start with docker-compose (includes Redis, API, MCP, and worker)
-docker-compose up
+# Start with docker-compose
+# Note: Both 'api' and 'api-for-task-worker' services use port 8000
+# Choose one depending on your needs:
+
+# Option 1: Development mode (no worker, immediate task execution)
+docker compose up api redis
+
+# Option 2: Production-like mode (with background worker)
+docker compose up api-for-task-worker task-worker redis mcp
 
 # Or run just the API server (requires separate Redis)
 docker run -p 8000:8000 \
@@ -206,8 +213,9 @@ uv run pytest
 uv run ruff format
 uv run ruff check
 
-# Start development stack
-docker-compose up
+# Start development stack (choose one based on your needs)
+docker compose up api redis                               # Development mode
+docker compose up api-for-task-worker task-worker redis   # Production-like mode
 ```
 ## License
 

agent_memory_server/config.py

@@ -4,7 +4,7 @@
 
 import yaml
 from dotenv import load_dotenv
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 from pydantic_settings import BaseSettings
 
 
@@ -290,6 +290,71 @@ class Settings(BaseSettings):
     # Compaction settings
     compaction_every_minutes: int = 10
 
+    # Progressive summarization prompt template
+    progressive_summarization_prompt: str = """You are a precise summarization assistant. Your task is to progressively
+summarize conversation history while maintaining critical context and accuracy.
+
+INSTRUCTIONS:
+1. Build upon the previous summary by incorporating new information chronologically
+2. Preserve key details: names, technical terms, code references, and important decisions
+3. Maintain the temporal sequence of events and discussions
+4. For technical discussions, keep specific terms, versions, and implementation details
+5. For code-related content, preserve function names, file paths, and important parameters
+6. If the new content is irrelevant or doesn't add value, return "NONE"
+7. Keep the summary concise but complete - aim for 2-3 sentences unless more detail is crucial
+8. Use neutral, factual language
+
+EXAMPLE
+Current summary:
+The user inquires about retirement investment options, specifically comparing
+traditional IRAs and Roth IRAs. The assistant explains the key differences in
+tax treatment, with traditional IRAs offering immediate tax deductions and Roth
+IRAs providing tax-free withdrawals in retirement.
+
+New lines of conversation:
+Human: What factors should I consider when deciding between the two?
+Assistant: Several key factors influence this decision: 1) Your current tax
+bracket vs. expected retirement tax bracket, 2) Time horizon until retirement,
+3) Current income and eligibility for Roth IRA contributions, and 4) Desire for
+flexibility in retirement withdrawals. For example, if you expect to be in a
+higher tax bracket during retirement, a Roth IRA might be more advantageous
+since qualified withdrawals are tax-free. Additionally, Roth IRAs don't have
+required minimum distributions (RMDs) during your lifetime, offering more
+flexibility in estate planning.
+
+New summary:
+The discussion covers retirement investment options, comparing traditional and
+Roth IRAs' tax implications, with traditional IRAs offering immediate deductions
+and Roth IRAs providing tax-free withdrawals. The conversation expands to cover
+decision factors including current vs. future tax brackets, retirement timeline,
+income eligibility, and withdrawal flexibility, with specific emphasis on Roth
+IRA advantages for those expecting higher retirement tax brackets and the
+benefit of no required minimum distributions. END OF EXAMPLE
+
+Current summary:
+{prev_summary}
+
+New lines of conversation:
+{messages_joined}
+
+New summary:
+"""
+
+    @field_validator("progressive_summarization_prompt")
+    @classmethod
+    def validate_progressive_summarization_prompt(cls, v: str) -> str:
+        """Validate that the progressive summarization prompt contains required placeholders."""
+        required_vars = ["prev_summary", "messages_joined"]
+        missing_vars = [var for var in required_vars if f"{{{var}}}" not in v]
+
+        if missing_vars:
+            raise ValueError(
+                f"progressive_summarization_prompt must contain the following placeholders: "
+                f"{', '.join(f'{{{var}}}' for var in missing_vars)}"
+            )
+
+        return v
+
     class Config:
         env_file = ".env"
         env_file_encoding = "utf-8"

agent_memory_server/summarization.py

@@ -42,56 +42,10 @@ async def _incremental_summary(
     messages_joined = "\n".join(messages)
     prev_summary = context or ""
 
-    # Prompt template for progressive summarization
-    progressive_prompt = f"""
-You are a precise summarization assistant. Your task is to progressively
-summarize conversation history while maintaining critical context and accuracy.
-
-INSTRUCTIONS:
-1. Build upon the previous summary by incorporating new information chronologically
-2. Preserve key details: names, technical terms, code references, and important decisions
-3. Maintain the temporal sequence of events and discussions
-4. For technical discussions, keep specific terms, versions, and implementation details
-5. For code-related content, preserve function names, file paths, and important parameters
-6. If the new content is irrelevant or doesn't add value, return "NONE"
-7. Keep the summary concise but complete - aim for 2-3 sentences unless more detail is crucial
-8. Use neutral, factual language
-
-EXAMPLE
-Current summary:
-The user inquires about retirement investment options, specifically comparing
-traditional IRAs and Roth IRAs. The assistant explains the key differences in
-tax treatment, with traditional IRAs offering immediate tax deductions and Roth
-IRAs providing tax-free withdrawals in retirement.
-
-New lines of conversation:
-Human: What factors should I consider when deciding between the two?
-Assistant: Several key factors influence this decision: 1) Your current tax
-bracket vs. expected retirement tax bracket, 2) Time horizon until retirement,
-3) Current income and eligibility for Roth IRA contributions, and 4) Desire for
-flexibility in retirement withdrawals. For example, if you expect to be in a
-higher tax bracket during retirement, a Roth IRA might be more advantageous
-since qualified withdrawals are tax-free. Additionally, Roth IRAs don't have
-required minimum distributions (RMDs) during your lifetime, offering more
-flexibility in estate planning.
-
-New summary:
-The discussion covers retirement investment options, comparing traditional and
-Roth IRAs' tax implications, with traditional IRAs offering immediate deductions
-and Roth IRAs providing tax-free withdrawals. The conversation expands to cover
-decision factors including current vs. future tax brackets, retirement timeline,
-income eligibility, and withdrawal flexibility, with specific emphasis on Roth
-IRA advantages for those expecting higher retirement tax brackets and the
-benefit of no required minimum distributions. END OF EXAMPLE
-
-Current summary:
-{prev_summary}
-
-New lines of conversation:
-{messages_joined}
-
-New summary:
-"""
+    # Use configurable prompt template for progressive summarization
+    progressive_prompt = settings.progressive_summarization_prompt.format(
+        prev_summary=prev_summary, messages_joined=messages_joined
+    )
 
     try:
         # Get completion from client

tests/test_summarization.py

@@ -2,7 +2,9 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from pydantic import ValidationError
 
+from agent_memory_server.config import Settings, settings
 from agent_memory_server.summarization import (
     _incremental_summary,
     summarize_session,
@@ -174,7 +176,7 @@ async def test_summarize_session(
     @pytest.mark.asyncio
     @patch("agent_memory_server.summarization._incremental_summary")
     async def test_handle_summarization_no_messages(
-        self, mock_summarization, mock_openai_client, mock_async_redis_client
+        self, mock_summarization, mock_async_redis_client
     ):
         """Test summarize_session when no messages need summarization"""
         session_id = "test-session"
@@ -218,3 +220,77 @@ async def test_handle_summarization_no_messages(
         assert pipeline_mock.hmset.call_count == 0
         assert pipeline_mock.ltrim.call_count == 0
         assert pipeline_mock.execute.call_count == 0
+
+    @pytest.mark.asyncio
+    async def test_configurable_summarization_prompt(self, mock_openai_client):
+        """Test that the summarization prompt can be configured"""
+        model = "gpt-3.5-turbo"
+        context = "Previous context"
+        messages = ["User: Hello", "Assistant: Hi there"]
+
+        # Create a custom prompt template
+        custom_prompt = "Custom prompt: {prev_summary} | {messages_joined}"
+
+        mock_response = MagicMock()
+        mock_choices = MagicMock()
+        mock_choices.message = MagicMock()
+        mock_choices.message.content = "Custom summary"
+        mock_response.choices = [mock_choices]
+        mock_response.total_tokens = 100
+
+        mock_openai_client.create_chat_completion.return_value = mock_response
+
+        # Temporarily override the prompt setting
+        original_prompt = settings.progressive_summarization_prompt
+        try:
+            settings.progressive_summarization_prompt = custom_prompt
+
+            summary, tokens_used = await _incremental_summary(
+                model, mock_openai_client, context, messages
+            )
+
+            assert summary == "Custom summary"
+            assert tokens_used == 100
+
+            # Verify the custom prompt was used
+            mock_openai_client.create_chat_completion.assert_called_once()
+            args = mock_openai_client.create_chat_completion.call_args[0]
+            assert "Custom prompt:" in args[1]
+            assert "Previous context" in args[1]
+        finally:
+            # Restore original prompt
+            settings.progressive_summarization_prompt = original_prompt
+
+    def test_prompt_validation_missing_prev_summary(self):
+        """Test that validation fails when {prev_summary} is missing"""
+        with pytest.raises(ValidationError) as exc_info:
+            Settings(
+                progressive_summarization_prompt="Template with {messages_joined} only"
+            )
+
+        assert "prev_summary" in str(exc_info.value)
+
+    def test_prompt_validation_missing_messages_joined(self):
+        """Test that validation fails when {messages_joined} is missing"""
+        with pytest.raises(ValidationError) as exc_info:
+            Settings(
+                progressive_summarization_prompt="Template with {prev_summary} only"
+            )
+
+        assert "messages_joined" in str(exc_info.value)
+
+    def test_prompt_validation_missing_both_variables(self):
+        """Test that validation fails when both required variables are missing"""
+        with pytest.raises(ValidationError) as exc_info:
+            Settings(progressive_summarization_prompt="Template with no variables")
+
+        error_str = str(exc_info.value)
+        assert "prev_summary" in error_str
+        assert "messages_joined" in error_str
+
+    def test_prompt_validation_with_both_variables(self):
+        """Test that validation passes when both variables are present"""
+        custom_prompt = "Summary: {prev_summary} Messages: {messages_joined}"
+        config = Settings(progressive_summarization_prompt=custom_prompt)
+
+        assert config.progressive_summarization_prompt == custom_prompt