diff --git a/scripts/mma_exec.py b/scripts/mma_exec.py index d1ece92..61c4ba9 100644 --- a/scripts/mma_exec.py +++ b/scripts/mma_exec.py @@ -1,5 +1,6 @@ import argparse import subprocess +import json def get_role_documents(role: str) -> list[str]: if role == 'tier1': @@ -15,9 +16,24 @@ def execute_agent(role: str, prompt: str, docs: list[str]) -> str: for doc in docs: command_text += f" @{doc}" - cmd = ["gemini", command_text] - process = subprocess.run(cmd, capture_output=True, text=True) - return process.stdout + cmd = ['gemini', '-p', command_text, '--output-format', 'json', '--model', 'gemini-2.5-flash-lite'] + try: + process = subprocess.run(cmd, capture_output=True, text=True, shell=True) + if not process.stdout and process.stderr: + return f"Error: {process.stderr}" + + stdout = process.stdout + start_index = stdout.find('{') + if start_index != -1: + json_str = stdout[start_index:] + try: + data = json.loads(json_str) + return data.get('response', stdout) + except json.JSONDecodeError: + return stdout + return stdout + except Exception as e: + return f"Execution failed: {str(e)}" def create_parser(): parser = argparse.ArgumentParser(description="MMA Execution Script") diff --git a/tests/test_mma_exec.py b/tests/test_mma_exec.py index 15e1ae1..56964e6 100644 --- a/tests/test_mma_exec.py +++ b/tests/test_mma_exec.py @@ -64,7 +64,8 @@ def test_execute_agent(): cmd_list = args[0] assert cmd_list[0] == "gemini" - assert cmd_list[1] == expected_gemini_arg + assert cmd_list[1] == "-p" + assert cmd_list[2] == expected_gemini_arg assert kwargs.get("capture_output") is True assert kwargs.get("text") is True