track related files through memory
This commit is contained in:
parent
14abec9735
commit
33f5f9bbbb
|
|
@ -1,5 +1,5 @@
|
||||||
from .shell import run_shell_command
|
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 .expert import ask_expert, emit_expert_context
|
||||||
from .read_file import read_file_tool
|
from .read_file import read_file_tool
|
||||||
from .fuzzy_find import fuzzy_find_project_files
|
from .fuzzy_find import fuzzy_find_project_files
|
||||||
|
|
@ -9,7 +9,7 @@ from .note_tech_debt import note_tech_debt
|
||||||
from .memory import (
|
from .memory import (
|
||||||
emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
|
emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
|
||||||
request_implementation, skip_implementation, delete_key_facts, emit_research_subtask,
|
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__ = [
|
__all__ = [
|
||||||
|
|
@ -32,7 +32,6 @@ __all__ = [
|
||||||
'run_shell_command',
|
'run_shell_command',
|
||||||
'skip_implementation',
|
'skip_implementation',
|
||||||
'emit_research_subtask',
|
'emit_research_subtask',
|
||||||
'fuzzy_find_project_files',
|
|
||||||
'ripgrep_search',
|
'ripgrep_search',
|
||||||
'note_tech_debt'
|
'note_tech_debt'
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -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 ra_aid.exceptions import TaskCompletedException
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
from rich.markdown import Markdown
|
from rich.markdown import Markdown
|
||||||
|
|
@ -15,7 +15,7 @@ class SnippetInfo(TypedDict):
|
||||||
console = Console()
|
console = Console()
|
||||||
|
|
||||||
# Global memory store
|
# 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': [],
|
'research_notes': [],
|
||||||
'plans': [],
|
'plans': [],
|
||||||
'tasks': [],
|
'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_snippets': {}, # Dict[int, SnippetInfo] - ID to snippet mapping
|
||||||
'key_snippet_id_counter': 0, # Counter for generating unique snippet IDs
|
'key_snippet_id_counter': 0, # Counter for generating unique snippet IDs
|
||||||
'implementation_requested': [],
|
'implementation_requested': [],
|
||||||
'implementation_skipped': []
|
'implementation_skipped': [],
|
||||||
|
'related_files': set()
|
||||||
}
|
}
|
||||||
|
|
||||||
@tool("emit_research_notes")
|
@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 ValueError("Cannot complete in one shot - implementation was requested")
|
||||||
raise TaskCompletedException(message)
|
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:
|
def get_memory_value(key: str) -> str:
|
||||||
"""Get a value from global memory.
|
"""Get a value from global memory.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -8,46 +8,10 @@ from rich.markdown import Markdown
|
||||||
from rich.text import Text
|
from rich.text import Text
|
||||||
from ra_aid.proc.interactive import run_interactive_command
|
from ra_aid.proc.interactive import run_interactive_command
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
from .memory import get_memory_value
|
|
||||||
from ra_aid.text.processing import truncate_output
|
from ra_aid.text.processing import truncate_output
|
||||||
|
|
||||||
console = Console()
|
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):
|
class RunProgrammingTaskInput(BaseModel):
|
||||||
instructions: str = Field(description="Instructions for the programming task")
|
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)
|
command.append(input.instructions)
|
||||||
|
|
||||||
# Use both input files and related files
|
if input.files:
|
||||||
files_to_use = set(related_files) # Start with related files
|
command.extend(input.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))
|
|
||||||
|
|
||||||
# Create a pretty display of what we're doing
|
# Create a pretty display of what we're doing
|
||||||
task_display = [
|
task_display = [
|
||||||
|
|
@ -102,10 +61,10 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str,
|
||||||
f"{input.instructions}\n"
|
f"{input.instructions}\n"
|
||||||
]
|
]
|
||||||
|
|
||||||
if files_to_use:
|
if input.files:
|
||||||
task_display.extend([
|
task_display.extend([
|
||||||
"\n## Files\n",
|
"\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)
|
markdown_content = "".join(task_display)
|
||||||
|
|
@ -138,4 +97,4 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str,
|
||||||
}
|
}
|
||||||
|
|
||||||
# Export the functions
|
# Export the functions
|
||||||
__all__ = ['run_programming_task', 'emit_related_files']
|
__all__ = ['run_programming_task']
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,9 @@ from ra_aid.tools.memory import (
|
||||||
emit_key_facts,
|
emit_key_facts,
|
||||||
delete_key_facts,
|
delete_key_facts,
|
||||||
emit_key_snippets,
|
emit_key_snippets,
|
||||||
delete_key_snippets
|
delete_key_snippets,
|
||||||
|
emit_related_files,
|
||||||
|
get_related_files
|
||||||
)
|
)
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
@ -20,6 +22,7 @@ def reset_memory():
|
||||||
_global_memory['plans'] = []
|
_global_memory['plans'] = []
|
||||||
_global_memory['tasks'] = []
|
_global_memory['tasks'] = []
|
||||||
_global_memory['research_subtasks'] = []
|
_global_memory['research_subtasks'] = []
|
||||||
|
_global_memory['related_files'] = set()
|
||||||
yield
|
yield
|
||||||
# Clean up after test
|
# Clean up after test
|
||||||
_global_memory['key_facts'] = {}
|
_global_memory['key_facts'] = {}
|
||||||
|
|
@ -207,6 +210,35 @@ def test_delete_key_snippets_empty(reset_memory):
|
||||||
# Verify snippet still exists
|
# Verify snippet still exists
|
||||||
assert 0 in _global_memory['key_snippets']
|
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):
|
def test_key_snippets_integration(reset_memory):
|
||||||
"""Integration test for key snippets functionality"""
|
"""Integration test for key snippets functionality"""
|
||||||
# Initial snippets to add
|
# Initial snippets to add
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue