"""
OpenSearchNG integration for OllamaTools.
This module extends the web_tools capabilities by adding integration with
a locally running OpenSearchNG instance for improved search functionality.
"""
import requests
import json
import logging
from typing import Dict, Any, Optional, List, Union, Set
from urllib.parse import quote_plus
# Import URL validation utilities
from ollamatools.core.utils.validators.url import (
validate_url,
URLValidationError,
configure_domain_lists,
)
# Import error handling utilities
from ollamatools.core.errors import (
WebError,
ERROR_CODES,
handle_exception
)
# Set up logging
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,
permission_manager = None,
domain_allowlist: Optional[Set[str]] = None,
domain_blocklist: Optional[Set[str]] = 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)
permission_manager: Optional permission manager to request permission
domain_allowlist: Optional set of allowed domains for results
domain_blocklist: Optional set of blocked domains for results
Returns:
dict: Search results with metadata
Raises:
ValueError: If parameters are invalid
WebError: If the search fails
PermissionError: If permission is denied
"""
# 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)
# Request permission if permission manager is provided
if permission_manager:
args = {
"query": query,
"config": search_config,
}
if not permission_manager.request_permission("opensearch_query", args):
logger.warning(f"Permission denied for OpenSearchNG query: {query}")
raise PermissionError(f"Permission denied for OpenSearchNG query: {query}")
# Construct the search URL
base_url = search_config["base_url"]
search_path = search_config["search_path"]
search_url = f"{base_url}{search_path}"
# Validate the OpenSearchNG instance URL
try:
# Allow localhost for OpenSearchNG access
local_allowlist = {"localhost", "127.0.0.1"}
if domain_allowlist:
local_allowlist.update(domain_allowlist)
# Extract the host from the base URL for validation
from urllib.parse import urlparse
parsed_url = urlparse(base_url)
hostname = parsed_url.netloc.split(':')[0]
if hostname not in local_allowlist:
validate_url(
base_url,
domain_allowlist=local_allowlist,
domain_blocklist=domain_blocklist
)
except URLValidationError as e:
# Make a special exception for localhost
if "localhost" not in base_url and "127.0.0.1" not in base_url:
raise ValueError(f"Invalid OpenSearchNG URL: {str(e)}")
# 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:
raise WebError(
f"OpenSearchNG query failed with status code: {response.status_code}",
url=search_url,
status_code=response.status_code,
code=ERROR_CODES["URL_RESPONSE_ERROR"]
)
# Parse the response based on format
if search_config["format"] == "json":
try:
result_data = response.json()
except json.JSONDecodeError:
raise WebError(
"Failed to parse JSON response from OpenSearchNG",
url=search_url,
code=ERROR_CODES["URL_RESPONSE_ERROR"]
)
else:
# For other formats, return the raw content
result_data = {"raw_content": response.text}
# Validate and process the results
processed_results = process_opensearch_results(result_data, query)
return processed_results
except requests.RequestException as e:
raise WebError(
f"OpenSearchNG request failed: {str(e)}",
url=search_url,
code=ERROR_CODES["URL_CONNECTION_ERROR"]
)
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 get_opensearch_info(
config: Optional[Dict[str, Any]] = None,
permission_manager = None
) -> Dict[str, Any]:
"""
Get information about the OpenSearchNG instance.
Args:
config: OpenSearchNG configuration (defaults used if None)
permission_manager: Optional permission manager to request permission
Returns:
dict: Information about the OpenSearchNG instance
Raises:
WebError: If the request fails
PermissionError: If permission is denied
"""
# Use default config if none provided
search_config = DEFAULT_OPENSEARCH_CONFIG.copy()
if config:
search_config.update(config)
# Request permission if permission manager is provided
if permission_manager:
if not permission_manager.request_permission("get_opensearch_info", {"config": search_config}):
logger.warning("Permission denied for getting OpenSearchNG info")
raise PermissionError("Permission denied for getting OpenSearchNG info")
# 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:
raise WebError(
f"OpenSearchNG info request failed with status code: {response.status_code}",
url=config_url,
status_code=response.status_code,
code=ERROR_CODES["URL_RESPONSE_ERROR"]
)
# Parse the response
try:
config_data = response.json()
return config_data
except json.JSONDecodeError:
raise WebError(
"Failed to parse JSON response from OpenSearchNG",
url=config_url,
code=ERROR_CODES["URL_RESPONSE_ERROR"]
)
except requests.RequestException as e:
raise WebError(
f"OpenSearchNG info request failed: {str(e)}",
url=config_url,
code=ERROR_CODES["URL_CONNECTION_ERROR"]
)
def opensearch_fetch_url(
url: str,
config: Optional[Dict[str, Any]] = None,
permission_manager = None,
domain_allowlist: Optional[Set[str]] = None,
domain_blocklist: Optional[Set[str]] = None
) -> Dict[str, Any]:
"""
Fetch a URL through the OpenSearchNG proxy to avoid direct connections.
This provides an additional security layer for web requests.
Args:
url: URL to fetch
config: OpenSearchNG configuration (defaults used if None)
permission_manager: Optional permission manager to request permission
domain_allowlist: Optional set of allowed domains
domain_blocklist: Optional set of blocked domains
Returns:
dict: Response content and metadata
Raises:
ValueError: If URL is invalid
WebError: If the fetch fails
PermissionError: If permission is denied
"""
# Validate URL
try:
validated_url = validate_url(
url,
domain_allowlist=domain_allowlist,
domain_blocklist=domain_blocklist
)
except URLValidationError as e:
raise ValueError(f"URL validation failed: {str(e)}")
# Use default config if none provided
search_config = DEFAULT_OPENSEARCH_CONFIG.copy()
if config:
search_config.update(config)
# Request permission if permission manager is provided
if permission_manager:
args = {
"url": url,
"config": search_config
}
if not permission_manager.request_permission("opensearch_fetch_url", args):
logger.warning(f"Permission denied for fetching URL via OpenSearchNG: {url}")
raise PermissionError(f"Permission denied for fetching URL via OpenSearchNG: {url}")
# Construct the URL for the OpenSearchNG proxy
base_url = search_config["base_url"]
proxy_url = f"{base_url}/url_proxy"
# Prepare query parameters
params = {
"url": validated_url
}
try:
logger.info(f"Fetching URL via OpenSearchNG proxy: {validated_url}")
response = requests.get(
proxy_url,
params=params,
timeout=search_config["timeout"]
)
# Check for successful response
if response.status_code != 200:
raise WebError(
f"URL fetch via OpenSearchNG failed with status code: {response.status_code}",
url=proxy_url,
status_code=response.status_code,
code=ERROR_CODES["URL_RESPONSE_ERROR"]
)
# Create result with metadata
result = {
"status": response.status_code,
"headers": dict(response.headers),
"content": response.text,
"url": validated_url
}
return result
except requests.RequestException as e:
raise WebError(
f"URL fetch via OpenSearchNG failed: {str(e)}",
url=proxy_url,
code=ERROR_CODES["URL_CONNECTION_ERROR"]
)