#!/usr/bin/env python3
"""
Test script to verify OpenSearchNG functionality in a Matrix client context.
"""
import logging
import json
import sys
import requests
from typing import Dict, Any, Optional, List
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Default OpenSearchNG configuration
DEFAULT_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: Optional[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 = DEFAULT_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)
async def execute_search_for_message(message: str, opensearch_config: Dict[str, Any]) -> str:
"""
Execute a search for a message and return the search context.
Args:
message: The message to search for
opensearch_config: The OpenSearchNG configuration
Returns:
str: The search context, or an empty string if search failed
"""
factual_intent = detect_factual_intent(message)
if factual_intent:
logger.info(f"Factual intent detected in message: {message}")
try:
results = opensearch_query(message, opensearch_config)
search_context = extract_search_context(results)
logger.info(f"Search context: {search_context[:100]}...")
return search_context
except Exception as e:
logger.error(f"Error executing search: {e}")
return ""
else:
logger.debug(f"No factual intent detected in message: {message}")
return ""
class MockMatrixClient:
"""Mock Matrix client for testing."""
def __init__(self, opensearch_config: Dict[str, Any] = None):
"""Initialize the mock Matrix client."""
self.opensearch_config = opensearch_config or DEFAULT_OPENSEARCH_CONFIG
self.enable_search = True
self.messages = []
async def send_message(self, room_id: str, message: str) -> None:
"""
Send a message to a room.
Args:
room_id: The ID of the room to send the message to
message: The message to send
"""
logger.info(f"Sending message to room {room_id}: {message[:50]}...")
self.messages.append({"room_id": room_id, "message": message})
async def handle_search_command(client: MockMatrixClient, query: str, room_id: str) -> None:
"""
Handle the /search command.
Args:
client: The MockMatrixClient instance
query: The search query
room_id: The ID of the room where the command was sent
"""
if not client.enable_search:
await client.send_message(room_id, "Search functionality is not enabled.")
return
await client.send_message(room_id, f"Searching for: {query}...")
try:
results = opensearch_query(query, client.opensearch_config)
if not results or not results.get("results"):
await client.send_message(room_id, f"No results found for: {query}")
return
# Format the results
formatted_results = ["Search Results:"]
for i, result in enumerate(results["results"][:5], 1):
title = result.get("title", "No title")
url = result.get("url", "No URL")
content = result.get("content", "No content")
# Truncate content if it's too long
if len(content) > 200:
content = content[:200] + "..."
formatted_results.append(f"{i}. {title}")
formatted_results.append(f" URL: {url}")
formatted_results.append(f" {content}")
formatted_results.append("")
# Send the formatted results
await client.send_message(room_id, "\n".join(formatted_results))
except Exception as e:
logger.error(f"Error executing search: {e}")
await client.send_message(room_id, f"Error executing search: {str(e)}")
async def test_matrix_search():
"""Test the Matrix search functionality."""
try:
# Create a mock Matrix client
client = MockMatrixClient()
# Test the search command
logger.info("Testing search command...")
await handle_search_command(client, "python programming", "test_room")
# Check if we got a response
if not client.messages:
logger.error("No messages sent by the client")
return False
# Print the last message
last_message = client.messages[-1]["message"]
logger.info(f"Last message: {last_message[:100]}...")
# Test search for message
logger.info("Testing search for message...")
search_context = await execute_search_for_message("What is Python programming?", client.opensearch_config)
if not search_context:
logger.error("No search context returned")
return False
logger.info(f"Search context: {search_context[:100]}...")
return True
except Exception as e:
logger.error(f"Error: {e}")
return False
if __name__ == "__main__":
import asyncio
success = asyncio.run(test_matrix_search())
if success:
logger.info("Matrix search test successful!")
sys.exit(0)
else:
logger.error("Matrix search test failed!")
sys.exit(1)