task_generator.py

"""
Task Generator - LLM-powered task refinement for browser automation.

Converts vague user inputs into structured, numbered task steps using Gemini.
"""

import asyncio
from typing import Optional
from google import genai
from google.genai import types


async def generate_task(
    user_input: str,
    api_key: str,
    model: Optional[str] = None,
    verbose: bool = False
) -> str:
    """
    Convert vague user input into structured numbered task steps.

    Args:
        user_input: The user's original task description
        api_key: Google API key for Gemini
        model: Gemini model name (defaults to gemini-3-flash-preview)
        verbose: Print generation progress

    Returns:
        str: Refined task with numbered steps, or original input if generation fails
    """
    # Default model
    if not model:
        model = "gemini-3-flash-preview"

    # Validate API key
    if not api_key:
        if verbose:
            print("Warning: No API key provided for task generation")
        return user_input

    # System prompt for task refinement
    prompt = f"""You are a task planning assistant for a browser automation agent.

Available browser actions:
- Navigate to a URL (must be explicit, e.g., https://example.com)
- Take snapshot to see page elements
- Click on elements by index
- Type text into input fields
- Save current page as PDF
- Go back to previous page
- Wait for specified seconds

Convert the user's request into a numbered list of clear, specific steps (3-10 steps).

Requirements:
1. Use full URLs (https://example.com, not "search for URL")
2. Include snapshot steps after navigation to see elements
3. Be explicit about filenames for PDFs
4. Keep steps actionable and clear
5. Output ONLY the numbered list, no extra commentary

USER TASK: {user_input}

OUTPUT (numbered steps only):"""

    try:
        if verbose:
            print(f"Generating structured task with {model}...")

        # Initialize client
        client = genai.Client(api_key=api_key)

        # Generate with timeout
        response = await asyncio.wait_for(
            asyncio.to_thread(
                client.models.generate_content,
                model=model,
                contents=[types.Content(
                    role="user",
                    parts=[types.Part.from_text(text=prompt)]
                )],
                config=types.GenerateContentConfig(
                    temperature=0.3,  # Low temperature for consistent output
                    max_output_tokens=500
                )
            ),
            timeout=10.0  # 10 second timeout
        )

        # Extract text from response
        if not response.candidates:
            if verbose:
                print("Warning: No response from LLM, using original input")
            return user_input

        refined_task = response.candidates[0].content.parts[0].text.strip()

        # Validate that we got a numbered list
        if not refined_task or not any(line.strip().startswith(('1.', '1)', '1 ')) for line in refined_task.split('\n')[:3]):
            if verbose:
                print("Warning: Response not in numbered format, using original input")
            return user_input

        return refined_task

    except asyncio.TimeoutError:
        if verbose:
            print("Warning: Task generation timed out, retrying once...")

        # Retry once
        try:
            response = await asyncio.wait_for(
                asyncio.to_thread(
                    client.models.generate_content,
                    model=model,
                    contents=[types.Content(
                        role="user",
                        parts=[types.Part.from_text(text=prompt)]
                    )],
                    config=types.GenerateContentConfig(
                        temperature=0.3,
                        max_output_tokens=500
                    )
                ),
                timeout=10.0
            )

            if response.candidates:
                refined_task = response.candidates[0].content.parts[0].text.strip()
                if refined_task and any(line.strip().startswith(('1.', '1)', '1 ')) for line in refined_task.split('\n')[:3]):
                    return refined_task
        except Exception as retry_error:
            if verbose:
                print(f"Warning: Retry failed ({retry_error}), using original input")

        return user_input

    except Exception as e:
        if verbose:
            print(f"Warning: Task generation failed ({e}), using original input")
        return user_input