ds_zagres_ai/test_ai_service.py

"""
Test script for the AI service API.
"""

import requests
import json
import argparse

def test_health(api_url):
    """Test the health endpoint."""
    print(f"Testing health endpoint at: {api_url}/health")

    try:
        response = requests.get(f"{api_url}/health", timeout=30)
        response.raise_for_status()

        print("Health check successful!")
        print(f"Response: {response.json()}")

        return True
    except Exception as e:
        print(f"ERROR: Health check failed: {str(e)}")
        return False

def test_config(api_url):
    """Test the config endpoint."""
    print(f"Testing config endpoint at: {api_url}/config")

    try:
        response = requests.get(f"{api_url}/config", timeout=30)
        response.raise_for_status()

        print("Config check successful!")
        print("Configuration:")
        config = response.json()
        for key, value in config.items():
            print(f"  {key}: {value}")

        return True, config
    except Exception as e:
        print(f"ERROR: Config check failed: {str(e)}")
        return False, None

def test_ollama_connection(api_url):
    """Test the Ollama connection."""
    print(f"Testing Ollama connection at: {api_url}/test-ollama")

    try:
        response = requests.get(f"{api_url}/test-ollama", timeout=60)
        response.raise_for_status()

        print("Ollama connection test successful!")
        result = response.json()
        print(f"Status: {result.get('status')}")
        print(f"Message: {result.get('message')}")
        print(f"Ollama URL: {result.get('ollama_url')}")

        if 'models' in result:
            print("Available models:")
            for model in result.get('models', {}).get('models', []):
                print(f"  - {model.get('name')}")

        return True
    except Exception as e:
        print(f"ERROR: Ollama connection test failed: {str(e)}")
        return False

def test_chat_completion(api_url, model_id, prompt):
    """Test the chat completion endpoint."""
    print(f"Testing chat completion at: {api_url}/test-ollama-direct")
    print(f"Model: {model_id}")
    print(f"Prompt: {prompt}")

    try:
        # First, create a chat
        chat_response = requests.post(
            f"{api_url}/chats",
            headers={"Content-Type": "application/json"},
            json={
                "user_id": "test_user",
                "title": "Test Chat",
                "model_id": model_id
            },
            timeout=60
        )

        if chat_response.status_code != 200:
            print(f"ERROR: Failed to create chat: {chat_response.status_code}")
            print(chat_response.text)
            # Try direct Ollama test instead
            print("Trying direct Ollama test...")
            response = requests.post(
                f"{api_url}/test-ollama-direct",
                headers={"Content-Type": "application/json"},
                json={
                    "model": model_id,
                    "prompt": prompt
                },
                timeout=120
            )

            response.raise_for_status()

            result = response.json()
            print("\nResponse:")
            print(result.get('response', 'No response'))

            return True

        chat_id = chat_response.json().get("id")
        print(f"Chat created with ID: {chat_id}")

        # Send a message to the chat
        message_response = requests.post(
            f"{api_url}/chats/{chat_id}/messages",
            headers={"Content-Type": "application/json"},
            json={
                "message": prompt,
                "user_id": "test_user",
                "temperature": 0.7,
                "max_tokens": 500
            },
            timeout=120
        )

        message_response.raise_for_status()

        result = message_response.json()
        print("\nResponse:")
        print(result.get('content', 'No content in response'))

        return True
    except Exception as e:
        print(f"ERROR: Chat completion test failed: {str(e)}")
        return False

def test_rag_completion(api_url, prompt):
    """Test the RAG completion endpoint."""
    print(f"Testing RAG completion at: {api_url}/ask-rag")
    print(f"Prompt: {prompt}")

    try:
        response = requests.post(
            f"{api_url}/ask-rag",
            headers={"Content-Type": "application/json"},
            json={"query": prompt},
            timeout=120
        )

        response.raise_for_status()

        result = response.json()
        print("\nResponse:")
        print(result.get('response', 'No response'))

        return True
    except Exception as e:
        print(f"ERROR: RAG completion test failed: {str(e)}")
        return False

def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Test the AI service API')
    parser.add_argument('--api-url', type=str, default='http://localhost:5252', help='AI service API URL')
    parser.add_argument('--model', type=str, default='llama3.1', help='Model to use for testing')
    parser.add_argument('--prompt', type=str, default='What is the capital of France?', help='Prompt to use for testing')
    parser.add_argument('--rag-prompt', type=str, default='What information do you have in your knowledge base?', help='Prompt to use for RAG testing')

    args = parser.parse_args()

    print("=== AI Service API Test ===")
    print(f"API URL: {args.api_url}")
    print(f"Model: {args.model}")
    print()

    # Test health endpoint
    print("=== Testing Health Endpoint ===")
    health_success = test_health(args.api_url)
    print()

    # Test config endpoint
    print("=== Testing Config Endpoint ===")
    config_success, config = test_config(args.api_url)
    print()

    # Test Ollama connection
    print("=== Testing Ollama Connection ===")
    ollama_success = test_ollama_connection(args.api_url)
    print()

    # Test chat completion
    if ollama_success:
        print("=== Testing Chat Completion ===")
        chat_success = test_chat_completion(args.api_url, args.model, args.prompt)
        print()

        # Test RAG completion
        print("=== Testing RAG Completion ===")
        rag_success = test_rag_completion(args.api_url, args.rag_prompt)
        print()
    else:
        chat_success = False
        rag_success = False

    # Print summary
    print("=== Test Summary ===")
    print(f"Health Endpoint: {'SUCCESS' if health_success else 'FAILED'}")
    print(f"Config Endpoint: {'SUCCESS' if config_success else 'FAILED'}")
    print(f"Ollama Connection: {'SUCCESS' if ollama_success else 'FAILED'}")
    print(f"Chat Completion: {'SUCCESS' if chat_success else 'FAILED'}")
    print(f"RAG Completion: {'SUCCESS' if rag_success else 'FAILED'}")

if __name__ == "__main__":
    main()