test_agent_tools.py

"""
Test Agent Loop Tool Calling Logic
Tests tool determination and execution with different LLMs and inputs
Measures accuracy and performance metrics
"""
import asyncio
import logging
import time
import json
from typing import Dict, List, Any
from pathlib import Path
from datetime import datetime
from config import get_config
from ollama_client import OllamaClient
from tool_processor import ToolProcessor
from system_prompt import SystemPrompt

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class AgentTestResult:
    """Container for test result data"""
    def __init__(self, test_input: str, model: str):
        self.test_input = test_input
        self.model = model
        self.start_time = None
        self.end_time = None
        self.duration = None
        self.result = None
        self.tool_called = None
        self.expected_tool = None
        self.success = False
        self.error = None
        self.iterations = 0

    def start(self):
        """Mark test start time"""
        self.start_time = time.time()

    def end(self):
        """Mark test end time and calculate duration"""
        self.end_time = time.time()
        self.duration = self.end_time - self.start_time

    def set_result(self, result: str, tool_called: str, iterations: int):
        """Set test result"""
        self.result = result
        self.tool_called = tool_called
        self.iterations = iterations

    def set_expected_tool(self, tool: str):
        """Set expected tool for validation"""
        self.expected_tool = tool

    def check_success(self):
        """Check if correct tool was called"""
        if self.tool_called and self.expected_tool:
            self.success = self.tool_called.lower() == self.expected_tool.lower()
        return self.success

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for reporting"""
        return {
            'test_input': self.test_input,
            'model': self.model,
            'duration': round(self.duration, 3) if self.duration else None,
            'tool_called': self.tool_called,
            'expected_tool': self.expected_tool,
            'success': self.success,
            'iterations': self.iterations,
            'result': self.result[:100] + '...' if self.result and len(self.result) > 100 else self.result,
            'error': self.error
        }


class AgentToolTester:
    """Test harness for agent tool calling"""

    def __init__(self):
        self.config = get_config()
        self.ollama_client = OllamaClient(base_url=self.config.OLLAMA_URL)
        self.tool_processor = ToolProcessor()
        self.system_prompts = SystemPrompt(self.config)
        self.results: List[AgentTestResult] = []

        # Test configuration - using short commands for speed
        self.test_inputs = [
            "timenow",
            "fortune"
        ]

        self.test_models = [
            "exaone3.5:2.4b",
            "qwen3:1.7b",
            "gemma3:1b",
            "granite3.1-moe:1b",
            "granite3.3:2b",
            "ingu627/exaone4.0:1.2b"
        ]

        # Expected tool mappings
        self.expected_tools = {
            "timenow": "get_current_time",
            "fortune": "get_fortune",
            "what time is it": "get_current_time",
            "tell me my fortune": "get_fortune",
            "give me a fortune cookie": "get_fortune",
            "what is the current time": "get_current_time",
            "show me the time": "get_current_time",
            "get me a random fortune": "get_fortune"
        }

        logger.info("Agent Tool Tester initialized")
        logger.info(f"Available tools: {list(self.tool_processor.tools.keys())}")

    async def test_single_input_model(self, test_input: str, model: str, run_number: int = 1) -> AgentTestResult:
        """Test a single input with a single model"""
        result = AgentTestResult(test_input, model)
        result.set_expected_tool(self.expected_tools.get(test_input.lower(), "unknown"))
        result.start()

        logger.info(f"\n{'='*80}")
        logger.info(f"Testing: '{test_input}' with model: {model} (Run {run_number}/3)")
        logger.info(f"Expected tool: {result.expected_tool}")
        logger.info(f"{'='*80}")

        try:
            # Get agent system prompt
            agent_config = self.system_prompts.agent
            system_prompt = agent_config['prompt_text']

            # Run agent loop
            agent_result = await self.tool_processor.agent_loop(
                test_input,
                self.ollama_client,
                system_prompt,
                max_iterations=self.config.AGENT_MAX_ITERATIONS,
                model=model
            )

            result.end()

            # Parse result to extract tool information
            tool_called = self._extract_tool_from_result(agent_result)
            iterations = self._extract_iterations_from_result(agent_result)

            result.set_result(agent_result, tool_called, iterations)
            result.check_success()

            # Log result
            status = "✓ SUCCESS" if result.success else "✗ FAILED"
            logger.info(f"{status} - Tool: {tool_called}, Duration: {result.duration:.3f}s, Iterations: {iterations}")
            logger.info(f"Result: {agent_result[:200]}...")

        except Exception as e:
            result.end()
            result.error = str(e)
            logger.error(f"✗ ERROR: {e}")

        return result

    async def test_single_input_model_with_retries(self, test_input: str, model: str) -> AgentTestResult:
        """Test a single input with a model, running successful tests 3 times to get best timing"""

        # First attempt
        result = await self.test_single_input_model(test_input, model, run_number=1)

        # If first attempt failed or had error, return it (no retries for failures)
        if not result.success or result.error:
            logger.info(f"Test failed on first attempt - skipping retries")
            return result

        # Test succeeded - run 2 more times and keep the best (fastest) result
        logger.info(f"✓ Test succeeded! Running 2 more times to get best timing...")

        best_result = result

        for run_num in range(2, 4):  # runs 2 and 3
            await asyncio.sleep(0.5)  # Brief delay between runs

            retry_result = await self.test_single_input_model(test_input, model, run_number=run_num)

            # Only compare if retry also succeeded
            if retry_result.success and not retry_result.error:
                if retry_result.duration < best_result.duration:
                    logger.info(f"  Run {run_num}: {retry_result.duration:.3f}s - NEW BEST TIME!")
                    best_result = retry_result
                else:
                    logger.info(f"  Run {run_num}: {retry_result.duration:.3f}s")
            else:
                logger.warning(f"  Run {run_num}: Failed (inconsistent)")

        logger.info(f"🏆 Best time: {best_result.duration:.3f}s (from 3 runs)")
        return best_result

    def _extract_tool_from_result(self, result: str) -> str:
        """Extract tool name from agent result"""
        result_lower = result.lower()

        # Check for each known tool
        if "get_current_time" in result_lower or "current time" in result_lower:
            return "get_current_time"
        elif "get_fortune" in result_lower or "fortune" in result_lower:
            return "get_fortune"
        elif "no tasks executed" in result_lower:
            return "none"
        else:
            return "unknown"

    def _extract_iterations_from_result(self, result: str) -> int:
        """Extract iteration count from result (if available)"""
        # This is a heuristic - could be improved by having tool_processor return this info
        # For now, assume 1 iteration if tool was executed, 0 if not
        if "no tasks executed" in result.lower():
            return 0
        return 1

    async def run_all_tests(self):
        """Run all test combinations"""
        logger.info("\n" + "="*80)
        logger.info("STARTING AGENT TOOL CALLING TEST SUITE")
        logger.info("="*80)
        logger.info(f"Test inputs: {len(self.test_inputs)}")
        logger.info(f"Test models: {len(self.test_models)}")
        logger.info(f"Total tests: {len(self.test_inputs) * len(self.test_models)}")
        logger.info("Note: Successful tests will run 3 times, keeping best timing")
        logger.info("="*80 + "\n")

        start_time = time.time()

        # Run all test combinations
        for test_input in self.test_inputs:
            for model in self.test_models:
                # Use retry logic to run successful tests 3 times
                result = await self.test_single_input_model_with_retries(test_input, model)
                self.results.append(result)

                # Small delay between test groups
                await asyncio.sleep(1.0)

        total_duration = time.time() - start_time

        logger.info("\n" + "="*80)
        logger.info("ALL TESTS COMPLETE")
        logger.info(f"Total duration: {total_duration:.2f}s")
        logger.info("="*80 + "\n")

        return total_duration

    def generate_report(self) -> str:
        """Generate comprehensive test report"""
        report = []
        report.append("\n" + "="*100)
        report.append("AGENT TOOL CALLING TEST REPORT")
        report.append("="*100)
        report.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        report.append(f"Total tests: {len(self.results)}")
        report.append("="*100 + "\n")

        # Overall statistics
        total_tests = len(self.results)
        successful_tests = sum(1 for r in self.results if r.success)
        failed_tests = sum(1 for r in self.results if not r.success and not r.error)
        error_tests = sum(1 for r in self.results if r.error)

        report.append("OVERALL STATISTICS")
        report.append("-" * 100)
        report.append(f"Total tests:      {total_tests}")
        report.append(f"Successful:       {successful_tests} ({successful_tests/total_tests*100:.1f}%)")
        report.append(f"Failed:           {failed_tests} ({failed_tests/total_tests*100:.1f}%)")
        report.append(f"Errors:           {error_tests} ({error_tests/total_tests*100:.1f}%)")
        report.append("")

        # Model performance comparison
        report.append("MODEL PERFORMANCE COMPARISON")
        report.append("-" * 100)
        report.append(f"{'Model':<20} {'Accuracy':<12} {'Avg Duration':<15} {'Min Duration':<15} {'Max Duration':<15}")
        report.append("-" * 100)

        model_stats = {}
        for model in self.test_models:
            model_results = [r for r in self.results if r.model == model and not r.error]
            if model_results:
                accuracy = sum(1 for r in model_results if r.success) / len(model_results) * 100
                avg_duration = sum(r.duration for r in model_results) / len(model_results)
                min_duration = min(r.duration for r in model_results)
                max_duration = max(r.duration for r in model_results)

                model_stats[model] = {
                    'accuracy': accuracy,
                    'avg_duration': avg_duration,
                    'min_duration': min_duration,
                    'max_duration': max_duration
                }

                report.append(f"{model:<20} {accuracy:>6.1f}%      {avg_duration:>10.3f}s      {min_duration:>10.3f}s      {max_duration:>10.3f}s")

        report.append("")

        # Best performers
        if model_stats:
            best_accuracy_model = max(model_stats.items(), key=lambda x: x[1]['accuracy'])
            fastest_model = min(model_stats.items(), key=lambda x: x[1]['avg_duration'])

            report.append("BEST PERFORMERS")
            report.append("-" * 100)
            report.append(f"Highest Accuracy: {best_accuracy_model[0]} ({best_accuracy_model[1]['accuracy']:.1f}%)")
            report.append(f"Fastest Average:  {fastest_model[0]} ({fastest_model[1]['avg_duration']:.3f}s)")
            report.append("")

        # Tool accuracy breakdown
        report.append("TOOL ACCURACY BREAKDOWN")
        report.append("-" * 100)

        tools = set(r.expected_tool for r in self.results if r.expected_tool)
        for tool in sorted(tools):
            tool_results = [r for r in self.results if r.expected_tool == tool and not r.error]
            if tool_results:
                tool_success = sum(1 for r in tool_results if r.success)
                tool_total = len(tool_results)
                tool_accuracy = tool_success / tool_total * 100
                report.append(f"{tool:<20} {tool_success}/{tool_total} tests passed ({tool_accuracy:.1f}%)")

        report.append("")

        # Detailed results per input
        report.append("DETAILED RESULTS BY INPUT")
        report.append("-" * 100)

        for test_input in self.test_inputs:
            input_results = [r for r in self.results if r.test_input == test_input]
            if input_results:
                report.append(f"\nInput: '{test_input}'")
                report.append(f"Expected tool: {input_results[0].expected_tool}")
                report.append(f"{'Model':<20} {'Status':<10} {'Tool Called':<20} {'Duration':<12} {'Iterations'}")
                report.append("-" * 100)

                for result in input_results:
                    status = "✓ PASS" if result.success else ("ERROR" if result.error else "✗ FAIL")
                    tool = result.tool_called or "none"
                    duration = f"{result.duration:.3f}s" if result.duration else "N/A"
                    iterations = str(result.iterations) if result.iterations is not None else "N/A"

                    report.append(f"{result.model:<20} {status:<10} {tool:<20} {duration:<12} {iterations}")

        report.append("")
        report.append("="*100)
        report.append("END OF REPORT")
        report.append("="*100 + "\n")

        return "\n".join(report)

    def save_results(self, filename: str = "agent_test_results.json"):
        """Save detailed results to JSON file"""
        output_data = {
            'test_date': datetime.now().isoformat(),
            'test_inputs': self.test_inputs,
            'test_models': self.test_models,
            'results': [r.to_dict() for r in self.results]
        }

        output_file = Path(filename)
        with open(output_file, 'w') as f:
            json.dump(output_data, f, indent=2)

        logger.info(f"Detailed results saved to: {output_file}")
        return output_file


async def main():
    """Main test execution"""
    logger.info("Initializing Agent Tool Tester...")

    try:
        tester = AgentToolTester()

        # Run all tests
        total_duration = await tester.run_all_tests()

        # Generate and display report
        report = tester.generate_report()
        print(report)

        # Save detailed results
        results_file = tester.save_results()

        # Summary
        print(f"\n{'='*80}")
        print(f"Test suite completed in {total_duration:.2f}s")
        print(f"Detailed results saved to: {results_file}")
        print(f"{'='*80}\n")

    except Exception as e:
        logger.error(f"Fatal error in test suite: {e}", exc_info=True)
        raise


if __name__ == "__main__":
    """Run test suite"""
    asyncio.run(main())