ollama_semantic_function_calling / ollamatools-new / ollamatools / plugins / integrations / matrix / search.py

"""
Search functionality for the Enhanced Matrix integration in OllamaTools.

This module provides search-related functionality for the Enhanced Matrix client,
including integration with OpenSearchNG and the modular agent's search capabilities.

Dependencies:
- matrix-nio: For Matrix protocol support
- ollama: For Ollama API access
- ollamatools.agent.modular: For search integration and natural language understanding
- ollamatools.opensearch_web_tools: For web search functionality
"""

import logging
from typing import Dict, List, Optional, Any, Union

# Configure logging
logger = logging.getLogger(__name__)

# Import from modular agent (optional)
try:
    from ollamatools.plugins.standard.search import (
        detect_factual_intent,
        extract_search_context,
        execute_search
    )
    MODULAR_AGENT_AVAILABLE = True
except ImportError:
    MODULAR_AGENT_AVAILABLE = False
    # Only log this as debug since search might be disabled in config
    logging.debug(
        "The modular agent components are not available. "
        "Search functionality will be disabled."
    )

# Import from opensearch web tools (optional)
try:
    from ollamatools.plugins.integrations.opensearch.tools import (
        opensearch_query,
        opensearch_fetch_url,
        get_opensearch_info,
        DEFAULT_OPENSEARCH_CONFIG
    )
    OPENSEARCH_AVAILABLE = True
except ImportError:
    OPENSEARCH_AVAILABLE = False
    # Only log this as debug since search might be disabled in config
    logging.debug(
        "The opensearch_web_tools module is not available. "
        "Search functionality will be disabled."
    )

def get_default_opensearch_config() -> Dict[str, Any]:
    """
    Get the default OpenSearchNG configuration.
    
    Returns:
        Dict[str, Any]: The default configuration
    """
    if OPENSEARCH_AVAILABLE:
        return DEFAULT_OPENSEARCH_CONFIG.copy()
    else:
        return {}

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
    """
    if not MODULAR_AGENT_AVAILABLE or not OPENSEARCH_AVAILABLE:
        logger.warning("Search functionality is not available.")
        return ""
    
    try:
        factual_intent = detect_factual_intent(message)
        if factual_intent:
            logger.info(f"Factual intent detected in message: {message}")
            try:
                results = await execute_search_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 ""
    except Exception as e:
        logger.error(f"Error detecting factual intent: {e}")
        return ""

async def execute_search_query(query: str, opensearch_config: Dict[str, Any]) -> Dict[str, Any]:
    """
    Execute a search query using OpenSearchNG.
    
    Args:
        query: The search query
        opensearch_config: The OpenSearchNG configuration
        
    Returns:
        Dict[str, Any]: Search results
    """
    if not OPENSEARCH_AVAILABLE:
        logger.warning("OpenSearchNG is not available.")
        return {"error": "Search functionality is not available.", "results": []}
    
    try:
        # Execute the search
        results = opensearch_query(
            query=query,
            config=opensearch_config
        )
        
        return results
    except Exception as e:
        logger.error(f"Error executing search: {e}")
        return {"error": str(e), "results": []}

async def handle_search_command(client: Any, query: str, room_id: str) -> None:
    """
    Handle the /search command.
    
    Args:
        client: The EnhancedMatrixClient 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 = await execute_search_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)}")