test_llm_api.py

"""
Test script for LLM inference API endpoints.

Run the server first:
    python app.py

Then in another terminal:
    python test_llm_api.py
"""

import requests
import json
from typing import Dict, Any


BASE_URL = "http://localhost:7860"


def test_health() -> None:
    """Test the health endpoint."""
    print("\n=== Testing GET /api/v1/llm/health ===")
    response = requests.get(f"{BASE_URL}/api/v1/llm/health")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    assert response.status_code == 200 or response.status_code == 503


def test_chat_simple() -> None:
    """Test basic chat without context."""
    print("\n=== Testing POST /api/v1/llm/chat (simple) ===")
    payload = {
        "content": "What is drip irrigation?",
        "max_tokens": 150,
        "temperature": 0.7,
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/chat", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Model: {data.get('model')}")
    print(f"Content: {data.get('content')}")
    print(f"Latency (ms): {data.get('latency_ms')}")
    assert response.status_code == 200 or response.status_code == 500


def test_chat_with_context() -> None:
    """Test chat with farm context."""
    print("\n=== Testing POST /api/v1/llm/chat (with farm context) ===")
    payload = {
        "content": "How much water should I apply weekly?",
        "farm_context": {
            "farm_name": "Johnson Farm",
            "crop": "tomato",
            "area_ha": 2.5,
        },
        "max_tokens": 200,
        "temperature": 0.5,
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/chat", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Model: {data.get('model')}")
    print(f"Content: {data.get('content')}")
    print(f"Latency (ms): {data.get('latency_ms')}")
    print(f"Metadata: {json.dumps(data.get('metadata'), indent=2)}")
    assert response.status_code == 200 or response.status_code == 500


def test_chat_with_history() -> None:
    """Test chat with conversation history."""
    print("\n=== Testing POST /api/v1/llm/chat (with history) ===")
    payload = {
        "content": "What about pepper crops?",
        "conversation_history": [
            {"role": "user", "content": "How much water should I apply weekly?"},
            {"role": "assistant", "content": "For tomatoes, 25-40mm per week is typical."},
        ],
        "farm_context": {
            "farm_name": "Johnson Farm",
            "crop": "pepper",
            "area_ha": 1.5,
        },
        "max_tokens": 200,
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/chat", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Content: {data.get('content')}")
    print(f"Conversation history length: {data.get('metadata', {}).get('conversation_history_length')}")
    assert response.status_code == 200 or response.status_code == 500


def test_validate_context_valid() -> None:
    """Test context validation with valid data."""
    print("\n=== Testing POST /api/v1/llm/validate-context (valid) ===")
    payload = {
        "farm_context": {
            "farm_name": "Smith Farm",
            "crop": "lettuce",
            "area_ha": 0.5,
        }
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/validate-context", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Valid: {data.get('valid')}")
    print(f"Warnings: {data.get('warnings')}")
    print(f"Errors: {data.get('errors')}")
    assert response.status_code == 200
    assert data.get('valid') is True


def test_validate_context_invalid_crop() -> None:
    """Test context validation with invalid crop."""
    print("\n=== Testing POST /api/v1/llm/validate-context (invalid crop) ===")
    payload = {
        "farm_context": {
            "farm_name": "Smith Farm",
            "crop": "invalid_crop",
            "area_ha": 0.5,
        }
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/validate-context", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Valid: {data.get('valid')}")
    print(f"Warnings: {data.get('warnings')}")
    assert response.status_code == 200
    # Should still be valid (warnings don't fail validation)
    assert data.get('valid') is True


def test_validate_context_invalid_area() -> None:
    """Test context validation with invalid area type."""
    print("\n=== Testing POST /api/v1/llm/validate-context (invalid area) ===")
    payload = {
        "farm_context": {
            "farm_name": "Smith Farm",
            "crop": "tomato",
            "area_ha": "not a number",
        }
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/validate-context", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Valid: {data.get('valid')}")
    print(f"Errors: {data.get('errors')}")
    assert response.status_code == 200
    assert data.get('valid') is False


def test_batch_chat() -> None:
    """Test batch chat endpoint."""
    print("\n=== Testing POST /api/v1/llm/chat/batch ===")
    payload = [
        {
            "content": "What is drip irrigation?",
            "max_tokens": 100,
        },
        {
            "content": "How do I install valves?",
            "max_tokens": 100,
        },
        {
            "content": "What is emitter spacing?",
            "max_tokens": 100,
            "farm_context": {
                "crop": "tomato",
                "area_ha": 1.0,
            }
        },
    ]
    response = requests.post(f"{BASE_URL}/api/v1/llm/chat/batch", json=payload)
    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        print(f"Number of responses: {len(data)}")
        for i, resp in enumerate(data):
            print(f"\nResponse {i+1}:")
            print(f"  Content: {resp.get('content')[:100]}...")
            print(f"  Latency: {resp.get('latency_ms'):.1f}ms")


def test_chat_invalid_request() -> None:
    """Test chat with invalid request (missing required field)."""
    print("\n=== Testing POST /api/v1/llm/chat (invalid request) ===")
    payload = {
        "max_tokens": 100,
        # Missing required 'content' field
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/chat", json=payload)
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    assert response.status_code == 422


def test_chat_oversized_request() -> None:
    """Test chat with oversized content."""
    print("\n=== Testing POST /api/v1/llm/chat (oversized content) ===")
    payload = {
        "content": "x" * 3000,  # Exceeds max_length of 2000
    }
    response = requests.post(f"{BASE_URL}/api/v1/llm/chat", json=payload)
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    assert response.status_code == 422


if __name__ == "__main__":
    print("=" * 70)
    print("LLM Inference API Test Suite")
    print("=" * 70)
    
    try:
        # Health check first
        test_health()
        
        # Basic tests
        test_chat_simple()
        test_chat_with_context()
        test_chat_with_history()
        
        # Context validation tests
        test_validate_context_valid()
        test_validate_context_invalid_crop()
        test_validate_context_invalid_area()
        
        # Batch test
        test_batch_chat()
        
        # Error handling tests
        test_chat_invalid_request()
        test_chat_oversized_request()
        
        print("\n" + "=" * 70)
        print("✅ All tests completed!")
        print("=" * 70)
        
    except requests.exceptions.ConnectionError:
        print("\n❌ Error: Could not connect to server.")
        print("   Make sure the server is running: python app.py")
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
        import traceback
        traceback.print_exc()