From 9669f958654cc36e7a7a9acfb4a8218776d7ee40 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Thu, 10 Jul 2025 10:51:17 +0530 Subject: [PATCH 1/4] feat: Integrate summarization node and update dependencies --- pyproject.toml | 1 + src/katalyst/coding_agent/nodes/planner.py | 3 +- src/katalyst/coding_agent/nodes/summarizer.py | 43 +++++++++++++++++++ .../katalyst_core/config/llm_config.py | 2 + 4 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 src/katalyst/coding_agent/nodes/summarizer.py diff --git a/pyproject.toml b/pyproject.toml index e3be0de..0b5b167 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "langchain-anthropic>=0.3.16,<0.4.0", "simple-term-menu>=1.6.4", "langgraph-checkpoint-sqlite (>=2.0.10,<3.0.0)", + "langmem (==0.0.27)", ] [project.scripts] diff --git a/src/katalyst/coding_agent/nodes/planner.py b/src/katalyst/coding_agent/nodes/planner.py index f751649..6ad1ccc 100644 --- a/src/katalyst/coding_agent/nodes/planner.py +++ b/src/katalyst/coding_agent/nodes/planner.py @@ -16,7 +16,7 @@ from langchain_core.tools import StructuredTool from langchain_core.messages import HumanMessage from langgraph.prebuilt import create_react_agent - +from katalyst.coding_agent.nodes.summarizer import get_summarization_node # Simple planner prompt - no complex guidelines planner_prompt = ChatPromptTemplate.from_messages( @@ -143,6 +143,7 @@ def sync_wrapper(**kwargs): state.agent_executor = create_react_agent( model=agent_model, tools=tools, + pre_model_hook=get_summarization_node(), checkpointer=state.checkpointer if hasattr(state, 'checkpointer') else False ) diff --git a/src/katalyst/coding_agent/nodes/summarizer.py b/src/katalyst/coding_agent/nodes/summarizer.py new file mode 100644 index 0000000..e1f51f0 --- /dev/null +++ b/src/katalyst/coding_agent/nodes/summarizer.py @@ -0,0 +1,43 @@ +from katalyst.katalyst_core.state import KatalystState +from langmem.short_term import SummarizationNode +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.messages.utils import count_tokens_approximately +from katalyst.katalyst_core.services.llms import get_llm_client +from katalyst.katalyst_core.utils.logger import get_logger + +logger = get_logger() + +#TODO: Move to prompts.py +SUMMARIZATION_PROMPT = """ +>> +>> +>> +""" + + +#TODO: Explanatory Variable Names +MAX_TOKENS = 10000 +MAX_TOKENS_BEFORE_SUMMARY = 1000 +MAX_SUMMARY_TOKENS = 1000 + +async def get_summarization_node(state: KatalystState): + initial_summary_prompt = ChatPromptTemplate.from_messages( + [ + ("placeholder", "{messages}"), + ("user", SUMMARIZATION_PROMPT), + ] + ) + client = get_llm_client("summarizer") + # Summarization Node () + summarization_node = SummarizationNode( + token_counter=count_tokens_approximately, + # Advised to use gpt-4.1 for summarization + model=client, + max_tokens=MAX_TOKENS, + max_tokens_before_summary=MAX_TOKENS_BEFORE_SUMMARY, + initial_summary_prompt=initial_summary_prompt, + max_summary_tokens=MAX_SUMMARY_TOKENS, + # Output key "messages" replace the existing messages with the summarized messages + remaining messages + output_messages_key="messages", + ) + return summarization_node \ No newline at end of file diff --git a/src/katalyst/katalyst_core/config/llm_config.py b/src/katalyst/katalyst_core/config/llm_config.py index c1bff47..3cd57b3 100644 --- a/src/katalyst/katalyst_core/config/llm_config.py +++ b/src/katalyst/katalyst_core/config/llm_config.py @@ -18,6 +18,7 @@ "execution": "gpt-4.1", # Fast execution tasks (agent_react, tools) "fallback": "gpt-4o", # Fallback model "default_timeout": 45, + }, "anthropic": { "reasoning": "claude-3-opus-20240229", # High-reasoning tasks @@ -55,6 +56,7 @@ "tool_runner": "execution", # Default for any other component "default": "execution", + "summarizer": "execution", } From 5d54ef7e516372c8c60889c8cea67f98d2b167e2 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Fri, 18 Jul 2025 00:22:53 +0530 Subject: [PATCH 2/4] feat: Enhance summarization functionality and introduce configuration constants - Added new configuration constants for summarization: MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT, MAX_TOKENS_TO_TRIGGER_SUMMARY, and MAX_TOKENS_IN_SUMMARY_ONLY. - Updated the summarization prompt to include detailed analysis structure and ensure comprehensive coverage of user requests and intents. - Refactored the get_summarization_node function to utilize the new configuration constants for token limits. - Introduced unit tests for the summarization node to validate its functionality and configuration. --- src/katalyst/app/config.py | 8 + src/katalyst/coding_agent/nodes/planner.py | 18 +- src/katalyst/coding_agent/nodes/summarizer.py | 87 +++++++-- tests/unit/test_summarizer.py | 167 ++++++++++++++++++ 4 files changed, 263 insertions(+), 17 deletions(-) create mode 100644 tests/unit/test_summarizer.py diff --git a/src/katalyst/app/config.py b/src/katalyst/app/config.py index 80e93cf..7e2b2f4 100644 --- a/src/katalyst/app/config.py +++ b/src/katalyst/app/config.py @@ -56,3 +56,11 @@ # Project specific ".katalyst", } + +#TODO: Explanatory Variable Names (Config Variables) +# Maximum number of tokens to return in the final output. Will be enforced only after summarization. +MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT = 50000 # 50k +# Maximum number of tokens to accumulate before triggering summarization. +MAX_TOKENS_TO_TRIGGER_SUMMARY = 40000 # 40k +# Maximum number of tokens to budget for the summary. +MAX_TOKENS_IN_SUMMARY_ONLY = 8000 # 8k \ No newline at end of file diff --git a/src/katalyst/coding_agent/nodes/planner.py b/src/katalyst/coding_agent/nodes/planner.py index 6ad1ccc..f54cae2 100644 --- a/src/katalyst/coding_agent/nodes/planner.py +++ b/src/katalyst/coding_agent/nodes/planner.py @@ -17,6 +17,15 @@ from langchain_core.messages import HumanMessage from langgraph.prebuilt import create_react_agent from katalyst.coding_agent.nodes.summarizer import get_summarization_node +from langgraph.prebuilt.chat_agent_executor import AgentState as LangGraphAgentState +from typing import Any + +class ReactAgentState(LangGraphAgentState): + """ + Custom state for the ReactAgent. + """ + context: dict[str, Any] + # Simple planner prompt - no complex guidelines planner_prompt = ChatPromptTemplate.from_messages( @@ -138,13 +147,16 @@ def sync_wrapper(**kwargs): timeout=timeout, api_base=api_base ) - + summarization_node = get_summarization_node() # Create the agent executor with checkpointer if available state.agent_executor = create_react_agent( model=agent_model, tools=tools, - pre_model_hook=get_summarization_node(), - checkpointer=state.checkpointer if hasattr(state, 'checkpointer') else False + pre_model_hook=summarization_node, + state_schema=ReactAgentState, + checkpointer=state.checkpointer if hasattr(state, 'checkpointer') else False, + # uncomment for debugging + debug=False, ) # Initialize conversation with the plan diff --git a/src/katalyst/coding_agent/nodes/summarizer.py b/src/katalyst/coding_agent/nodes/summarizer.py index e1f51f0..2388278 100644 --- a/src/katalyst/coding_agent/nodes/summarizer.py +++ b/src/katalyst/coding_agent/nodes/summarizer.py @@ -4,23 +4,82 @@ from langchain_core.messages.utils import count_tokens_approximately from katalyst.katalyst_core.services.llms import get_llm_client from katalyst.katalyst_core.utils.logger import get_logger - +from katalyst.app.config import MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT, MAX_TOKENS_TO_TRIGGER_SUMMARY, MAX_TOKENS_IN_SUMMARY_ONLY logger = get_logger() -#TODO: Move to prompts.py +#(Reference: https://www.reddit.com/r/ClaudeAI/comments/1jr52qj/here_is_claude_codes_compact_prompt/) SUMMARIZATION_PROMPT = """ ->> ->> ->> -""" +Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions. +This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context. + +Before providing your final summary, wrap your analysis in tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process: + +1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify: + - The user's explicit requests and intents + - Your approach to addressing the user's requests + - Key decisions, technical concepts and code patterns + - Specific details like file names, full code snippets, function signatures, file edits, etc +2. Double-check for technical accuracy and completeness, addressing each required element thoroughly. + +Your summary should include the following sections: + +1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail +2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed. +3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important. +4. Problem Solving: Document problems solved and any ongoing troubleshooting efforts. +5. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on. +6. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable. +7. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first. +8. If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation. + +Here's an example of how your output should be structured: + + + +[Your thought process, ensuring all points are covered thoroughly and accurately] + + + +1. Primary Request and Intent: + [Detailed description] +2. Key Technical Concepts: + - [Concept 1] + - [Concept 2] + - [...] + +3. Files and Code Sections: + - [File Name 1] + - [Summary of why this file is important] + - [Summary of the changes made to this file, if any] + - [Important Code Snippet] + - [File Name 2] + - [Important Code Snippet] + - [...] + +4. Problem Solving: + [Description of solved problems and ongoing troubleshooting] + +5. Pending Tasks: + - [Task 1] + - [Task 2] + - [...] + +6. Current Work: + [Precise description of current work] + +7. Optional Next Step: + [Optional Next step to take] + + + + +Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response. + +""" -#TODO: Explanatory Variable Names -MAX_TOKENS = 10000 -MAX_TOKENS_BEFORE_SUMMARY = 1000 -MAX_SUMMARY_TOKENS = 1000 -async def get_summarization_node(state: KatalystState): +def get_summarization_node(): initial_summary_prompt = ChatPromptTemplate.from_messages( [ ("placeholder", "{messages}"), @@ -33,10 +92,10 @@ async def get_summarization_node(state: KatalystState): token_counter=count_tokens_approximately, # Advised to use gpt-4.1 for summarization model=client, - max_tokens=MAX_TOKENS, - max_tokens_before_summary=MAX_TOKENS_BEFORE_SUMMARY, + max_tokens=MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT, + max_tokens_before_summary=MAX_TOKENS_TO_TRIGGER_SUMMARY, initial_summary_prompt=initial_summary_prompt, - max_summary_tokens=MAX_SUMMARY_TOKENS, + max_summary_tokens=MAX_TOKENS_IN_SUMMARY_ONLY, # Output key "messages" replace the existing messages with the summarized messages + remaining messages output_messages_key="messages", ) diff --git a/tests/unit/test_summarizer.py b/tests/unit/test_summarizer.py new file mode 100644 index 0000000..1299e38 --- /dev/null +++ b/tests/unit/test_summarizer.py @@ -0,0 +1,167 @@ +"""Tests for the summarizer node module.""" + +import pytest +from unittest.mock import MagicMock, patch, call +from langchain_core.messages import HumanMessage, AIMessage +from langchain_core.messages.utils import count_tokens_approximately + +from katalyst.coding_agent.nodes.summarizer import ( + get_summarization_node, + SUMMARIZATION_PROMPT, +) +from katalyst.app.config import ( + MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT, + MAX_TOKENS_TO_TRIGGER_SUMMARY, + MAX_TOKENS_IN_SUMMARY_ONLY, +) + + +class TestSummarizationNode: + """Test the summarization node functionality.""" + + def setup_method(self): + """Set up test fixtures.""" + self.mock_llm = MagicMock() + self.mock_llm.invoke.return_value = MagicMock(content="Test summary") + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_get_summarization_node_creation(self, mock_get_llm_client): + """Test that get_summarization_node creates a properly configured node.""" + mock_get_llm_client.return_value = self.mock_llm + + node = get_summarization_node() + + # Verify the LLM client was requested for summarizer + mock_get_llm_client.assert_called_once_with("summarizer") + + # Verify the node is properly configured + assert node.token_counter == count_tokens_approximately + assert node.model == self.mock_llm + assert node.max_tokens == MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT + assert node.max_tokens_before_summary == MAX_TOKENS_TO_TRIGGER_SUMMARY + assert node.max_summary_tokens == MAX_TOKENS_IN_SUMMARY_ONLY + assert node.output_messages_key == "messages" + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_summarization_prompt_structure(self, mock_get_llm_client): + """Test that the summarization prompt is properly structured.""" + mock_get_llm_client.return_value = self.mock_llm + + node = get_summarization_node() + + # Check that the prompt template has the expected structure + prompt_template = node.initial_summary_prompt + messages = prompt_template.messages + + assert len(messages) == 2 + assert messages[0].prompt.template == "{messages}" + assert messages[1].prompt.template == SUMMARIZATION_PROMPT + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_summarization_prompt_content(self, mock_get_llm_client): + """Test that the summarization prompt contains expected content.""" + mock_get_llm_client.return_value = self.mock_llm + + # Verify key sections are present in the prompt + assert "Primary Request and Intent" in SUMMARIZATION_PROMPT + assert "Key Technical Concepts" in SUMMARIZATION_PROMPT + assert "Files and Code Sections" in SUMMARIZATION_PROMPT + assert "Problem Solving" in SUMMARIZATION_PROMPT + assert "Pending Tasks" in SUMMARIZATION_PROMPT + assert "Current Work" in SUMMARIZATION_PROMPT + assert "Optional Next Step" in SUMMARIZATION_PROMPT + assert "" in SUMMARIZATION_PROMPT + assert "" in SUMMARIZATION_PROMPT + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_config_constants_used(self, mock_get_llm_client): + """Test that configuration constants are properly used.""" + mock_get_llm_client.return_value = self.mock_llm + + node = get_summarization_node() + + # Verify all config constants are used + assert node.max_tokens == MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT + assert node.max_tokens_before_summary == MAX_TOKENS_TO_TRIGGER_SUMMARY + assert node.max_summary_tokens == MAX_TOKENS_IN_SUMMARY_ONLY + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_multiple_node_instances(self, mock_get_llm_client): + """Test that multiple instances of the node can be created.""" + mock_get_llm_client.return_value = self.mock_llm + + node1 = get_summarization_node() + node2 = get_summarization_node() + + # Should be separate instances + assert node1 is not node2 + assert mock_get_llm_client.call_count == 2 + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_node_token_counter_integration(self, mock_get_llm_client): + """Test that the token counter function is properly integrated.""" + mock_get_llm_client.return_value = self.mock_llm + + node = get_summarization_node() + + # Test that the token counter is the expected function + assert callable(node.token_counter) + assert node.token_counter == count_tokens_approximately + + # Test that it works with sample messages + test_messages = [ + HumanMessage(content="Test message 1"), + AIMessage(content="Test response 1") + ] + + token_count = node.token_counter(test_messages) + assert isinstance(token_count, int) + assert token_count > 0 + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_output_messages_key_configuration(self, mock_get_llm_client): + """Test that the output messages key is correctly configured.""" + mock_get_llm_client.return_value = self.mock_llm + + node = get_summarization_node() + + # Verify the output key is set to "messages" + assert node.output_messages_key == "messages" + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + def test_summarization_node_type(self, mock_get_llm_client): + """Test that the correct type of summarization node is returned.""" + mock_get_llm_client.return_value = self.mock_llm + + node = get_summarization_node() + + # Should be a SummarizationNode from langmem + assert node.__class__.__name__ == "SummarizationNode" + + def test_prompt_constants_accessibility(self): + """Test that the prompt constant is accessible and properly formatted.""" + # Test that the prompt is a string + assert isinstance(SUMMARIZATION_PROMPT, str) + assert len(SUMMARIZATION_PROMPT) > 0 + + # Test that it contains the example structure + assert "" in SUMMARIZATION_PROMPT + assert "" in SUMMARIZATION_PROMPT + + # Test that it has the Reddit reference + assert "reddit.com" in SUMMARIZATION_PROMPT + + @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") + @patch("katalyst.coding_agent.nodes.summarizer.get_logger") + def test_logger_integration(self, mock_get_logger, mock_get_llm_client): + """Test that the logger is properly imported and available.""" + mock_get_llm_client.return_value = self.mock_llm + mock_logger = MagicMock() + mock_get_logger.return_value = mock_logger + + # Import should work without errors + from katalyst.coding_agent.nodes.summarizer import logger + + # Logger should be available + assert logger is not None + mock_get_logger.assert_called_once() \ No newline at end of file From 4838f7162865b9e4433aea668dd6e07def0a14de Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Fri, 18 Jul 2025 00:29:48 +0530 Subject: [PATCH 3/4] chore: Add Ruff linter configuration and update dependencies - Introduced Ruff as a development dependency for code linting and formatting. - Created a GitHub Actions workflow to automate Ruff checks on pull requests and pushes to the main branch. - Removed an unused import in the summarizer module to clean up the codebase. --- .github/workflows/ruff.yml | 24 +++++++++++++++++++ pyproject.toml | 3 +++ src/katalyst/coding_agent/nodes/summarizer.py | 1 - 3 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/ruff.yml diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml new file mode 100644 index 0000000..892036b --- /dev/null +++ b/.github/workflows/ruff.yml @@ -0,0 +1,24 @@ +name: Ruff + +on: + pull_request: + branches: [ main ] + push: + branches: [ main ] + +jobs: + ruff: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: '3.12' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install ruff + - name: Run Ruff linter + run: ruff check . + - name: Run Ruff formatter + run: ruff format --check . \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0b5b167..dda1aa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,9 @@ langchain-ollama = "^0.3.3" langchain-anthropic = "^0.3.16" simple-term-menu = "^1.6.4" +[tool.poetry.group.dev.dependencies] +ruff = "^0.12.4" + [tool.setuptools.packages.find] where = ["src"] diff --git a/src/katalyst/coding_agent/nodes/summarizer.py b/src/katalyst/coding_agent/nodes/summarizer.py index 2388278..36eef95 100644 --- a/src/katalyst/coding_agent/nodes/summarizer.py +++ b/src/katalyst/coding_agent/nodes/summarizer.py @@ -1,4 +1,3 @@ -from katalyst.katalyst_core.state import KatalystState from langmem.short_term import SummarizationNode from langchain_core.prompts import ChatPromptTemplate from langchain_core.messages.utils import count_tokens_approximately From 826826f0fc54dba511d43f54c558661106cdc12a Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Fri, 18 Jul 2025 01:06:12 +0530 Subject: [PATCH 4/4] refactor: Remove unused unit tests from summarization node tests - Deleted tests that verify the usage of configuration constants and output messages key in the summarization node. - Streamlined the test suite to focus on essential functionality and improve maintainability. --- tests/unit/test_summarizer.py | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/tests/unit/test_summarizer.py b/tests/unit/test_summarizer.py index 1299e38..f9e17c5 100644 --- a/tests/unit/test_summarizer.py +++ b/tests/unit/test_summarizer.py @@ -73,18 +73,6 @@ def test_summarization_prompt_content(self, mock_get_llm_client): assert "" in SUMMARIZATION_PROMPT assert "" in SUMMARIZATION_PROMPT - @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") - def test_config_constants_used(self, mock_get_llm_client): - """Test that configuration constants are properly used.""" - mock_get_llm_client.return_value = self.mock_llm - - node = get_summarization_node() - - # Verify all config constants are used - assert node.max_tokens == MAX_AGGREGATE_TOKENS_IN_SUMMARY_AND_OUTPUT - assert node.max_tokens_before_summary == MAX_TOKENS_TO_TRIGGER_SUMMARY - assert node.max_summary_tokens == MAX_TOKENS_IN_SUMMARY_ONLY - @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") def test_multiple_node_instances(self, mock_get_llm_client): """Test that multiple instances of the node can be created.""" @@ -118,16 +106,6 @@ def test_node_token_counter_integration(self, mock_get_llm_client): assert isinstance(token_count, int) assert token_count > 0 - @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") - def test_output_messages_key_configuration(self, mock_get_llm_client): - """Test that the output messages key is correctly configured.""" - mock_get_llm_client.return_value = self.mock_llm - - node = get_summarization_node() - - # Verify the output key is set to "messages" - assert node.output_messages_key == "messages" - @patch("katalyst.coding_agent.nodes.summarizer.get_llm_client") def test_summarization_node_type(self, mock_get_llm_client): """Test that the correct type of summarization node is returned."""