track related files through memory

This commit is contained in:
AI Christianson 2024-12-11 19:38:49 -05:00
parent 14abec9735
commit 33f5f9bbbb
4 changed files with 83 additions and 54 deletions

View File

@ -1,5 +1,5 @@
from .shell import run_shell_command
from .programmer import run_programming_task, emit_related_files
from .programmer import run_programming_task
from .expert import ask_expert, emit_expert_context
from .read_file import read_file_tool
from .fuzzy_find import fuzzy_find_project_files
@ -9,7 +9,7 @@ from .note_tech_debt import note_tech_debt
from .memory import (
emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
request_implementation, skip_implementation, delete_key_facts, emit_research_subtask,
emit_key_snippets, delete_key_snippets
emit_key_snippets, delete_key_snippets, emit_related_files
)
__all__ = [
@ -20,7 +20,7 @@ __all__ = [
'emit_key_facts',
'emit_key_snippets',
'emit_plan',
'emit_related_files',
'emit_related_files',
'emit_research_notes',
'emit_task',
'fuzzy_find_project_files',
@ -32,7 +32,6 @@ __all__ = [
'run_shell_command',
'skip_implementation',
'emit_research_subtask',
'fuzzy_find_project_files',
'ripgrep_search',
'note_tech_debt'
]

View File

@ -1,4 +1,4 @@
from typing import Dict, List, Any, Union, TypedDict, Optional, Sequence
from typing import Dict, List, Any, Union, TypedDict, Optional, Sequence, Set
from ra_aid.exceptions import TaskCompletedException
from rich.console import Console
from rich.markdown import Markdown
@ -15,7 +15,7 @@ class SnippetInfo(TypedDict):
console = Console()
# Global memory store
_global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo], int]] = {
_global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo], int, Set[str]]] = {
'research_notes': [],
'plans': [],
'tasks': [],
@ -25,7 +25,8 @@ _global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo
'key_snippets': {}, # Dict[int, SnippetInfo] - ID to snippet mapping
'key_snippet_id_counter': 0, # Counter for generating unique snippet IDs
'implementation_requested': [],
'implementation_skipped': []
'implementation_skipped': [],
'related_files': set()
}
@tool("emit_research_notes")
@ -259,6 +260,44 @@ def one_shot_completed(message: str) -> str:
raise ValueError("Cannot complete in one shot - implementation was requested")
raise TaskCompletedException(message)
def get_related_files() -> Set[str]:
"""Get the current set of related files.
Returns:
Set of file paths that have been marked as related
"""
return _global_memory['related_files']
@tool("emit_related_files")
def emit_related_files(files: List[str]) -> str:
"""Store multiple related files that tools should work with.
Args:
files: List of file paths to add
Returns:
Confirmation message
"""
results = []
added_files = []
# Process unique files
for file in set(files): # Remove duplicates in input
if file not in _global_memory['related_files']:
_global_memory['related_files'].add(file)
added_files.append(file)
results.append(f"Added related file: {file}")
# Rich output - single consolidated panel
if added_files:
files_added_md = '\n'.join(f"- `{file}`" for file in added_files)
md_content = f"**Files Noted:**\n{files_added_md}"
console.print(Panel(Markdown(md_content),
title="📁 Related Files Noted",
border_style="green"))
return "Files noted."
def get_memory_value(key: str) -> str:
"""Get a value from global memory.

View File

@ -8,46 +8,10 @@ from rich.markdown import Markdown
from rich.text import Text
from ra_aid.proc.interactive import run_interactive_command
from pydantic import BaseModel, Field
from .memory import get_memory_value
from ra_aid.text.processing import truncate_output
console = Console()
# Keep track of related files globally
related_files: List[str] = []
related_files_set: Set[str] = set()
@tool("emit_related_files")
def emit_related_files(files: List[str]) -> List[str]:
"""Store multiple related files that the programmer tool should work with.
Args:
files: List of file paths to add
Returns:
List of confirmation messages for added files
"""
global related_files, related_files_set
results = []
added_files = []
# Process unique files
for file in set(files): # Remove duplicates in input
if file not in related_files_set:
related_files.append(file)
related_files_set.add(file)
added_files.append(file)
results.append(f"Added related file: {file}")
# Rich output - single consolidated panel
if added_files:
files_added_md = '\n'.join(f"- `{file}`" for file in added_files)
md_content = f"**Files Noted:**\n{files_added_md}"
console.print(Panel(Markdown(md_content),
title="📁 Related Files Noted",
border_style="green"))
return results
class RunProgrammingTaskInput(BaseModel):
instructions: str = Field(description="Instructions for the programming task")
@ -88,13 +52,8 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str,
command.append(input.instructions)
# Use both input files and related files
files_to_use = set(related_files) # Start with related files
if input.files: # Add any additional input files
files_to_use.update(input.files)
if files_to_use:
command.extend(list(files_to_use))
if input.files:
command.extend(input.files)
# Create a pretty display of what we're doing
task_display = [
@ -102,10 +61,10 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str,
f"{input.instructions}\n"
]
if files_to_use:
if input.files:
task_display.extend([
"\n## Files\n",
*[f"- `{file}`\n" for file in files_to_use]
*[f"- `{file}`\n" for file in input.files]
])
markdown_content = "".join(task_display)
@ -138,4 +97,4 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str,
}
# Export the functions
__all__ = ['run_programming_task', 'emit_related_files']
__all__ = ['run_programming_task']

View File

@ -6,7 +6,9 @@ from ra_aid.tools.memory import (
emit_key_facts,
delete_key_facts,
emit_key_snippets,
delete_key_snippets
delete_key_snippets,
emit_related_files,
get_related_files
)
@pytest.fixture
@ -20,6 +22,7 @@ def reset_memory():
_global_memory['plans'] = []
_global_memory['tasks'] = []
_global_memory['research_subtasks'] = []
_global_memory['related_files'] = set()
yield
# Clean up after test
_global_memory['key_facts'] = {}
@ -207,6 +210,35 @@ def test_delete_key_snippets_empty(reset_memory):
# Verify snippet still exists
assert 0 in _global_memory['key_snippets']
def test_emit_related_files_basic(reset_memory):
"""Test basic adding of files"""
# Test adding single file
result = emit_related_files.invoke({"files": ["test.py"]})
assert result == "Files noted."
assert get_related_files() == {"test.py"}
# Test adding multiple files
result = emit_related_files.invoke({"files": ["main.py", "utils.py"]})
assert result == "Files noted."
assert get_related_files() == {"test.py", "main.py", "utils.py"}
def test_get_related_files_empty(reset_memory):
"""Test getting related files when none added"""
assert get_related_files() == set()
def test_emit_related_files_duplicates(reset_memory):
"""Test that duplicate files are handled correctly"""
# Add initial files
result = emit_related_files.invoke({"files": ["test.py", "main.py"]})
assert result == "Files noted."
assert get_related_files() == {"test.py", "main.py"}
# Try adding duplicates
result = emit_related_files.invoke({"files": ["test.py", "main.py", "test.py"]})
assert result == "Files noted."
# Set should still only contain unique entries
assert get_related_files() == {"test.py", "main.py"}
def test_key_snippets_integration(reset_memory):
"""Integration test for key snippets functionality"""
# Initial snippets to add