Nexa_Labs / examples /test_agent.py
Allanatrix's picture
Upload 57 files
d8328bf verified
"""Simple CLI to exercise the NexaSci agent and tool server."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Dict, Sequence
# Add project root to Python path
project_root = Path(__file__).resolve().parents[1]
if str(project_root) not in sys.path:
sys.path.insert(0, str(project_root))
from agent.controller import AgentController
SAMPLE_PROMPTS: Dict[str, str] = {
"experiment": "Design a reproducible procedure to measure battery degradation in solid-state lithium cells over 200 charge cycles.",
"simulation": "Use the python sandbox to simulate the trajectory of a projectile launched at 45 degrees with air resistance.",
"literature": "Summarise recent advances in room-temperature superconductivity and cite key papers.",
}
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
"""Parse command line arguments."""
parser = argparse.ArgumentParser(description="Run sample prompts through the NexaSci agent.")
parser.add_argument(
"--prompt",
choices=SAMPLE_PROMPTS.keys(),
help="Run a single named sample prompt.",
)
parser.add_argument(
"--raw",
help="Run a custom prompt provided on the command line.",
)
parser.add_argument(
"--max-turns",
type=int,
default=8,
help="Maximum LLM/tool turns before aborting (default: %(default)s).",
)
return parser.parse_args(argv)
def run_prompt(controller: AgentController, label: str, prompt: str) -> None:
"""Execute the agent for a single prompt and print a compact report."""
print("=" * 80)
print(f"Prompt: {label}")
print("-" * 80)
result = controller.run(prompt)
print("Final Response:")
print(result.pretty())
if result.tool_results:
print("\nTool Trace:")
print(json.dumps([tool.output for tool in result.tool_results], indent=2))
print("=" * 80)
def main() -> None:
"""Entry point for the CLI utility."""
args = parse_args()
controller = AgentController(max_turns=args.max_turns)
if args.raw:
run_prompt(controller, "custom", args.raw)
return
if args.prompt:
run_prompt(controller, args.prompt, SAMPLE_PROMPTS[args.prompt])
return
for label, prompt in SAMPLE_PROMPTS.items():
run_prompt(controller, label, prompt)
if __name__ == "__main__": # pragma: no cover - CLI helper
main()