eveAI/tests/interactive_client/specialist_client.py

#!/usr/bin/env python3
import json
import logging
import sys
import time
import requests  # Used for calling the auth API
from datetime import datetime
import yaml      # For loading the YAML configuration
from urllib.parse import urlparse

import socketio  # Official python-socketio client

# ----------------------------
# Constants for authentication and specialist selection
# ----------------------------
API_KEY = "EveAI-8342-2966-4731-6578-1010-8903-4230-4378"
TENANT_ID = 2
SPECIALIST_ID = 2
BASE_API_URL = "http://macstudio.ask-eve-ai-local.com:8080/api/api/v1"
BASE_SOCKET_URL = "http://macstudio.ask-eve-ai-local.com:8080"
CONFIG_FILE = "config/specialists/SPIN_SPECIALIST/1.0.0.yaml"  # Path to specialist configuration

# ----------------------------
# Logging Configuration
# ----------------------------
LOG_FILENAME = "specialist_client.log"
logging.basicConfig(
    filename=LOG_FILENAME,
    level=logging.DEBUG,
    format="%(asctime)s %(levelname)s: %(message)s"
)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
logging.getLogger('').addHandler(console_handler)

# ----------------------------
# Create the Socket.IO client using the official python-socketio client
# ----------------------------
sio = socketio.Client(logger=True, engineio_logger=True)
room = None  # Global variable to store the assigned room

# ----------------------------
# Event Handlers
# ----------------------------
@sio.event
def connect():
    logging.info("Connected to Socket.IO server.")
    print("Connected to server.")

@sio.event
def disconnect():
    logging.info("Disconnected from Socket.IO server.")
    print("Disconnected from server.")

@sio.on("connect_error")
def on_connect_error(data):
    logging.error("Connect error: %s", data)
    print("Connect error:", data)

@sio.on("authenticated")
def on_authenticated(data):
    global room
    room = data.get("room")
    logging.info("Authenticated. Room: %s", room)
    print("Authenticated. Room:", room)

@sio.on("room_join")
def on_room_join(data):
    global room
    room = data.get("room")
    logging.info("Room join event received. Room: %s", room)
    print("Joined room:", room)

@sio.on("token_expired")
def on_token_expired(data):
    logging.warning("Token expired.")
    print("Token expired. Please refresh your session.")

@sio.on("reconnect_attempt")
def on_reconnect_attempt(attempt):
    logging.info("Reconnect attempt #%s", attempt)
    print(f"Reconnect attempt #{attempt}")

@sio.on("reconnect")
def on_reconnect():
    logging.info("Reconnected successfully.")
    print("Reconnected to server.")

@sio.on("reconnect_failed")
def on_reconnect_failed():
    logging.error("Reconnection failed.")
    print("Reconnection failed. Please refresh.")

@sio.on("room_rejoin_result")
def on_room_rejoin_result(data):
    if data.get("success"):
        global room
        room = data.get("room")
        logging.info("Successfully rejoined room: %s", room)
        print("Rejoined room:", room)
    else:
        logging.error("Failed to rejoin room.")
        print("Failed to rejoin room.")

@sio.on("bot_response")
def on_bot_response(data):
    logging.info("Received bot response: %s", data)
    print("Bot response received:")
    print(json.dumps(data, indent=2))

@sio.on("task_status")
def on_task_status(data):
    logging.info("Received task status: %s", data)
    print("Task status:")
    print(json.dumps(data, indent=2))

# ----------------------------
# Helper: Retrieve token from REST API
# ----------------------------
def retrieve_token(api_url: str) -> str:
    payload = {
        "tenant_id": TENANT_ID,
        "api_key": API_KEY
    }
    try:
        logging.info("Requesting token from %s with payload: %s", api_url, payload)
        response = requests.post(api_url, json=payload)
        response.raise_for_status()
        token = response.json()["access_token"]
        logging.info("Token retrieved successfully.")
        return token
    except Exception as e:
        logging.error("Failed to retrieve token: %s", e)
        raise e

# ----------------------------
# Main Interactive UI Function
# ----------------------------
def main():
    global room

    # Retrieve the token
    auth_url = f"{BASE_API_URL}/auth/token"
    try:
        token = retrieve_token(auth_url)
        print("Token retrieved successfully.")
    except Exception as e:
        print("Error retrieving token. Check logs for details.")
        sys.exit(1)

    # Parse the BASE_SOCKET_URL
    parsed_url = urlparse(BASE_SOCKET_URL)
    host_url = f"{parsed_url.scheme}://{parsed_url.netloc}"

    # Connect to the Socket.IO server.
    # Note: Use `auth` instead of `query_string` (the official client uses the `auth` parameter)
    try:
        sio.connect(
            host_url,
            socketio_path='/chat/socket.io',
            auth={"token": token},
        )
    except Exception as e:
        logging.error("Failed to connect to Socket.IO server: %s", e)
        print("Failed to connect to Socket.IO server:", e)
        sys.exit(1)

    # Allow time for authentication and room assignment.
    time.sleep(2)
    if not room:
        logging.warning("No room assigned. Exiting.")
        print("No room assigned by the server. Exiting.")
        sio.disconnect()
        sys.exit(1)

    # Load specialist configuration from YAML.
    try:
        with open(CONFIG_FILE, "r") as f:
            specialist_config = yaml.safe_load(f)
        arg_config = specialist_config.get("arguments", {})
        logging.info("Loaded specialist argument configuration: %s", arg_config)
    except Exception as e:
        logging.error("Failed to load specialist configuration: %s", e)
        print("Failed to load specialist configuration. Exiting.")
        sys.exit(1)

    # Dictionary to store default values for static arguments (except "query")
    static_defaults = {}

    print("\nInteractive Specialist Client")
    print("For each iteration, you will be prompted for the following arguments:")
    for key, details in arg_config.items():
        print(f" - {details.get('name', key)}: {details.get('description', '')}")
    print("Type 'quit' or 'exit' as the query to end the session.\n")

    # Interactive loop: prompt for arguments and send user message.
    while True:
        current_arguments = {}
        for arg_key, arg_details in arg_config.items():
            prompt_msg = f"Enter {arg_details.get('name', arg_key)}"
            desc = arg_details.get("description", "")
            if desc:
                prompt_msg += f" ({desc})"
            if arg_key != "query":
                default_value = static_defaults.get(arg_key, "")
                if default_value:
                    prompt_msg += f" [default: {default_value}]"
                prompt_msg += ": "
                value = input(prompt_msg).strip()
                if not value:
                    value = default_value
                static_defaults[arg_key] = value
            else:
                prompt_msg += " (required): "
                value = input(prompt_msg).strip()
                while not value:
                    print("Query is required. Please enter a value.")
                    value = input(prompt_msg).strip()
            current_arguments[arg_key] = value

        if current_arguments.get("query", "").lower() in ["quit", "exit"]:
            break

        try:
            timezone = datetime.now().astimezone().tzname()
        except Exception:
            timezone = "UTC"

        payload = {
            "token": token,
            "tenant_id": TENANT_ID,
            "specialist_id": SPECIALIST_ID,
            "arguments": current_arguments,
            "timezone": timezone,
            "room": room
        }

        logging.info("Sending user_message with payload: %s", payload)
        print("Sending message to specialist...")
        sio.emit("user_message", payload)
        time.sleep(1)

    print("Exiting interactive session.")
    sio.disconnect()

if __name__ == "__main__":
    main()