test_local_agent.py

#!/usr/bin/env python3
"""
Local Agent Test - End-to-end testing with Mock LLM and test server
Validates the complete ReAct pipeline without GPU resources
"""

import os
import sys
import json
import time
import subprocess
import signal
import logging
from pathlib import Path
from typing import Dict, List

# Add project root to path
sys.path.insert(0, os.path.dirname(__file__))

from agent.mock_llm import get_mock_llm
from agent.react_loop import ReActAgent, AgentState
from agent.tools import ToolExecutor
from agent.trace_logger import TraceLogger
from utils.format_validator import get_format_validator

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class LocalTestEnvironment:
    """
    Local testing environment for the ReAct agent
    """
    
    def __init__(self, server_url: str = "http://localhost:8000"):
        self.server_url = server_url
        self.server_process = None
        self.mock_llm = get_mock_llm()
        self.tool_executor = ToolExecutor(server_url)
        self.trace_logger = TraceLogger("test_traces.json")
        self.validator = get_format_validator()
        
    def start_test_server(self) -> bool:
        """Start the test server in background"""
        try:
            # Start the test server
            self.server_process = subprocess.Popen(
                [sys.executable, "test_server.py"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
            
            # Wait for server to start
            time.sleep(2)
            
            # Check if server is running
            import requests
            response = requests.get(f"{self.server_url}/health", timeout=5)
            
            if response.status_code == 200:
                logger.info("Test server started successfully")
                return True
            else:
                logger.error("Test server health check failed")
                return False
                
        except Exception as e:
            logger.error(f"Failed to start test server: {e}")
            return False
    
    def stop_test_server(self):
        """Stop the test server"""
        if self.server_process:
            self.server_process.terminate()
            self.server_process.wait()
            logger.info("Test server stopped")
    
    def create_test_scenarios(self) -> List[Dict]:
        """Create test scenarios for local testing"""
        return [
            {
                "id": "test_001",
                "description": "Network connectivity issue between Node-A and Node-B",
                "symptoms": ["ping fails", "high latency"],
                "environment": "production",
                "scenario_type": "connectivity_issue"
            },
            {
                "id": "test_002", 
                "description": "High CPU usage on Node-C affecting performance",
                "symptoms": ["slow response", "high latency"],
                "environment": "staging",
                "scenario_type": "performance_issue"
            },
            {
                "id": "test_003",
                "description": "Web service not responding on Node-D",
                "symptoms": ["service down", "connection refused"],
                "environment": "development",
                "scenario_type": "service_down"
            },
            {
                "id": "test_004",
                "description": "General network diagnosis required",
                "symptoms": ["intermittent issues", "unknown cause"],
                "environment": "production",
                "scenario_type": "general_diagnosis"
            }
        ]
    
    def test_http_tools(self) -> bool:
        """Test HTTP tool-calling functions against local server"""
        logger.info("Testing HTTP tools...")
        
        test_cases = [
            {
                "tool": "get_node_status",
                "params": {"node_id": "Node-A"},
                "expected_status": "success"
            },
            {
                "tool": "check_connectivity",
                "params": {"source_node": "Node-A", "destination_node": "Node-B"},
                "expected_status": "success"
            },
            {
                "tool": "execute_command",
                "params": {"node_id": "Node-A", "command": "ping 8.8.8.8"},
                "expected_status": "success"
            },
            {
                "tool": "get_logs",
                "params": {"node_id": "Node-A", "log_type": "syslog", "lines": 10},
                "expected_status": "success"
            },
            {
                "tool": "ping_test",
                "params": {"target": "8.8.8.8", "count": 3},
                "expected_status": "success"
            }
        ]
        
        for test_case in test_cases:
            tool = test_case["tool"]
            params = test_case["params"]
            expected = test_case["expected_status"]
            
            try:
                result = self.tool_executor.execute_tool(tool, params)
                actual = result.get("status")
                
                # For successful responses, status should be "success"
                if actual == expected:
                    logger.info(f"  {tool}: PASS")
                else:
                    logger.error(f"  {tool}: FAIL - Expected {expected}, got {actual}")
                    return False
                    
            except Exception as e:
                logger.error(f"  {tool}: ERROR - {e}")
                return False
        
        logger.info("All HTTP tools tests passed")
        return True
    
    def test_mock_llm_responses(self) -> bool:
        """Test Mock LLM response generation"""
        logger.info("Testing Mock LLM responses...")
        
        test_prompts = [
            ("Network connectivity issue", "connectivity_issue"),
            ("High latency problem", "performance_issue"),
            ("Service is down", "service_down"),
            ("General diagnosis", "general_diagnosis")
        ]
        
        for prompt, expected_type in test_prompts:
            response = self.mock_llm.generate_json_response(prompt, expected_type)
            
            # Check response structure
            required_keys = ["thought", "action", "action_input"]
            if not all(key in response for key in required_keys):
                logger.error(f"Mock LLM missing required keys for: {prompt}")
                return False
            
            # Check action is valid
            valid_actions = ["get_node_status", "check_connectivity", "execute_command", 
                           "get_logs", "analyze_metrics", "get_topology", "restart_service", "ping_test"]
            if response["action"] not in valid_actions:
                logger.error(f"Mock LLM invalid action: {response['action']}")
                return False
            
            logger.info(f"  {prompt}: {response['action']} - {response['thought'][:30]}...")
        
        logger.info("Mock LLM responses test passed")
        return True
    
    def run_react_test(self, scenario: Dict) -> Dict:
        """Run ReAct agent test with Mock LLM"""
        logger.info(f"Running ReAct test for: {scenario['description']}")
        
        # Initialize agent state
        state = AgentState(
            conversation_history=[],
            current_scenario=scenario,
            tool_results={},
            error_count=0,
            max_errors=3
        )
        
        start_time = time.time()
        iterations = 0
        max_iterations = 5  # Limited for testing
        
        # Build initial prompt
        prompt = f"""You are an expert telecom network troubleshooting agent. Your task is to diagnose and resolve network issues using the available tools.

You must follow this exact format for every response:
{{
    "thought": "Your reasoning about the current situation and what you need to do next",
    "action": "The specific tool to execute",
    "action_input": {{"parameter": "value"}}
}}

Available tools:
- get_node_status: Get status of a specific network node
- execute_command: Execute CLI commands on network devices
- check_connectivity: Test connectivity between nodes
- get_logs: Retrieve system logs
- analyze_metrics: Analyze network performance metrics
- get_topology: Get network topology information
- restart_service: Restart a network service
- ping_test: Perform ping test

Current scenario: {json.dumps(scenario, indent=2)}

What should I do first?"""
        
        while iterations < max_iterations and state.error_count < state.max_errors:
            # Generate response with Mock LLM
            response = self.mock_llm.generate_json_response(prompt, scenario.get("scenario_type"))
            
            # Extract components
            thought = response.get("thought", "")
            action = response.get("action", "")
            action_input = response.get("action_input", {})
            
            # Execute action
            if action != "think" and action != "error":
                tool_start = time.time()
                observation = self.tool_executor.execute_tool(action, action_input)
                execution_time = time.time() - tool_start
                
                # Log trace
                self.trace_logger.add_trace_entry(
                    iterations, thought, action, action_input, observation, execution_time
                )
                
                # Store result
                state.tool_results[action] = observation
                
                # Check for errors
                if observation.get("status") == "error":
                    state.error_count += 1
                    logger.warning(f"Action failed: {action} - {observation.get('error')}")
                else:
                    logger.info(f"Action success: {action}")
            else:
                observation = {"status": "success", "message": "Thought completed"}
                execution_time = 0.1
            
            # Add to conversation history
            state.conversation_history.append({
                "iteration": iterations,
                "thought": thought,
                "action": action,
                "action_input": action_input,
                "observation": observation
            })
            
            # Update prompt for next iteration
            prompt += f"\n\nThought: {thought}\nAction: {action}\nAction Input: {json.dumps(action_input)}\nObservation: {json.dumps(observation)}\n\nWhat should I do next?"
            
            iterations += 1
        
        total_time = time.time() - start_time
        
        # Extract solution from conversation
        solution = self._extract_solution(state.conversation_history)
        
        return {
            "scenario": scenario,
            "iterations": iterations,
            "execution_time": total_time,
            "error_count": state.error_count,
            "conversation_history": state.conversation_history,
            "tool_results": state.tool_results,
            "final_state": {
                "completed": state.error_count < state.max_errors,
                "solution": solution
            }
        }
    
    def _extract_solution(self, history: List[Dict]) -> Dict:
        """Extract solution from conversation history"""
        if not history:
            return {}
        
        # Get last successful action
        for entry in reversed(history):
            obs = entry.get("observation", {})
            if obs.get("status") == "success":
                return {
                    "track_a_answer": self._extract_track_a(entry),
                    "track_b_answer": self._extract_track_b(entry)
                }
        
        return {}
    
    def _extract_track_a(self, entry: Dict) -> str:
        """Extract Track A answer from entry"""
        # Look for cell patterns in thought
        import re
        cells = re.findall(r'Node-\w+', entry.get("thought", ""))
        if cells:
            # Convert to cell format (simplified)
            return "|".join([f"C{i+1}" for i in range(len(cells))])
        return "C1|C2"  # Default
    
    def _extract_track_b(self, entry: Dict) -> str:
        """Extract Track B answer from entry"""
        thought = entry.get("thought", "")
        action = entry.get("action", "")
        
        # Create diagnosis string
        diagnosis = f"{action};{thought[:100]}"
        return diagnosis
    
    def run_end_to_end_test(self) -> Dict:
        """Run complete end-to-end test"""
        logger.info("Starting end-to-end test...")
        
        # Test scenarios
        scenarios = self.create_test_scenarios()
        results = []
        
        for scenario in scenarios:
            # Reset Mock LLM for each scenario
            self.mock_llm.reset()
            
            # Run ReAct test
            result = self.run_react_test(scenario)
            result["scenario_index"] = scenarios.index(scenario)
            results.append(result)
            
            logger.info(f"Scenario {scenario['id']}: {result['iterations']} iterations, {result['execution_time']:.2f}s")
        
        # Generate submission file
        submission_data = {
            "results": results
        }
        
        submission_success = self.generate_submission(submission_data, "test_result.csv")
        
        # Validate traces
        traces_valid, traces_result = self.validator.validate_traces_format("test_traces.json")
        
        # Validate submission
        csv_valid, csv_result = self.validator.validate_csv_format("test_result.csv")
        
        return {
            "total_scenarios": len(scenarios),
            "total_execution_time": sum(r["execution_time"] for r in results),
            "results": results,
            "submission_generated": submission_success,
            "traces_valid": traces_valid,
            "csv_valid": csv_valid,
            "overall_success": traces_valid and csv_valid and submission_success
        }
    
    def generate_submission(self, results: Dict, output_file: str) -> bool:
        """Generate submission file from test results"""
        try:
            import pandas as pd
            
            submission_data = []
            
            for result in results.get("results", []):
                scenario = result.get("scenario", {})
                solution = result.get("final_state", {}).get("solution", {})
                
                submission_data.append({
                    "ID": scenario.get("id", ""),
                    "Track A": solution.get("track_a_answer", ""),
                    "Track B": solution.get("track_b_answer", "")
                })
            
            # Create DataFrame and save
            df = pd.DataFrame(submission_data, columns=["ID", "Track A", "Track B"])
            df = df.fillna("")
            df.to_csv(output_file, index=False)
            
            logger.info(f"Test submission generated: {output_file}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to generate test submission: {e}")
            return False
    
    def cleanup(self):
        """Clean up test environment"""
        self.stop_test_server()
        
        # Clean up test files
        test_files = ["test_result.csv", "test_traces.json"]
        for file in test_files:
            if os.path.exists(file):
                os.remove(file)
                logger.info(f"Cleaned up: {file}")


def main():
    """Main test execution"""
    print("=" * 60)
    print("LOCAL AGENT TEST SUITE")
    print("=" * 60)
    
    test_env = LocalTestEnvironment()
    
    try:
        # Start test server
        if not test_env.start_test_server():
            print("FAILED: Could not start test server")
            return False
        
        # Test HTTP tools
        if not test_env.test_http_tools():
            print("FAILED: HTTP tools test")
            return False
        
        # Test Mock LLM
        if not test_env.test_mock_llm_responses():
            print("FAILED: Mock LLM test")
            return False
        
        # Run end-to-end test
        result = test_env.run_end_to_end_test()
        
        # Print results
        print("\n" + "=" * 60)
        print("TEST RESULTS")
        print("=" * 60)
        print(f"Total scenarios: {result['total_scenarios']}")
        print(f"Total execution time: {result['total_execution_time']:.2f}s")
        print(f"Submission generated: {'YES' if result['submission_generated'] else 'NO'}")
        print(f"Traces valid: {'YES' if result['traces_valid'] else 'NO'}")
        print(f"CSV valid: {'YES' if result['csv_valid'] else 'NO'}")
        print(f"Overall success: {'YES' if result['overall_success'] else 'NO'}")
        
        if result['overall_success']:
            print("\nSUCCESS: All tests passed!")
            print("The local agent is ready for Kaggle deployment.")
        else:
            print("\nFAILED: Some tests failed.")
            print("Please fix issues before deploying to Kaggle.")
        
        return result['overall_success']
        
    except KeyboardInterrupt:
        print("\nTest interrupted by user")
        return False
    except Exception as e:
        print(f"\nTest failed with error: {e}")
        return False
    finally:
        test_env.cleanup()


if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)