client.py

#!/usr/bin/env python3

import requests
import json
import sys
from typing import Dict, Any

class AnswerBotClient:
    def __init__(self, router_url: str = "http://localhost:8000"):
        self.router_url = router_url
        self.session = requests.Session()
    
    def ask_question(self, question: str) -> Dict[str, Any]:
        try:
            response = self.session.post(
                f"{self.router_url}/ask",
                json={"question": question},
                timeout=30
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"error": f"Failed to connect to AnswerBot: {str(e)}"}
    
    def get_health(self) -> Dict[str, Any]:
        try:
            response = self.session.get(f"{self.router_url}/health", timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"error": f"Health check failed: {str(e)}"}
    
    def get_agents(self) -> Dict[str, Any]:
        try:
            response = self.session.get(f"{self.router_url}/agents", timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"error": f"Failed to get agents info: {str(e)}"}

def print_response(response: Dict[str, Any]) -> None:
    if "error" in response:
        print(f"❌ Error: {response['error']}")
        return
    
    print("\n" + "="*50)
    
    if "agent" in response:
        agent_name = response["agent"]
        print(f"🤖 Agent: {agent_name}")
    
    if "question" in response:
        print(f"❓ Question: {response['question']}")
    
    # Handle response types
    if agent_name == "MathAgent" and "result" in response:
        print(f"🔢 Expression: {response.get('expression', 'N/A')}")
        print(f"✅ Result: {response['result']}")
    elif agent_name == "WeatherAgent" and "weather" in response:
        weather = response["weather"]
        city = response.get("city", "Unknown")
        print(f"🌍 City: {city}")
        print(f"🌡️  Temperature: {weather.get('temperature', 'N/A')}")
        print(f"☁️  Condition: {weather.get('condition', 'N/A')}")
        print(f"💧 Humidity: {weather.get('humidity', 'N/A')}")
        if "response" in response:
            print(f"📝 Summary: {response['response']}")
    elif agent_name == "QnAAgent" and "response" in response:
        print(f"💬 Answer: {response['response']}")
    
    # Routing info
    if "router_info" in response:
        router_info = response["router_info"]
        print(f"\n🔀 Routing Info:")
        print(f"   Selected Agent: {router_info.get('selected_agent', 'N/A')}")
        print(f"   Reason: {router_info.get('routing_reason', 'N/A')}")
    
    print("="*50)

def interactive_mode(client: AnswerBotClient) -> None:
    print("🤖 AI AnswerBot Interactive Client")
    print("Type 'help' for commands, 'quit' to exit")
    print("-" * 40)
    
    while True:
        try:
            try:
                question = input("\n💭 Ask me anything: ").strip()
            except EOFError:
                print("\n👋 Session ended (EOF detected)")
                break
            
            if not question:
                continue
            
            if question.lower() in ['quit', 'exit', 'q']:
                print("👋 Goodbye!")
                break
            
            if question.lower() == 'help':
                print_help()
                continue
            
            if question.lower() == 'health':
                health = client.get_health()
                print(f"🏥 Health Status: {json.dumps(health, indent=2)}")
                continue
            
            if question.lower() == 'agents':
                agents = client.get_agents()
                print(f"🔍 Available Agents: {json.dumps(agents, indent=2)}")
                continue
            
            print("🔄 Processing your question...")
            response = client.ask_question(question)
            print_response(response)
            
        except KeyboardInterrupt:
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"❌ Unexpected error: {str(e)}")

def print_help() -> None:
    print("\n📚 Available Commands:")
    print("  help    - Show this help message")
    print("  health  - Check router health")
    print("  agents  - List available agents")
    print("  quit    - Exit the client")
    print("\n🎯 Example Questions:")
    print("  Math: 'What is 25 + 17?', 'Calculate 100 / 4'")
    print("  Weather: 'Weather in London?', 'Temperature in Tokyo?'")
    print("  General: 'What is AI?', 'Hello, how are you?'")

def main():
    if len(sys.argv) > 1:
        # Non-interactive mode
        router_url = sys.argv[2] if len(sys.argv) > 2 else "http://localhost:8000"
        client = AnswerBotClient(router_url)
        
        question = sys.argv[1]
        print(f"🤖 AI AnswerBot - Single Query Mode")
        print(f"Question: {question}")
        
        response = client.ask_question(question)
        print_response(response)
    else:
        # Interactive mode
        client = AnswerBotClient()
        
        # Check service availability
        health = client.get_health()
        if "error" in health:
            print(f"❌ Cannot connect to AnswerBot service: {health['error']}")
            print("💡 Make sure the service is running on http://localhost:8000")
            sys.exit(1)
        
        interactive_mode(client)

if __name__ == "__main__":
    main()