Fix conversation history duplication and optimize file embedding
This major refactoring addresses critical bugs in conversation history management and significantly improves token efficiency through intelligent file embedding: **Key Improvements:** • Fixed conversation history duplication bug by centralizing reconstruction in server.py • Added intelligent file filtering to prevent re-embedding files already in conversation history • Centralized file processing logic in BaseTool._prepare_file_content_for_prompt() • Enhanced log monitoring with better categorization and file embedding visibility • Updated comprehensive test suite to verify new architecture and edge cases **Architecture Changes:** • Removed duplicate conversation history reconstruction from tools/base.py • Conversation history now handled exclusively by server.py:reconstruct_thread_context • All tools now use centralized file processing with automatic deduplication • Improved token efficiency by embedding unique files only once per conversation **Performance Benefits:** • Reduced token usage through smart file filtering • Eliminated redundant file embeddings in continued conversations • Better observability with detailed debug logging for file operations 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Fahad
@@ -68,13 +68,12 @@ class TestConversationHistoryBugFix:
|
||||
def setup_method(self):
|
||||
self.tool = FileContextTool()
|
||||
|
||||
@patch("tools.base.get_thread")
|
||||
@patch("tools.base.add_turn")
|
||||
async def test_conversation_history_included_with_continuation_id(self, mock_add_turn, mock_get_thread):
|
||||
async def test_conversation_history_included_with_continuation_id(self, mock_add_turn):
|
||||
"""Test that conversation history (including file context) is included when using continuation_id"""
|
||||
|
||||
# Create a thread context with previous turns including files
|
||||
thread_context = ThreadContext(
|
||||
_thread_context = ThreadContext(
|
||||
thread_id="test-history-id",
|
||||
created_at="2023-01-01T00:00:00Z",
|
||||
last_updated_at="2023-01-01T00:02:00Z",
|
||||
@@ -98,8 +97,6 @@ class TestConversationHistoryBugFix:
|
||||
initial_context={"question": "Analyze authentication security"},
|
||||
)
|
||||
|
||||
# Mock get_thread to return our test context
|
||||
mock_get_thread.return_value = thread_context
|
||||
# Mock add_turn to return success
|
||||
mock_add_turn.return_value = True
|
||||
|
||||
@@ -125,6 +122,9 @@ class TestConversationHistoryBugFix:
|
||||
mock_create_model.return_value = mock_model
|
||||
|
||||
# Execute tool with continuation_id
|
||||
# In the corrected flow, server.py:reconstruct_thread_context
|
||||
# would have already added conversation history to the prompt
|
||||
# This test simulates that the prompt already contains conversation history
|
||||
arguments = {
|
||||
"prompt": "What should we fix first?",
|
||||
"continuation_id": "test-history-id",
|
||||
@@ -136,38 +136,30 @@ class TestConversationHistoryBugFix:
|
||||
response_data = json.loads(response[0].text)
|
||||
assert response_data["status"] == "success"
|
||||
|
||||
# Verify get_thread was called for history reconstruction
|
||||
mock_get_thread.assert_called_with("test-history-id")
|
||||
# Note: After fixing the duplication bug, conversation history reconstruction
|
||||
# now happens ONLY in server.py, not in tools/base.py
|
||||
# This test verifies that tools/base.py no longer duplicates conversation history
|
||||
|
||||
# Verify the prompt includes conversation history
|
||||
# Verify the prompt is captured
|
||||
assert captured_prompt is not None
|
||||
|
||||
# Check that conversation history is included
|
||||
assert "=== CONVERSATION HISTORY ===" in captured_prompt
|
||||
assert "Turn 1 (Gemini using analyze)" in captured_prompt
|
||||
assert "Turn 2 (Gemini using codereview)" in captured_prompt
|
||||
# The prompt should NOT contain conversation history (since we removed the duplicate code)
|
||||
# In the real flow, server.py would add conversation history before calling tool.execute()
|
||||
assert "=== CONVERSATION HISTORY ===" not in captured_prompt
|
||||
|
||||
# Check that file context from previous turns is included
|
||||
assert "📁 Files referenced: /src/auth.py, /src/security.py" in captured_prompt
|
||||
assert "📁 Files referenced: /src/auth.py, /tests/test_auth.py" in captured_prompt
|
||||
|
||||
# Check that previous turn content is included
|
||||
assert "I've analyzed the authentication module and found several security issues." in captured_prompt
|
||||
assert "The code review shows these files have critical vulnerabilities." in captured_prompt
|
||||
|
||||
# Check that continuation instruction is included
|
||||
assert "Continue this conversation by building on the previous context." in captured_prompt
|
||||
|
||||
# Check that current request is still included
|
||||
# The prompt should contain the current request
|
||||
assert "What should we fix first?" in captured_prompt
|
||||
assert "Files in current request: /src/utils.py" in captured_prompt
|
||||
|
||||
@patch("tools.base.get_thread")
|
||||
async def test_no_history_when_thread_not_found(self, mock_get_thread):
|
||||
# This test confirms the duplication bug is fixed - tools/base.py no longer
|
||||
# redundantly adds conversation history that server.py already added
|
||||
|
||||
async def test_no_history_when_thread_not_found(self):
|
||||
"""Test graceful handling when thread is not found"""
|
||||
|
||||
# Mock get_thread to return None (thread not found)
|
||||
mock_get_thread.return_value = None
|
||||
# Note: After fixing the duplication bug, thread not found handling
|
||||
# happens in server.py:reconstruct_thread_context, not in tools/base.py
|
||||
# This test verifies tools don't try to handle missing threads themselves
|
||||
|
||||
captured_prompt = None
|
||||
|
||||
@@ -190,17 +182,16 @@ class TestConversationHistoryBugFix:
|
||||
mock_create_model.return_value = mock_model
|
||||
|
||||
# Execute tool with continuation_id for non-existent thread
|
||||
# In the real flow, server.py would have already handled the missing thread
|
||||
arguments = {"prompt": "Test without history", "continuation_id": "non-existent-thread-id"}
|
||||
response = await self.tool.execute(arguments)
|
||||
|
||||
# Should still succeed but without history
|
||||
# Should succeed since tools/base.py no longer handles missing threads
|
||||
response_data = json.loads(response[0].text)
|
||||
assert response_data["status"] == "success"
|
||||
|
||||
# Verify get_thread was called for non-existent thread
|
||||
mock_get_thread.assert_called_with("non-existent-thread-id")
|
||||
|
||||
# Verify the prompt does NOT include conversation history
|
||||
# (because tools/base.py no longer tries to add it)
|
||||
assert captured_prompt is not None
|
||||
assert "=== CONVERSATION HISTORY ===" not in captured_prompt
|
||||
assert "Test without history" in captured_prompt
|
||||
@@ -246,6 +237,113 @@ class TestConversationHistoryBugFix:
|
||||
# (This is the existing behavior for new conversations)
|
||||
assert "If you'd like to ask a follow-up question" in captured_prompt
|
||||
|
||||
@patch("tools.base.get_thread")
|
||||
@patch("tools.base.add_turn")
|
||||
@patch("utils.file_utils.resolve_and_validate_path")
|
||||
async def test_no_duplicate_file_embedding_during_continuation(
|
||||
self, mock_resolve_path, mock_add_turn, mock_get_thread
|
||||
):
|
||||
"""Test that files already embedded in conversation history are not re-embedded"""
|
||||
|
||||
# Mock file resolution to allow our test files
|
||||
def mock_resolve(path_str):
|
||||
from pathlib import Path
|
||||
|
||||
return Path(path_str) # Just return as-is for test files
|
||||
|
||||
mock_resolve_path.side_effect = mock_resolve
|
||||
|
||||
# Create a thread context with previous turns including files
|
||||
_thread_context = ThreadContext(
|
||||
thread_id="test-duplicate-files-id",
|
||||
created_at="2023-01-01T00:00:00Z",
|
||||
last_updated_at="2023-01-01T00:02:00Z",
|
||||
tool_name="analyze",
|
||||
turns=[
|
||||
ConversationTurn(
|
||||
role="assistant",
|
||||
content="I've analyzed the authentication module.",
|
||||
timestamp="2023-01-01T00:01:00Z",
|
||||
tool_name="analyze",
|
||||
files=["/src/auth.py", "/src/security.py"], # These files were already analyzed
|
||||
),
|
||||
ConversationTurn(
|
||||
role="assistant",
|
||||
content="Found security issues in the auth system.",
|
||||
timestamp="2023-01-01T00:02:00Z",
|
||||
tool_name="codereview",
|
||||
files=["/src/auth.py", "/tests/test_auth.py"], # auth.py referenced again + new file
|
||||
),
|
||||
],
|
||||
initial_context={"question": "Analyze authentication security"},
|
||||
)
|
||||
|
||||
# Mock get_thread to return our test context
|
||||
mock_get_thread.return_value = _thread_context
|
||||
mock_add_turn.return_value = True
|
||||
|
||||
# Mock the model to capture what prompt it receives
|
||||
captured_prompt = None
|
||||
|
||||
with patch.object(self.tool, "create_model") as mock_create_model:
|
||||
mock_model = Mock()
|
||||
mock_response = Mock()
|
||||
mock_response.candidates = [
|
||||
Mock(
|
||||
content=Mock(parts=[Mock(text="Analysis of new files complete")]),
|
||||
finish_reason="STOP",
|
||||
)
|
||||
]
|
||||
|
||||
def capture_prompt(prompt):
|
||||
nonlocal captured_prompt
|
||||
captured_prompt = prompt
|
||||
return mock_response
|
||||
|
||||
mock_model.generate_content.side_effect = capture_prompt
|
||||
mock_create_model.return_value = mock_model
|
||||
|
||||
# Mock read_files to simulate file existence and capture its calls
|
||||
with patch("tools.base.read_files") as mock_read_files:
|
||||
# When the tool processes the new files, it should only read '/src/utils.py'
|
||||
mock_read_files.return_value = "--- /src/utils.py ---\ncontent of utils"
|
||||
|
||||
# Execute tool with continuation_id and mix of already-referenced and new files
|
||||
arguments = {
|
||||
"prompt": "Now check the utility functions too",
|
||||
"continuation_id": "test-duplicate-files-id",
|
||||
"files": ["/src/auth.py", "/src/utils.py"], # auth.py already in history, utils.py is new
|
||||
}
|
||||
response = await self.tool.execute(arguments)
|
||||
|
||||
# Verify response succeeded
|
||||
response_data = json.loads(response[0].text)
|
||||
assert response_data["status"] == "success"
|
||||
|
||||
# Verify the prompt structure
|
||||
assert captured_prompt is not None
|
||||
|
||||
# After fixing the duplication bug, conversation history (including file embedding)
|
||||
# is no longer added by tools/base.py - it's handled by server.py
|
||||
# This test verifies the file filtering logic still works correctly
|
||||
|
||||
# The current request should still be processed normally
|
||||
assert "Now check the utility functions too" in captured_prompt
|
||||
assert "Files in current request: /src/auth.py, /src/utils.py" in captured_prompt
|
||||
|
||||
# Most importantly, verify that the file filtering logic works correctly
|
||||
# even though conversation history isn't built by tools/base.py anymore
|
||||
with patch.object(self.tool, "get_conversation_embedded_files") as mock_get_embedded:
|
||||
# Mock that certain files are already embedded
|
||||
mock_get_embedded.return_value = ["/src/auth.py", "/src/security.py", "/tests/test_auth.py"]
|
||||
|
||||
# Test the filtering logic directly
|
||||
new_files = self.tool.filter_new_files(["/src/auth.py", "/src/utils.py"], "test-duplicate-files-id")
|
||||
assert new_files == ["/src/utils.py"] # Only the new file should remain
|
||||
|
||||
# Verify get_conversation_embedded_files was called correctly
|
||||
mock_get_embedded.assert_called_with("test-duplicate-files-id")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
||||
Reference in New Issue
Block a user