utils.py

"""
GPT Researcher MCP Server Utilities

This module provides utility functions and helpers for the GPT Researcher MCP Server.
"""

import re
import html
import sys
from typing import Dict, List, Optional, Tuple, Any
from loguru import logger

try:
    from markdownify import markdownify as md
    HAS_MARKDOWNIFY = True
except ImportError:
    HAS_MARKDOWNIFY = False

# Configure logging for console only (no file logging)
logger.configure(handlers=[{"sink": sys.stderr, "level": "INFO"}])

# Research store to track ongoing research topics and contexts
research_store = {}


# HTML/Markdown Utilities
def html_to_markdown(text: str) -> str:
    """
    Convert HTML to Markdown, preserving structure (tables, headers, lists, etc.)
    
    Falls back to simple tag stripping if markdownify is not available.
    
    Args:
        text: Text that may contain HTML
        
    Returns:
        Clean Markdown text
    """
    if not text:
        return text
    
    # Decode HTML entities first
    text = html.unescape(text)
    
    # Check if text contains HTML tags
    if '<' not in text or '>' not in text:
        return text
    
    # Use markdownify if available (preserves structure)
    if HAS_MARKDOWNIFY:
        try:
            result = md(text, heading_style="atx", strip=['script', 'style'])
            # Clean up excessive newlines
            result = re.sub(r'\n{3,}', '\n\n', result)
            return result.strip()
        except Exception:
            pass  # Fall through to basic stripping
    
    # Fallback: strip tags (loses structure but removes HTML)
    text = re.sub(r'<[^>]+>', '', text)
    text = re.sub(r'\s+', ' ', text)
    return text.strip()


def strip_html(text: str) -> str:
    """
    Strip HTML tags and decode HTML entities from text.
    
    Deprecated: Use html_to_markdown() for better results.
    
    Args:
        text: Text that may contain HTML tags or entities
        
    Returns:
        Clean text with HTML removed
    """
    if not text:
        return text
    
    # Decode HTML entities (e.g., &amp; -> &, &#160; -> space, \u00b7 -> ·)
    text = html.unescape(text)
    
    # Remove HTML tags
    text = re.sub(r'<[^>]+>', '', text)
    
    # Remove excess whitespace
    text = re.sub(r'\s+', ' ', text)
    
    return text.strip()


def clean_search_results(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Convert HTML to Markdown in search results (typically from DuckDuckGo body field).
    
    Args:
        results: List of search result dictionaries
        
    Returns:
        Cleaned search results with Markdown
    """
    cleaned = []
    for result in results:
        cleaned_result = {}
        for key, value in result.items():
            if isinstance(value, str):
                cleaned_result[key] = html_to_markdown(value)
            else:
                cleaned_result[key] = value
        cleaned.append(cleaned_result)
    return cleaned


def clean_context(context: Any) -> str:
    """
    Convert HTML to Markdown from research context.
    
    Args:
        context: Research context (string or list of strings)
        
    Returns:
        Clean Markdown text
    """
    if isinstance(context, list):
        # Context is a list of strings
        cleaned_parts = [html_to_markdown(part) if isinstance(part, str) else str(part) for part in context]
        return '\n\n'.join(cleaned_parts)
    elif isinstance(context, str):
        return html_to_markdown(context)
    else:
        return str(context)

# API Response Utilities
def create_error_response(message: str) -> Dict[str, Any]:
    """Create a standardized error response"""
    return {"status": "error", "message": message}


def create_success_response(data: Dict[str, Any]) -> Dict[str, Any]:
    """Create a standardized success response"""
    return {"status": "success", **data}


def handle_exception(e: Exception, operation: str) -> Dict[str, Any]:
    """Handle exceptions in a consistent way"""
    error_message = str(e)
    logger.error(f"{operation} failed: {error_message}")
    return create_error_response(error_message)


def get_researcher_by_id(researchers_dict: Dict, research_id: str) -> Tuple[bool, Any, Dict[str, Any]]:
    """
    Helper function to retrieve a researcher by ID.
    
    Args:
        researchers_dict: Dictionary of research objects
        research_id: The ID of the research session
        
    Returns:
        Tuple containing (success, researcher_object, error_response)
    """
    if not researchers_dict or research_id not in researchers_dict:
        return False, None, create_error_response("Research ID not found. Please conduct research first.")
    return True, researchers_dict[research_id], {}


def format_sources_for_response(sources: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Format source information for API responses.
    
    Args:
        sources: List of source dictionaries
        
    Returns:
        Formatted source list for API responses
    """
    return [
        {
            "title": source.get("title", "Unknown"),
            "url": source.get("url", ""),
            "content_length": len(source.get("content", ""))
        }
        for source in sources
    ]


def format_context_with_sources(topic: str, context: str, sources: List[Dict[str, Any]]) -> str:
    """
    Format research context with sources for display.
    
    Args:
        topic: Research topic
        context: Research context
        sources: List of sources
        
    Returns:
        Formatted context string with sources
    """
    formatted_context = f"## Research: {topic}\n\n{context}\n\n"
    formatted_context += "## Sources:\n"
    for i, source in enumerate(sources):
        formatted_context += f"{i+1}. {source.get('title', 'Unknown')}: {source.get('url', '')}\n"
    return formatted_context


def store_research_results(topic: str, context: str, sources: List[Dict[str, Any]], 
                           source_urls: List[str], formatted_context: Optional[str] = None):
    """
    Store research results in the research store.
    
    Args:
        topic: Research topic
        context: Research context
        sources: List of sources
        source_urls: List of source URLs
        formatted_context: Optional pre-formatted context
    """
    research_store[topic] = {
        "context": formatted_context or context,
        "sources": sources,
        "source_urls": source_urls
    }


def create_research_prompt(topic: str, goal: str, report_format: str = "research_report") -> str:
    """
    Create a research query prompt for GPT Researcher.
    
    Args:
        topic: The topic to research
        goal: The goal or specific question to answer
        report_format: The format of the report to generate
        
    Returns:
        A formatted prompt for research
    """
    return f"""
    Please research the following topic: {topic}
    
    Goal: {goal}
    
    You have two methods to access web-sourced information:
    
    1. Use the "research://{topic}" resource to directly access context about this topic if it exists
       or if you want to get straight to the information without tracking a research ID.
       
    2. Use the deep_research tool to perform new research and get a research_id for later use.
       This tool also returns the context directly in its response, which you can use immediately.
    
    After getting context, you can:
    - Use it directly in your response
    - Use the write_report tool with a custom prompt to generate a structured {report_format}
    
    You can also use get_research_sources to view additional details about the information sources.
    """