OverviewFilesCommitsIssuesDocsReleasesPeers

ollama_semantic_function_calling / ollamatools-new / tests / test_direct_opensearch.py

328 lines · python ref:
#!/usr/bin/env python3
"""
Test script to directly verify OpenSearchNG functionality.
"""

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

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

# OpenSearchNG configuration
OPENSEARCH_CONFIG = {
    "base_url": "http://localhost:8081",  # Default local OpenSearchNG URL
    "search_path": "/search",             # Path for search endpoint
    "preferences_path": "/preferences",   # Path for preferences endpoint
    "categories": ["general"],            # Default search categories
    "engines": ["google", "bing", "duckduckgo"], # Default search engines
    "language": "en-US",                  # Default language
    "format": "json",                     # Response format (html, json, csv, rss)
    "timeout": 10,                        # Default timeout in seconds
    "num_results": 10,                    # Default number of results
    "safe_search": 0                      # 0=off, 1=moderate, 2=strict
}

def opensearch_query(query: str, config: Dict[str, Any] = None) -> Dict[str, Any]:
    """
    Perform a search query using a local OpenSearchNG instance.
    
    Args:
        query: Search query string
        config: OpenSearchNG configuration (defaults used if None)
        
    Returns:
        dict: Search results with metadata
    """
    # Validate input
    if not query or not isinstance(query, str):
        raise ValueError("Query must be a non-empty string")
    
    # Use default config if none provided
    search_config = OPENSEARCH_CONFIG.copy()
    if config:
        search_config.update(config)
    
    # 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"],
        "safesearch": search_config["safe_search"],
        "time_range": "",  # Optional time range
        "pageno": 1,       # First page of results
    }
    
    # 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 based on format
            if search_config["format"] == "json":
                try:
                    result_data = response.json()
                except json.JSONDecodeError:
                    logger.error("Failed to parse JSON response from OpenSearchNG")
                    return {"error": "JSON decode error", "results": []}
            else:
                # For other formats, return the raw content
                result_data = {"raw_content": response.text}
            
            # Process the results
            processed_results = process_opensearch_results(result_data, query)
            
            return processed_results
            
        except requests.RequestException as e:
            logger.error(f"OpenSearchNG request failed: {str(e)}")
            return {"error": str(e), "results": []}

def process_opensearch_results(result_data: Dict[str, Any], query: str) -> Dict[str, Any]:
    """
    Process and format the OpenSearchNG results.
    
    Args:
        result_data: Raw result data from OpenSearchNG
        query: Original search query
        
    Returns:
        dict: Processed search results with metadata
    """
    processed_results = {
        "query": query,
        "results": [],
        "suggestion": None,
        "answer": None,
        "infobox": None,
        "engines": [],
        "timing": {}
    }
    
    # Extract search results
    if "results" in result_data:
        for result in result_data["results"]:
            processed_result = {
                "title": result.get("title", ""),
                "url": result.get("url", ""),
                "content": result.get("content", ""),
                "engine": result.get("engine", ""),
                "score": result.get("score", 0),
                "category": result.get("category", ""),
                "pretty_url": result.get("pretty_url", "")
            }
            processed_results["results"].append(processed_result)
    
    # Extract suggestions
    if "suggestions" in result_data and result_data["suggestions"]:
        try:
            suggestion = result_data["suggestions"][0]
            # Handle case where suggestion is a dict instead of a string
            if isinstance(suggestion, dict):
                processed_results["suggestion"] = suggestion.get("value", str(suggestion))
            else:
                processed_results["suggestion"] = suggestion
        except (IndexError, TypeError) as e:
            logger.warning(f"Error processing suggestions: {str(e)}")
            processed_results["suggestion"] = None
    
    # Extract answer if available
    if "answers" in result_data and result_data["answers"]:
        try:
            answer = result_data["answers"][0]
            # Handle case where answer is a dict instead of a string
            if isinstance(answer, dict):
                processed_results["answer"] = answer.get("content", str(answer))
            else:
                processed_results["answer"] = answer
        except (IndexError, TypeError) as e:
            logger.warning(f"Error processing answers: {str(e)}")
            processed_results["answer"] = None
    
    # Extract infobox if available
    if "infoboxes" in result_data and result_data["infoboxes"]:
        try:
            infobox = result_data["infoboxes"][0]
            # Handle case where infobox is a dict (which is expected)
            processed_results["infobox"] = infobox
        except (IndexError, TypeError) as e:
            logger.warning(f"Error processing infoboxes: {str(e)}")
            processed_results["infobox"] = None
    
    # Extract engines used
    if "engines" in result_data:
        processed_results["engines"] = result_data["engines"]
    
    # Extract timing information
    if "timing" in result_data:
        processed_results["timing"] = result_data["timing"]
    
    return processed_results

def detect_factual_intent(message: str) -> bool:
    """
    Detect if a message has factual intent (i.e., asking for information).
    
    Args:
        message: The message to check
        
    Returns:
        bool: True if the message has factual intent, False otherwise
    """
    # Simple heuristic: check if the message contains question words or phrases
    question_words = [
        "what", "who", "where", "when", "why", "how", "is", "are", "can", "could",
        "would", "should", "tell me about", "explain", "search", "find", "look up"
    ]
    
    message_lower = message.lower()
    
    # Check for question words
    for word in question_words:
        if word in message_lower:
            return True
    
    # Check for question marks
    if "?" in message:
        return True
    
    return False

def extract_search_context(results: Dict[str, Any]) -> str:
    """
    Extract search context from OpenSearchNG results.
    
    Args:
        results: OpenSearchNG results
        
    Returns:
        str: Search context
    """
    if not results or "results" not in results or not results["results"]:
        return ""
    
    # Extract information from results
    context_parts = []
    
    # Add answer if available
    if results.get("answer"):
        context_parts.append(f"Answer: {results['answer']}")
    
    # Add infobox if available
    if results.get("infobox"):
        infobox = results["infobox"]
        if isinstance(infobox, dict):
            infobox_content = infobox.get("content", "")
            if infobox_content:
                context_parts.append(f"Infobox: {infobox_content}")
    
    # Add top results
    for i, result in enumerate(results["results"][:5], 1):
        title = result.get("title", "No title")
        content = result.get("content", "No content")
        url = result.get("url", "No URL")
        
        context_parts.append(f"Result {i}: {title}")
        context_parts.append(f"Content: {content}")
        context_parts.append(f"URL: {url}")
        context_parts.append("")
    
    # Add suggestion if available
    if results.get("suggestion"):
        context_parts.append(f"Did you mean: {results['suggestion']}?")
    
    return "\n".join(context_parts)

def test_direct_opensearch():
    """Test OpenSearchNG functionality directly."""
    try:
        # Test the search command
        query = "python programming"
        logger.info(f"Searching for: {query}")
        
        # Execute the search query
        results = opensearch_query(query)
        
        # Check if we got results
        if not results or not results.get("results"):
            logger.error(f"No results found for: {query}")
            return False
        
        # Print the number of results
        num_results = len(results.get("results", []))
        logger.info(f"Got {num_results} results for '{query}'")
        
        # Print the first result
        if num_results > 0:
            first_result = results["results"][0]
            logger.info(f"First result: {json.dumps(first_result, indent=2)}")
        
        # Test factual intent detection and search context extraction
        logger.info("Testing factual intent detection and search context extraction...")
        
        message = "What is Python programming?"
        logger.info(f"Message: {message}")
        
        # Check if the message has factual intent
        factual_intent = detect_factual_intent(message)
        if factual_intent:
            logger.info(f"Factual intent detected in message: {message}")
            
            # Execute the search query
            results = opensearch_query(message)
            
            # Extract search context
            search_context = extract_search_context(results)
            
            if search_context:
                logger.info(f"Search context: {search_context[:100]}...")
                logger.info("Factual intent detection and search context extraction test completed")
            else:
                logger.error("No search context extracted")
                return False
        else:
            logger.error("No factual intent detected")
            return False
        
        return True
        
    except Exception as e:
        logger.error(f"Error: {e}")
        return False

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