test_lm_studio_client.py

#!/usr/bin/env python3
"""Test script for LM Studio API client.

This script validates the LM Studio client implementation by:
1. Testing connection validation
2. Making a simple chat completion request
3. Demonstrating error handling

Run with: python test_lm_studio_client.py
"""

import logging
import sys

from src.paribusai.api import LMStudioClient, LMStudioConfig

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)

logger = logging.getLogger(__name__)


def main() -> int:
    """Run LM Studio client tests."""
    logger.info("Starting LM Studio client tests...")

    # Initialize client
    try:
        config = LMStudioConfig()
        logger.info(f"Configuration loaded: {config.lm_studio_url}")
    except Exception as e:
        logger.error(f"Failed to load configuration: {e}")
        return 1

    client = LMStudioClient(config)

    # Test 1: Connection validation
    logger.info("\n=== Test 1: Connection Validation ===")
    if not client.validate_connection():
        logger.error(
            "Failed to connect to LM Studio. "
            "Make sure LM Studio is running and a model is loaded."
        )
        return 1
    logger.info("✓ Connection validated successfully")

    # Test 2: List available models
    logger.info("\n=== Test 2: List Available Models ===")
    try:
        models = client.get_available_models()
        logger.info(f"✓ Available models: {models}")
    except Exception as e:
        logger.error(f"Failed to list models: {e}")
        return 1

    # Test 3: Simple chat completion
    logger.info("\n=== Test 3: Simple Chat Completion ===")
    try:
        messages = [
            {
                "role": "system",
                "content": "You are a helpful assistant. Keep responses concise.",
            },
            {
                "role": "user",
                "content": "What is economics in one sentence?",
            },
        ]

        response = client.chat_completion(messages)
        logger.info(f"✓ Chat completion successful!")
        logger.info(f"  Model: {response.model}")
        logger.info(f"  Finish reason: {response.finish_reason}")
        logger.info(f"  Token usage: {response.usage}")
        logger.info(f"  Response content: {response.content[:200]}...")

    except Exception as e:
        logger.error(f"Chat completion failed: {e}")
        return 1

    # Test 4: Temperature variation
    logger.info("\n=== Test 4: Temperature Variation ===")
    try:
        messages = [
            {
                "role": "user",
                "content": "List 3 fundamental economic concepts.",
            }
        ]

        # High temperature for more creative response
        response_creative = client.chat_completion(
            messages,
            temperature=1.0,
            max_tokens=150,
        )
        logger.info("✓ High temperature (1.0) response received")
        logger.info(f"  Response: {response_creative.content[:150]}...")

        # Zero temperature for deterministic response
        response_deterministic = client.chat_completion(
            messages,
            temperature=0.0,
            max_tokens=150,
        )
        logger.info("✓ Zero temperature (0.0) response received")
        logger.info(f"  Response: {response_deterministic.content[:150]}...")

    except Exception as e:
        logger.error(f"Temperature variation test failed: {e}")
        return 1

    # Test 5: Error handling (invalid messages)
    logger.info("\n=== Test 5: Error Handling ===")
    try:
        # This should raise ValueError
        client.chat_completion([])
        logger.error("Expected ValueError for empty messages!")
        return 1
    except ValueError as e:
        logger.info(f"✓ Correctly handled invalid input: {e}")

    try:
        # This should raise ValueError
        client.chat_completion([{"invalid": "message"}])
        logger.error("Expected ValueError for invalid message format!")
        return 1
    except ValueError as e:
        logger.info(f"✓ Correctly handled invalid message format: {e}")

    # Clean up
    logger.info("\n=== Cleanup ===")
    client.close()
    logger.info("✓ Client closed successfully")

    logger.info("\n" + "=" * 50)
    logger.info("All tests passed successfully!")
    logger.info("=" * 50)
    return 0


if __name__ == "__main__":
    sys.exit(main())