#!/usr/bin/env python3 """ Atlas Functionality Test Suite Tests the chat endpoint and search decision waterfall logic """ import asyncio import json import requests from typing import Dict, Any, Optional, List import time # Test configuration BASE_URL = "http://localhost:7860" CHAT_ENDPOINT = f"{BASE_URL}/chat" HEALTH_ENDPOINT = f"{BASE_URL}/" class AtlasTestSuite: def __init__(self): self.results = [] self.session_id = None def log_result(self, test_name: str, success: bool, details: str = "", response_data: Dict = None): """Log test result""" result = { "test": test_name, "success": success, "details": details, "timestamp": time.time() } if response_data: result["response_data"] = response_data self.results.append(result) status = "āœ… PASS" if success else "āŒ FAIL" print(f"{status} {test_name}: {details}") def test_health_endpoint(self) -> bool: """Test the health endpoint""" try: response = requests.get(HEALTH_ENDPOINT, timeout=10) if response.status_code == 200: data = response.json() self.log_result("Health Check", True, "Server responding correctly", data) return True else: self.log_result("Health Check", False, f"Status code: {response.status_code}") return False except Exception as e: self.log_result("Health Check", False, f"Exception: {str(e)}") return False def test_anonymous_chat_no_search(self) -> bool: """Test anonymous chat request without search""" try: payload = { "prompt": "Hello, how are you?", "max_new_tokens": 100, "use_search": False, "temperature": 0.7 } response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30) if response.status_code == 200: data = response.json() # Check response structure if "response" in data and len(data["response"]) > 0: # Store session ID for follow-up tests self.session_id = response.headers.get("X-Session-ID") self.log_result("Anonymous Chat (No Search)", True, f"Response length: {len(data['response'])} chars, Session: {self.session_id}", data) return True else: self.log_result("Anonymous Chat (No Search)", False, "Empty or missing response") return False else: self.log_result("Anonymous Chat (No Search)", False, f"Status: {response.status_code}") return False except Exception as e: self.log_result("Anonymous Chat (No Search)", False, f"Exception: {str(e)}") return False def test_anonymous_chat_with_search(self) -> bool: """Test anonymous chat request with search enabled""" try: payload = { "prompt": "What is artificial intelligence?", "max_new_tokens": 150, "use_search": True, "temperature": 0.7 } response = requests.post(CHAT_ENDPOINT, json=payload, timeout=45) if response.status_code == 200: data = response.json() # Check response structure has_response = "response" in data and len(data["response"]) > 0 has_search_decision = "search_decision" in data has_cache_info = "cache_info" in data details = f"Response: {has_response}, Decision: {has_search_decision}, Cache: {has_cache_info}" if has_response: self.log_result("Anonymous Chat (With Search)", True, details, data) return True else: self.log_result("Anonymous Chat (With Search)", False, "Missing response") return False else: self.log_result("Anonymous Chat (With Search)", False, f"Status: {response.status_code}") return False except Exception as e: self.log_result("Anonymous Chat (With Search)", False, f"Exception: {str(e)}") return False def test_authenticated_chat(self) -> bool: """Test authenticated chat request""" try: payload = { "prompt": "Hello, I'm a test user. Can you help me?", "max_new_tokens": 100, "use_search": False, "temperature": 0.7, "user_id": "test_user_001" } response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30) if response.status_code == 200: data = response.json() if "response" in data and len(data["response"]) > 0: self.log_result("Authenticated Chat", True, f"Response length: {len(data['response'])} chars", data) return True else: self.log_result("Authenticated Chat", False, "Empty or missing response") return False else: self.log_result("Authenticated Chat", False, f"Status: {response.status_code}") return False except Exception as e: self.log_result("Authenticated Chat", False, f"Exception: {str(e)}") return False def test_conversation_follow_up(self) -> bool: """Test conversation follow-up to check search decision waterfall""" if not self.session_id: self.log_result("Conversation Follow-up", False, "No session ID available") return False try: # First establish conversation history payload1 = { "prompt": "What is machine learning?", "max_new_tokens": 100, "use_search": True, "temperature": 0.7 } headers = {"X-Session-ID": self.session_id} response1 = requests.post(CHAT_ENDPOINT, json=payload1, headers=headers, timeout=45) if response1.status_code != 200: self.log_result("Conversation Follow-up", False, f"First request failed: {response1.status_code}") return False # Now test follow-up question (should trigger different search logic) payload2 = { "prompt": "Can you tell me more about that?", "max_new_tokens": 100, "use_search": True, "temperature": 0.7, "history": [ {"role": "user", "content": "What is machine learning?"}, {"role": "assistant", "content": response1.json()["response"][:200]} ] } response2 = requests.post(CHAT_ENDPOINT, json=payload2, headers=headers, timeout=45) if response2.status_code == 200: data = response2.json() has_search_decision = "search_decision" in data decision_info = "" if has_search_decision: decision = data["search_decision"] decision_info = f"Should search: {decision.get('should_search')}, " \ f"Confidence: {decision.get('confidence')}, " \ f"Flow: {decision.get('flow_type', 'unknown')}" self.log_result("Conversation Follow-up", True, f"Follow-up successful. {decision_info}", data) return True else: self.log_result("Conversation Follow-up", False, f"Status: {response2.status_code}") return False except Exception as e: self.log_result("Conversation Follow-up", False, f"Exception: {str(e)}") return False def test_search_decision_patterns(self) -> bool: """Test different search decision patterns""" test_cases = [ { "name": "Elaboration Request", "prompt": "Tell me more about that", "history": [{"role": "assistant", "content": "AI is a field of computer science."}], "expected_search": False }, { "name": "New Information Request", "prompt": "What is the latest news about AI?", "history": [], "expected_search": True }, { "name": "Clarification Request", "prompt": "What do you mean by that?", "history": [{"role": "assistant", "content": "Machine learning uses algorithms."}], "expected_search": False } ] success_count = 0 for case in test_cases: try: payload = { "prompt": case["prompt"], "max_new_tokens": 50, "use_search": True, "temperature": 0.7, "history": case.get("history", []) } response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30) if response.status_code == 200: data = response.json() if "search_decision" in data: actual_search = data["search_decision"].get("should_search") expected = case["expected_search"] if actual_search == expected: self.log_result(f"Search Pattern: {case['name']}", True, f"Correctly decided {'to search' if actual_search else 'not to search'}") success_count += 1 else: self.log_result(f"Search Pattern: {case['name']}", False, f"Expected {expected}, got {actual_search}") else: self.log_result(f"Search Pattern: {case['name']}", False, "No search decision in response") else: self.log_result(f"Search Pattern: {case['name']}", False, f"Status: {response.status_code}") except Exception as e: self.log_result(f"Search Pattern: {case['name']}", False, f"Exception: {str(e)}") return success_count == len(test_cases) def test_force_search_parameter(self) -> bool: """Test force_search parameter override""" try: payload = { "prompt": "Tell me more", # Would normally not search "max_new_tokens": 50, "use_search": True, "force_search": True, # Should override decision "temperature": 0.7, "history": [{"role": "assistant", "content": "Here's some information about AI."}] } response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30) if response.status_code == 200: data = response.json() if "search_decision" in data: decision = data["search_decision"] if decision.get("should_search") and "forced" in decision.get("reason", "").lower(): self.log_result("Force Search Override", True, "Force search parameter worked correctly") return True else: self.log_result("Force Search Override", False, "Force search not detected in decision") return False else: self.log_result("Force Search Override", False, "No search decision in response") return False else: self.log_result("Force Search Override", False, f"Status: {response.status_code}") return False except Exception as e: self.log_result("Force Search Override", False, f"Exception: {str(e)}") return False def run_all_tests(self): """Run all tests and print summary""" print("šŸš€ Starting Atlas Functionality Tests") print("=" * 50) test_methods = [ self.test_health_endpoint, self.test_anonymous_chat_no_search, self.test_anonymous_chat_with_search, self.test_authenticated_chat, self.test_conversation_follow_up, self.test_search_decision_patterns, self.test_force_search_parameter ] passed = 0 total = len(test_methods) for test_method in test_methods: try: if test_method(): passed += 1 time.sleep(1) # Brief pause between tests except Exception as e: print(f"āŒ Test {test_method.__name__} crashed: {str(e)}") print("\n" + "=" * 50) print(f"šŸ“Š Test Summary: {passed}/{total} tests passed") if passed == total: print("šŸŽ‰ All tests passed! Atlas is functioning correctly.") else: print(f"āš ļø {total - passed} tests failed. Check the details above.") return passed == total if __name__ == "__main__": tester = AtlasTestSuite() success = tester.run_all_tests() # Save detailed results with open("test_results.json", "w") as f: json.dump(tester.results, f, indent=2) print(f"\nšŸ“„ Detailed results saved to test_results.json") exit(0 if success else 1)