feat(mma): Implement context injection using ASTParser in run_worker_lifecycle
This commit is contained in:
@@ -1,5 +1,7 @@
|
|||||||
import ai_client
|
import ai_client
|
||||||
|
from typing import List, Optional
|
||||||
from models import Ticket, Track, WorkerContext
|
from models import Ticket, Track, WorkerContext
|
||||||
|
from file_cache import ASTParser
|
||||||
|
|
||||||
class ConductorEngine:
|
class ConductorEngine:
|
||||||
"""
|
"""
|
||||||
@@ -34,15 +36,35 @@ class ConductorEngine:
|
|||||||
)
|
)
|
||||||
run_worker_lifecycle(ticket, context)
|
run_worker_lifecycle(ticket, context)
|
||||||
|
|
||||||
def run_worker_lifecycle(ticket: Ticket, context: WorkerContext):
|
def run_worker_lifecycle(ticket: Ticket, context: WorkerContext, context_files: List[str] = None):
|
||||||
"""
|
"""
|
||||||
Simulates the lifecycle of a single agent working on a ticket.
|
Simulates the lifecycle of a single agent working on a ticket.
|
||||||
Calls the AI client and updates the ticket status based on the response.
|
Calls the AI client and updates the ticket status based on the response.
|
||||||
"""
|
"""
|
||||||
|
context_injection = ""
|
||||||
|
if context_files:
|
||||||
|
parser = ASTParser(language="python")
|
||||||
|
for i, file_path in enumerate(context_files):
|
||||||
|
try:
|
||||||
|
with open(file_path, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
if i == 0:
|
||||||
|
view = parser.get_curated_view(content)
|
||||||
|
else:
|
||||||
|
view = parser.get_skeleton(content)
|
||||||
|
context_injection += f"\nFile: {file_path}\n{view}\n"
|
||||||
|
except Exception as e:
|
||||||
|
context_injection += f"\nError reading {file_path}: {e}\n"
|
||||||
|
|
||||||
# Build a prompt for the worker
|
# Build a prompt for the worker
|
||||||
user_message = (
|
user_message = (
|
||||||
f"You are assigned to Ticket {ticket.id}.\n"
|
f"You are assigned to Ticket {ticket.id}.\n"
|
||||||
f"Task Description: {ticket.description}\n"
|
f"Task Description: {ticket.description}\n"
|
||||||
|
)
|
||||||
|
if context_injection:
|
||||||
|
user_message += f"\nContext Files:\n{context_injection}\n"
|
||||||
|
|
||||||
|
user_message += (
|
||||||
"Please complete this task. If you are blocked and cannot proceed, "
|
"Please complete this task. If you are blocked and cannot proceed, "
|
||||||
"start your response with 'BLOCKED' and explain why."
|
"start your response with 'BLOCKED' and explain why."
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -69,6 +69,58 @@ def test_run_worker_lifecycle_calls_ai_client_send():
|
|||||||
# user_message is passed as a keyword argument
|
# user_message is passed as a keyword argument
|
||||||
assert ticket.description in kwargs["user_message"]
|
assert ticket.description in kwargs["user_message"]
|
||||||
|
|
||||||
|
def test_run_worker_lifecycle_context_injection():
|
||||||
|
"""
|
||||||
|
Test that run_worker_lifecycle can take a context_files list and injects AST views into the prompt.
|
||||||
|
"""
|
||||||
|
ticket = Ticket(id="T1", description="Task 1", status="todo", assigned_to="worker1")
|
||||||
|
context = WorkerContext(ticket_id="T1", model_name="test-model", messages=[])
|
||||||
|
context_files = ["primary.py", "secondary.py"]
|
||||||
|
|
||||||
|
from multi_agent_conductor import run_worker_lifecycle
|
||||||
|
|
||||||
|
# We mock ASTParser which is expected to be imported in multi_agent_conductor
|
||||||
|
with patch("ai_client.send") as mock_send, \
|
||||||
|
patch("multi_agent_conductor.ASTParser") as mock_ast_parser_class, \
|
||||||
|
patch("builtins.open", new_callable=MagicMock) as mock_open:
|
||||||
|
|
||||||
|
# Setup open mock to return different content for different files
|
||||||
|
file_contents = {
|
||||||
|
"primary.py": "def primary(): pass",
|
||||||
|
"secondary.py": "def secondary(): pass"
|
||||||
|
}
|
||||||
|
|
||||||
|
def mock_open_side_effect(file, *args, **kwargs):
|
||||||
|
content = file_contents.get(file, "")
|
||||||
|
mock_file = MagicMock()
|
||||||
|
mock_file.read.return_value = content
|
||||||
|
mock_file.__enter__.return_value = mock_file
|
||||||
|
return mock_file
|
||||||
|
|
||||||
|
mock_open.side_effect = mock_open_side_effect
|
||||||
|
|
||||||
|
# Setup ASTParser mock
|
||||||
|
mock_ast_parser = mock_ast_parser_class.return_value
|
||||||
|
mock_ast_parser.get_curated_view.return_value = "CURATED VIEW"
|
||||||
|
mock_ast_parser.get_skeleton.return_value = "SKELETON VIEW"
|
||||||
|
|
||||||
|
mock_send.return_value = "Success"
|
||||||
|
|
||||||
|
run_worker_lifecycle(ticket, context, context_files=context_files)
|
||||||
|
|
||||||
|
# Verify ASTParser calls:
|
||||||
|
# First file (primary) should get curated view, others (secondary) get skeleton
|
||||||
|
mock_ast_parser.get_curated_view.assert_called_once_with("def primary(): pass")
|
||||||
|
mock_ast_parser.get_skeleton.assert_called_once_with("def secondary(): pass")
|
||||||
|
|
||||||
|
# Verify user_message contains the views
|
||||||
|
_, kwargs = mock_send.call_args
|
||||||
|
user_message = kwargs["user_message"]
|
||||||
|
assert "CURATED VIEW" in user_message
|
||||||
|
assert "SKELETON VIEW" in user_message
|
||||||
|
assert "primary.py" in user_message
|
||||||
|
assert "secondary.py" in user_message
|
||||||
|
|
||||||
def test_run_worker_lifecycle_handles_blocked_response():
|
def test_run_worker_lifecycle_handles_blocked_response():
|
||||||
"""
|
"""
|
||||||
Test that run_worker_lifecycle marks the ticket as blocked if the AI indicates it cannot proceed.
|
Test that run_worker_lifecycle marks the ticket as blocked if the AI indicates it cannot proceed.
|
||||||
|
|||||||
Reference in New Issue
Block a user