ollama_semantic_function_calling / ollamatools-new / tests / test_opensearch.py

#!/usr/bin/env python3
"""
Test script to verify OpenSearchNG functionality.
"""

import requests
import json
import logging
from typing import Dict, Any

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

def test_opensearch_query(query: str = "test") -> Dict[str, Any]:
    """
    Test OpenSearchNG query functionality.
    
    Args:
        query: The search query
        
    Returns:
        Dict[str, Any]: Search results
    """
    # Default OpenSearchNG configuration
    search_config = {
        "base_url": "http://localhost:8081",  # Default local OpenSearchNG URL
        "search_path": "/search",             # Path for search endpoint
        "format": "json",                     # Response format
        "categories": ["general"],            # Default search categories
        "engines": ["google", "bing", "duckduckgo"], # Default search engines
        "language": "en-US",                  # Default language
        "timeout": 10,                        # Default timeout in seconds
    }
    
    # Construct the search URL
    base_url = search_config["base_url"]
    search_path = search_config["search_path"]
    search_url = f"{base_url}{search_path}"
    
    # Prepare query parameters
    params = {
        "q": query,
        "format": search_config["format"],
        "categories": ",".join(search_config["categories"]),
        "engines": ",".join(search_config["engines"]),
        "language": search_config["language"],
    }
    
    # Create session with appropriate headers
    with requests.Session() as session:
        headers = {
            "User-Agent": "OllamaTools/0.2.0",
            "Accept": "application/json"
        }
        session.headers.update(headers)
        
        try:
            logger.info(f"Making OpenSearchNG query: {query}")
            response = session.get(
                search_url,
                params=params,
                timeout=search_config["timeout"]
            )
            
            # Check for successful response
            if response.status_code != 200:
                logger.error(f"OpenSearchNG query failed with status code: {response.status_code}")
                return {"error": f"Status code: {response.status_code}", "results": []}
            
            # Parse the response
            try:
                result_data = response.json()
                logger.info(f"Got {len(result_data.get('results', []))} results")
                return result_data
            except json.JSONDecodeError:
                logger.error("Failed to parse JSON response from OpenSearchNG")
                return {"error": "JSON decode error", "results": []}
                
        except requests.RequestException as e:
            logger.error(f"OpenSearchNG request failed: {str(e)}")
            return {"error": str(e), "results": []}

def test_opensearch_info() -> Dict[str, Any]:
    """
    Test getting information about the OpenSearchNG instance.
    
    Returns:
        Dict[str, Any]: Information about the OpenSearchNG instance
    """
    # Default configuration
    search_config = {
        "base_url": "http://localhost:8081",
        "timeout": 10,
    }
    
    # Construct the URL
    base_url = search_config["base_url"]
    config_url = f"{base_url}/config"
    
    try:
        logger.info("Getting OpenSearchNG instance information")
        response = requests.get(
            config_url,
            timeout=search_config["timeout"]
        )
        
        # Check for successful response
        if response.status_code != 200:
            logger.error(f"OpenSearchNG info request failed with status code: {response.status_code}")
            return {"error": f"Status code: {response.status_code}"}
        
        # Parse the response
        try:
            config_data = response.json()
            return config_data
        except json.JSONDecodeError:
            logger.error("Failed to parse JSON response from OpenSearchNG")
            return {"error": "JSON decode error"}
            
    except requests.RequestException as e:
        logger.error(f"OpenSearchNG info request failed: {str(e)}")
        return {"error": str(e)}

if __name__ == "__main__":
    # Test OpenSearchNG info
    logger.info("Testing OpenSearchNG info...")
    info_result = test_opensearch_info()
    if "error" in info_result:
        logger.error(f"Error getting OpenSearchNG info: {info_result['error']}")
    else:
        logger.info(f"OpenSearchNG info: {json.dumps(info_result, indent=2)[:200]}...")
    
    # Test OpenSearchNG query
    logger.info("Testing OpenSearchNG query...")
    query_result = test_opensearch_query("python programming")
    if "error" in query_result:
        logger.error(f"Error executing query: {query_result['error']}")
    else:
        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)}")