findEthics/Atlas
0
1#!/usr/bin/env python32"""3Atlas Functionality Test Suite4Tests the chat endpoint and search decision waterfall logic5"""6 7import asyncio8import json9import requests10from typing import Dict, Any, Optional, List11import time12 13# Test configuration14BASE_URL = "http://localhost:7860"15CHAT_ENDPOINT = f"{BASE_URL}/chat"16HEALTH_ENDPOINT = f"{BASE_URL}/"17 18class AtlasTestSuite:19 def __init__(self):20 self.results = []21 self.session_id = None22 23 def log_result(self, test_name: str, success: bool, details: str = "", response_data: Dict = None):24 """Log test result"""25 result = {26 "test": test_name,27 "success": success,28 "details": details,29 "timestamp": time.time()30 }31 if response_data:32 result["response_data"] = response_data33 self.results.append(result)34 35 status = "✅ PASS" if success else "❌ FAIL"36 print(f"{status} {test_name}: {details}")37 38 def test_health_endpoint(self) -> bool:39 """Test the health endpoint"""40 try:41 response = requests.get(HEALTH_ENDPOINT, timeout=10)42 if response.status_code == 200:43 data = response.json()44 self.log_result("Health Check", True, "Server responding correctly", data)45 return True46 else:47 self.log_result("Health Check", False, f"Status code: {response.status_code}")48 return False49 except Exception as e:50 self.log_result("Health Check", False, f"Exception: {str(e)}")51 return False52 53 def test_anonymous_chat_no_search(self) -> bool:54 """Test anonymous chat request without search"""55 try:56 payload = {57 "prompt": "Hello, how are you?",58 "max_new_tokens": 100,59 "use_search": False,60 "temperature": 0.761 }62 63 response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30)64 65 if response.status_code == 200:66 data = response.json()67 # Check response structure68 if "response" in data and len(data["response"]) > 0:69 # Store session ID for follow-up tests70 self.session_id = response.headers.get("X-Session-ID")71 self.log_result("Anonymous Chat (No Search)", True, 72 f"Response length: {len(data['response'])} chars, Session: {self.session_id}", 73 data)74 return True75 else:76 self.log_result("Anonymous Chat (No Search)", False, "Empty or missing response")77 return False78 else:79 self.log_result("Anonymous Chat (No Search)", False, f"Status: {response.status_code}")80 return False81 82 except Exception as e:83 self.log_result("Anonymous Chat (No Search)", False, f"Exception: {str(e)}")84 return False85 86 def test_anonymous_chat_with_search(self) -> bool:87 """Test anonymous chat request with search enabled"""88 try:89 payload = {90 "prompt": "What is artificial intelligence?",91 "max_new_tokens": 150,92 "use_search": True,93 "temperature": 0.794 }95 96 response = requests.post(CHAT_ENDPOINT, json=payload, timeout=45)97 98 if response.status_code == 200:99 data = response.json()100 # Check response structure101 has_response = "response" in data and len(data["response"]) > 0102 has_search_decision = "search_decision" in data103 has_cache_info = "cache_info" in data104 105 details = f"Response: {has_response}, Decision: {has_search_decision}, Cache: {has_cache_info}"106 107 if has_response:108 self.log_result("Anonymous Chat (With Search)", True, details, data)109 return True110 else:111 self.log_result("Anonymous Chat (With Search)", False, "Missing response")112 return False113 else:114 self.log_result("Anonymous Chat (With Search)", False, f"Status: {response.status_code}")115 return False116 117 except Exception as e:118 self.log_result("Anonymous Chat (With Search)", False, f"Exception: {str(e)}")119 return False120 121 def test_authenticated_chat(self) -> bool:122 """Test authenticated chat request"""123 try:124 payload = {125 "prompt": "Hello, I'm a test user. Can you help me?",126 "max_new_tokens": 100,127 "use_search": False,128 "temperature": 0.7,129 "user_id": "test_user_001"130 }131 132 response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30)133 134 if response.status_code == 200:135 data = response.json()136 if "response" in data and len(data["response"]) > 0:137 self.log_result("Authenticated Chat", True, 138 f"Response length: {len(data['response'])} chars", data)139 return True140 else:141 self.log_result("Authenticated Chat", False, "Empty or missing response")142 return False143 else:144 self.log_result("Authenticated Chat", False, f"Status: {response.status_code}")145 return False146 147 except Exception as e:148 self.log_result("Authenticated Chat", False, f"Exception: {str(e)}")149 return False150 151 def test_conversation_follow_up(self) -> bool:152 """Test conversation follow-up to check search decision waterfall"""153 if not self.session_id:154 self.log_result("Conversation Follow-up", False, "No session ID available")155 return False156 157 try:158 # First establish conversation history159 payload1 = {160 "prompt": "What is machine learning?",161 "max_new_tokens": 100,162 "use_search": True,163 "temperature": 0.7164 }165 166 headers = {"X-Session-ID": self.session_id}167 response1 = requests.post(CHAT_ENDPOINT, json=payload1, headers=headers, timeout=45)168 169 if response1.status_code != 200:170 self.log_result("Conversation Follow-up", False, f"First request failed: {response1.status_code}")171 return False172 173 # Now test follow-up question (should trigger different search logic)174 payload2 = {175 "prompt": "Can you tell me more about that?",176 "max_new_tokens": 100,177 "use_search": True,178 "temperature": 0.7,179 "history": [180 {"role": "user", "content": "What is machine learning?"},181 {"role": "assistant", "content": response1.json()["response"][:200]}182 ]183 }184 185 response2 = requests.post(CHAT_ENDPOINT, json=payload2, headers=headers, timeout=45)186 187 if response2.status_code == 200:188 data = response2.json()189 has_search_decision = "search_decision" in data190 decision_info = ""191 192 if has_search_decision:193 decision = data["search_decision"]194 decision_info = f"Should search: {decision.get('should_search')}, " \195 f"Confidence: {decision.get('confidence')}, " \196 f"Flow: {decision.get('flow_type', 'unknown')}"197 198 self.log_result("Conversation Follow-up", True, 199 f"Follow-up successful. {decision_info}", data)200 return True201 else:202 self.log_result("Conversation Follow-up", False, f"Status: {response2.status_code}")203 return False204 205 except Exception as e:206 self.log_result("Conversation Follow-up", False, f"Exception: {str(e)}")207 return False208 209 def test_search_decision_patterns(self) -> bool:210 """Test different search decision patterns"""211 test_cases = [212 {213 "name": "Elaboration Request",214 "prompt": "Tell me more about that",215 "history": [{"role": "assistant", "content": "AI is a field of computer science."}],216 "expected_search": False217 },218 {219 "name": "New Information Request", 220 "prompt": "What is the latest news about AI?",221 "history": [],222 "expected_search": True223 },224 {225 "name": "Clarification Request",226 "prompt": "What do you mean by that?",227 "history": [{"role": "assistant", "content": "Machine learning uses algorithms."}],228 "expected_search": False229 }230 ]231 232 success_count = 0233 234 for case in test_cases:235 try:236 payload = {237 "prompt": case["prompt"],238 "max_new_tokens": 50,239 "use_search": True,240 "temperature": 0.7,241 "history": case.get("history", [])242 }243 244 response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30)245 246 if response.status_code == 200:247 data = response.json()248 if "search_decision" in data:249 actual_search = data["search_decision"].get("should_search")250 expected = case["expected_search"]251 252 if actual_search == expected:253 self.log_result(f"Search Pattern: {case['name']}", True, 254 f"Correctly decided {'to search' if actual_search else 'not to search'}")255 success_count += 1256 else:257 self.log_result(f"Search Pattern: {case['name']}", False, 258 f"Expected {expected}, got {actual_search}")259 else:260 self.log_result(f"Search Pattern: {case['name']}", False, "No search decision in response")261 else:262 self.log_result(f"Search Pattern: {case['name']}", False, f"Status: {response.status_code}")263 264 except Exception as e:265 self.log_result(f"Search Pattern: {case['name']}", False, f"Exception: {str(e)}")266 267 return success_count == len(test_cases)268 269 def test_force_search_parameter(self) -> bool:270 """Test force_search parameter override"""271 try:272 payload = {273 "prompt": "Tell me more", # Would normally not search274 "max_new_tokens": 50,275 "use_search": True,276 "force_search": True, # Should override decision277 "temperature": 0.7,278 "history": [{"role": "assistant", "content": "Here's some information about AI."}]279 }280 281 response = requests.post(CHAT_ENDPOINT, json=payload, timeout=30)282 283 if response.status_code == 200:284 data = response.json()285 if "search_decision" in data:286 decision = data["search_decision"]287 if decision.get("should_search") and "forced" in decision.get("reason", "").lower():288 self.log_result("Force Search Override", True, "Force search parameter worked correctly")289 return True290 else:291 self.log_result("Force Search Override", False, "Force search not detected in decision")292 return False293 else:294 self.log_result("Force Search Override", False, "No search decision in response")295 return False296 else:297 self.log_result("Force Search Override", False, f"Status: {response.status_code}")298 return False299 300 except Exception as e:301 self.log_result("Force Search Override", False, f"Exception: {str(e)}")302 return False303 304 def run_all_tests(self):305 """Run all tests and print summary"""306 print("🚀 Starting Atlas Functionality Tests")307 print("=" * 50)308 309 test_methods = [310 self.test_health_endpoint,311 self.test_anonymous_chat_no_search,312 self.test_anonymous_chat_with_search,313 self.test_authenticated_chat,314 self.test_conversation_follow_up,315 self.test_search_decision_patterns,316 self.test_force_search_parameter317 ]318 319 passed = 0320 total = len(test_methods)321 322 for test_method in test_methods:323 try:324 if test_method():325 passed += 1326 time.sleep(1) # Brief pause between tests327 except Exception as e:328 print(f"❌ Test {test_method.__name__} crashed: {str(e)}")329 330 print("\n" + "=" * 50)331 print(f"📊 Test Summary: {passed}/{total} tests passed")332 333 if passed == total:334 print("🎉 All tests passed! Atlas is functioning correctly.")335 else:336 print(f"⚠️ {total - passed} tests failed. Check the details above.")337 338 return passed == total339 340if __name__ == "__main__":341 tester = AtlasTestSuite()342 success = tester.run_all_tests()343 344 # Save detailed results345 with open("test_results.json", "w") as f:346 json.dump(tester.results, f, indent=2)347 348 print(f"\n📄 Detailed results saved to test_results.json")349 exit(0 if success else 1)