client_example.py

#!/usr/bin/env python3
"""
AIM Client Example - Connect to Remote MI300X Endpoint

This script demonstrates how to connect to an AIM inference service
running on a remote MI300X node from your laptop.

Usage:
    python client_example.py --endpoint http://remote-host:8000
    python client_example.py --endpoint http://localhost:8000  # if using SSH port forwarding
"""

import argparse
import json
import requests
import sys
from typing import Optional, Dict, Any, Iterator


class AIMClient:
    """Client for interacting with AIM inference service"""
    
    def __init__(self, base_url: str):
        """
        Initialize AIM client
        
        Args:
            base_url: Base URL of the AIM endpoint (e.g., http://remote-host:8000)
        """
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
        self.session.headers.update({
            'Content-Type': 'application/json'
        })
    
    def health_check(self) -> bool:
        """Check if the endpoint is healthy"""
        try:
            response = self.session.get(f"{self.base_url}/health", timeout=5)
            return response.status_code == 200
        except requests.exceptions.RequestException:
            return False
    
    def list_models(self) -> Dict[str, Any]:
        """List available models"""
        try:
            response = self.session.get(f"{self.base_url}/v1/models", timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error listing models: {e}", file=sys.stderr)
            return {}
    
    def chat_completion(
        self,
        messages: list,
        model: Optional[str] = None,
        max_tokens: int = 2048,
        temperature: float = 0.7,
        stream: bool = False
    ) -> Dict[str, Any]:
        """
        Send a chat completion request
        
        Args:
            messages: List of message dicts with 'role' and 'content'
            model: Model name (optional, will use default if not specified)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            stream: Whether to stream the response
        
        Returns:
            Response dict (or generator if streaming)
        """
        payload = {
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": stream
        }
        
        if model:
            payload["model"] = model
        
        try:
            if stream:
                return self._stream_chat_completion(payload)
            else:
                response = self.session.post(
                    f"{self.base_url}/v1/chat/completions",
                    json=payload,
                    timeout=300  # 5 minute timeout for long responses
                )
                response.raise_for_status()
                return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error in chat completion: {e}", file=sys.stderr)
            raise
    
    def _stream_chat_completion(self, payload: Dict[str, Any]) -> Iterator[Dict[str, Any]]:
        """Stream chat completion responses"""
        response = self.session.post(
            f"{self.base_url}/v1/chat/completions",
            json=payload,
            stream=True,
            timeout=300
        )
        response.raise_for_status()
        
        for line in response.iter_lines():
            if not line:
                continue
            
            line = line.decode('utf-8')
            if line.startswith('data: '):
                data = line[6:].strip()
                if data == '[DONE]':
                    break
                
                try:
                    chunk = json.loads(data)
                    yield chunk
                except json.JSONDecodeError:
                    continue
    
    def completion(
        self,
        prompt: str,
        max_tokens: int = 50,
        temperature: float = 0.7
    ) -> Dict[str, Any]:
        """
        Send a text completion request
        
        Args:
            prompt: Text prompt
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature
        
        Returns:
            Response dict
        """
        payload = {
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature
        }
        
        try:
            response = self.session.post(
                f"{self.base_url}/v1/completions",
                json=payload,
                timeout=300
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error in completion: {e}", file=sys.stderr)
            raise


def main():
    parser = argparse.ArgumentParser(
        description="AIM Client Example - Connect to Remote MI300X Endpoint"
    )
    parser.add_argument(
        '--endpoint',
        type=str,
        required=True,
        help='AIM endpoint URL (e.g., http://remote-host:8000 or http://localhost:8000)'
    )
    parser.add_argument(
        '--model',
        type=str,
        help='Model name (optional, will use default if not specified)'
    )
    parser.add_argument(
        '--prompt',
        type=str,
        default="Hello! Can you explain what AIM is?",
        help='Prompt/question to send'
    )
    parser.add_argument(
        '--max-tokens',
        type=int,
        default=2048,
        help='Maximum tokens to generate (default: 2048)'
    )
    parser.add_argument(
        '--stream',
        action='store_true',
        help='Stream the response'
    )
    parser.add_argument(
        '--list-models',
        action='store_true',
        help='List available models and exit'
    )
    
    args = parser.parse_args()
    
    # Initialize client
    client = AIMClient(args.endpoint)
    
    # Health check
    print(f"Connecting to {args.endpoint}...")
    if not client.health_check():
        print(f"⚠ Warning: Health check failed. Endpoint may not be ready.")
        print(f"  Continuing anyway...")
    else:
        print("✓ Endpoint is healthy")
    
    print()
    
    # List models if requested
    if args.list_models:
        print("Available models:")
        models = client.list_models()
        if models and 'data' in models:
            for model in models['data']:
                print(f"  - {model.get('id', 'unknown')}")
                if 'max_model_len' in model:
                    print(f"    Max context length: {model['max_model_len']}")
        else:
            print("  No models found or error retrieving models")
        return
    
    # Test chat completion
    print(f"Sending prompt: {args.prompt}")
    print()
    
    messages = [{"role": "user", "content": args.prompt}]
    
    if args.stream:
        print("Response (streaming):")
        print("-" * 80)
        
        try:
            for chunk in client.chat_completion(
                messages=messages,
                model=args.model,
                max_tokens=args.max_tokens,
                stream=True
            ):
                if 'choices' in chunk and len(chunk['choices']) > 0:
                    delta = chunk['choices'][0].get('delta', {})
                    # Qwen3 uses reasoning_content for thinking, content for response
                    reasoning = delta.get('reasoning_content', '')
                    content = delta.get('content', '')
                    
                    if reasoning:
                        print(reasoning, end='', flush=True)
                    if content:
                        print(content, end='', flush=True)
            
            print()  # Newline at end
            print("-" * 80)
        except Exception as e:
            print(f"\nError: {e}", file=sys.stderr)
            sys.exit(1)
    else:
        print("Response:")
        print("-" * 80)
        
        try:
            response = client.chat_completion(
                messages=messages,
                model=args.model,
                max_tokens=args.max_tokens,
                stream=False
            )
            
            if 'choices' in response and len(response['choices']) > 0:
                content = response['choices'][0]['message']['content']
                print(content)
                
                # Show usage if available
                if 'usage' in response:
                    usage = response['usage']
                    print()
                    print(f"Tokens used: {usage.get('total_tokens', 'N/A')} "
                          f"(prompt: {usage.get('prompt_tokens', 'N/A')}, "
                          f"completion: {usage.get('completion_tokens', 'N/A')})")
            else:
                print("No response received")
                print(json.dumps(response, indent=2))
            
            print("-" * 80)
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            sys.exit(1)


if __name__ == "__main__":
    main()