diff --git a/README.md b/README.md index 584b9b3..423ba78 100644 --- a/README.md +++ b/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 diff --git a/agent_memory_server/config.py b/agent_memory_server/config.py index d1e065c..1379e5a 100644 --- a/agent_memory_server/config.py +++ b/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" diff --git a/agent_memory_server/summarization.py b/agent_memory_server/summarization.py index bbffb9e..75adf10 100644 --- a/agent_memory_server/summarization.py +++ b/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 diff --git a/tests/test_summarization.py b/tests/test_summarization.py index 3c92c0e..1c496fd 100644 --- a/tests/test_summarization.py +++ b/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