ollama_semantic_function_calling / ollamatools-new / tests / test_ollamatools_opensearch_direct.py

#!/usr/bin/env python3
"""
Test script to directly verify OpenSearchNG integration with OllamaTools.
"""

import logging
import json
import sys

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_opensearch_integration():
    """Test the OpenSearchNG integration with OllamaTools."""
    try:
        # Try to import the OpenSearchNG tools directly
        from ollamatools.plugins.integrations.opensearch.tools import (
            opensearch_query,
            get_opensearch_info,
            DEFAULT_OPENSEARCH_CONFIG
        )
        
        logger.info("Successfully imported OpenSearchNG tools")
        
        # Test getting OpenSearchNG info
        logger.info("Testing OpenSearchNG info...")
        info_result = get_opensearch_info()
        logger.info(f"OpenSearchNG info: {json.dumps(info_result, indent=2)[:200]}...")
        
        # Test OpenSearchNG query
        logger.info("Testing OpenSearchNG query...")
        query_result = opensearch_query("python programming")
        
        num_results = len(query_result.get("results", []))
        logger.info(f"Got {num_results} results for 'python programming'")
        
        # Print the first result
        if num_results > 0:
            first_result = query_result["results"][0]
            logger.info(f"First result: {json.dumps(first_result, indent=2)}")
            
        return True
        
    except ImportError as e:
        logger.error(f"Import error: {e}")
        return False
    except Exception as e:
        logger.error(f"Error: {e}")
        return False

if __name__ == "__main__":
    success = test_opensearch_integration()
    if success:
        logger.info("OpenSearchNG integration test successful!")
        sys.exit(0)
    else:
        logger.error("OpenSearchNG integration test failed!")
        sys.exit(1)