parallel_tool_execution.py

#!/usr/bin/env python3
"""
Parallel tool execution optimization for ALS Research Agent
This module replaces sequential tool execution with parallel execution
to reduce response time by ~60-70% for multi-tool queries.
"""

import asyncio
from typing import List, Dict, Tuple, Any
import logging

logger = logging.getLogger(__name__)


async def execute_single_tool(
    tool_call: Dict,
    call_mcp_tool_func,
    index: int
) -> Tuple[int, str, Dict]:
    """
    Execute a single tool call asynchronously.
    Returns (index, progress_text, result_dict) to maintain order.
    """
    tool_name = tool_call["name"]
    tool_args = tool_call["input"]

    # Show search info in progress text
    tool_display = tool_name.replace('__', ' → ')
    search_info = ""
    if "query" in tool_args:
        search_info = f" `{tool_args['query'][:50]}{'...' if len(tool_args['query']) > 50 else ''}`"
    elif "condition" in tool_args:
        search_info = f" `{tool_args['condition'][:50]}{'...' if len(tool_args['condition']) > 50 else ''}`"

    try:
        # Call MCP tool
        start_time = asyncio.get_event_loop().time()
        tool_result = await call_mcp_tool_func(tool_name, tool_args)
        elapsed = asyncio.get_event_loop().time() - start_time

        logger.info(f"Tool {tool_name} completed in {elapsed:.2f}s")

        # Check for zero results to provide clear indicators
        has_results = True
        results_count = 0

        if isinstance(tool_result, str):
            result_lower = tool_result.lower()

            # Check for specific result counts
            import re
            count_matches = re.findall(r'found (\d+) (?:papers?|trials?|preprints?|results?)', result_lower)
            if count_matches:
                results_count = int(count_matches[0])

            # Also check JSON "total" field (from structured tool responses)
            if results_count == 0:
                total_matches = re.findall(r'"total":\s*(\d+)', result_lower)
                if total_matches:
                    results_count = int(total_matches[0])

            # Check for explicit no-results indicators
            no_result_phrases = [
                "no results found", "0 results", "no papers found",
                "no trials found", "no preprints found", "not found",
                "zero results", "no matches", "no als trials found",
                "no recruiting als trials", "no new or updated"
            ]
            if any(phrase in result_lower for phrase in no_result_phrases):
                has_results = False
            elif results_count == 0 and '"error"' not in result_lower:
                # Only mark as no-results if count is 0 AND no error (avoid false negatives)
                # If we simply couldn't detect a count, assume results exist
                has_results = True

        # Create clear success/failure indicator
        if has_results:
            if results_count > 0:
                progress_text = f"\n✅ **Found {results_count} results:** {tool_display}{search_info}"
            else:
                progress_text = f"\n✅ **Success:** {tool_display}{search_info}"
        else:
            progress_text = f"\n⚠️ **No results:** {tool_display}{search_info} - will try alternatives"

        # Add timing for long operations
        if elapsed > 5:
            progress_text += f" (took {elapsed:.1f}s)"

        # Check for zero results to enable self-correction
        if not has_results:
                # Add self-correction hint to the result
                tool_result += "\n\n**SELF-CORRECTION HINT:** No results found with this query. Consider:\n"
                tool_result += "1. Broadening search terms (remove qualifiers)\n"
                tool_result += "2. Using alternative terminology or synonyms\n"
                tool_result += "3. Searching related concepts\n"
                tool_result += "4. Checking for typos in search terms"

        result_dict = {
            "type": "tool_result",
            "tool_use_id": tool_call["id"],
            "content": tool_result
        }

        return index, progress_text, result_dict

    except Exception as e:
        logger.error(f"Error executing tool {tool_name}: {e}")

        # Clear failure indicator for errors
        progress_text = f"\n❌ **Failed:** {tool_display}{search_info} - {str(e)[:50]}"

        error_result = {
            "type": "tool_result",
            "tool_use_id": tool_call["id"],
            "content": f"Error executing tool: {str(e)}"
        }
        return index, progress_text, error_result


async def execute_tool_calls_parallel(
    tool_calls: List[Dict],
    call_mcp_tool_func,
    progress_callback=None
) -> Tuple[str, List[Dict]]:
    """
    Execute tool calls in parallel, reporting results as they arrive.
    Maintains the original order of tool calls in final results.

    Args:
        tool_calls: List of tool calls to execute
        call_mcp_tool_func: Function to call MCP tools
        progress_callback: Optional async callback(progress_text) called as each tool completes

    Returns: (progress_text, tool_results_content)
    """
    if not tool_calls:
        return "", []

    # Track execution time for progress reporting
    start_time = asyncio.get_event_loop().time()
    total_count = len(tool_calls)

    # Log parallel execution
    logger.info(f"Executing {total_count} tools in parallel")

    # Create named tasks for as_completed tracking
    tasks = {
        asyncio.create_task(execute_single_tool(tool_call, call_mcp_tool_func, i)): i
        for i, tool_call in enumerate(tool_calls)
    }

    # Collect results as they arrive
    completed_results = []
    completed_count = 0
    progress_text = ""

    for coro in asyncio.as_completed(tasks.keys()):
        try:
            result = await coro
            completed_count += 1
            index, prog_text, result_dict = result
            completed_results.append((index, prog_text, result_dict))

            # Report progress as each tool completes
            progress_line = f"\n📊 **Search Progress:** Completed {completed_count}/{total_count} searches"
            progress_line += prog_text
            if progress_callback:
                await progress_callback(progress_line)

        except Exception as e:
            completed_count += 1
            logger.error(f"Task failed with exception: {e}")

    # Sort results by original index to maintain order
    completed_results.sort(key=lambda x: x[0])

    # Build final progress text and results
    elapsed_time = asyncio.get_event_loop().time() - start_time
    timing_info = f" in {elapsed_time:.1f}s" if elapsed_time > 5 else ""

    progress_text = f"\n📊 **Search Progress:** Completed {len(completed_results)}/{total_count} searches{timing_info}\n"

    tool_results_content = []
    for index, prog_text, result_dict in completed_results:
        progress_text += prog_text
        tool_results_content.append(result_dict)

    # Handle any tasks that raised exceptions not caught above
    for task, original_index in tasks.items():
        if task.done() and task.exception() and original_index < len(tool_calls):
            error_already_added = any(r[0] == original_index for r in completed_results)
            if not error_already_added:
                tool_results_content.append({
                    "type": "tool_result",
                    "tool_use_id": tool_calls[original_index]["id"],
                    "content": f"Tool execution failed: {str(task.exception())}"
                })

    return progress_text, tool_results_content


# Backward compatibility wrapper
async def execute_tool_calls_optimized(
    tool_calls: List[Dict],
    call_mcp_tool_func,
    parallel: bool = True
) -> Tuple[str, List[Dict]]:
    """
    Execute tool calls with optional parallel execution.

    Args:
        tool_calls: List of tool calls to execute
        call_mcp_tool_func: Function to call MCP tools
        parallel: If True, execute tools in parallel; if False, execute sequentially

    Returns: (progress_text, tool_results_content)
    """
    if parallel and len(tool_calls) > 1:
        # Use parallel execution for multiple tools
        return await execute_tool_calls_parallel(tool_calls, call_mcp_tool_func)
    else:
        # Fall back to sequential execution (import from original)
        from refactored_helpers import execute_tool_calls
        return await execute_tool_calls(tool_calls, call_mcp_tool_func)


def estimate_time_savings(num_tools: int, avg_tool_time: float = 3.5) -> Dict[str, float]:
    """
    Estimate time savings from parallel execution.

    Args:
        num_tools: Number of tools to execute
        avg_tool_time: Average time per tool in seconds

    Returns: Dictionary with timing estimates
    """
    sequential_time = num_tools * avg_tool_time
    # Parallel time is roughly the time of the slowest tool plus overhead
    parallel_time = avg_tool_time + 0.5  # 0.5s overhead for coordination

    savings = sequential_time - parallel_time
    savings_percent = (savings / sequential_time) * 100 if sequential_time > 0 else 0

    return {
        "sequential_time": sequential_time,
        "parallel_time": parallel_time,
        "time_saved": savings,
        "savings_percent": savings_percent
    }


# Test the optimization
if __name__ == "__main__":
    # Test time savings estimation
    for n in [2, 3, 4, 5]:
        estimates = estimate_time_savings(n)
        print(f"\n{n} tools:")
        print(f"  Sequential: {estimates['sequential_time']:.1f}s")
        print(f"  Parallel: {estimates['parallel_time']:.1f}s")
        print(f"  Savings: {estimates['time_saved']:.1f}s ({estimates['savings_percent']:.0f}%)")